import struct
import tarfile
import os
import sys
import subprocess

import pyarrow as pa

from pypaimon import CatalogFactory, Schema

FTP_URL = "ftp://ftp.irisa.fr/local/texmex/corpus/siftsmall.tar.gz"
DOWNLOAD_DIR = "/tmp/siftsmall"

BATCH_SIZE = 1000

DLF_URI = "http://dlf-regres-test-cn-wulanchabu-vpc.taobao.net"
DLF_WAREHOUSE = "index_test"
DATABASE = "default"
BASE_TABLE = "sift1k_base"
QUERY_TABLE = "sift1k_query"


def download_dataset():
    os.makedirs(DOWNLOAD_DIR, exist_ok=True)
    tar_path = os.path.join(DOWNLOAD_DIR, "siftsmall.tar.gz")
    if os.path.exists(tar_path):
        print(f"Already downloaded: {tar_path}")
    else:
        print(f"Downloading from {FTP_URL} ...")
        ret = subprocess.run(["curl", "-o", tar_path, FTP_URL],
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if ret.returncode != 0:
            print("curl failed, trying wget...")
            ret = subprocess.run(["wget", "-O", tar_path, FTP_URL],
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if ret.returncode != 0:
                print("Download failed. Try manually: wget " + FTP_URL)
                sys.exit(1)
        print(f"Downloaded: {tar_path} ({os.path.getsize(tar_path)} bytes)")

    extract_dir = os.path.join(DOWNLOAD_DIR, "extracted")
    if not os.path.exists(extract_dir):
        print("Extracting...")
        with tarfile.open(tar_path, "r:gz") as tf:
            tf.extractall(extract_dir)
    return os.path.join(extract_dir, "siftsmall")


def read_fvecs(filename):
    vectors = []
    with open(filename, "rb") as f:
        while True:
            dim_bytes = f.read(4)
            if len(dim_bytes) < 4:
                break
            dim = struct.unpack("<i", dim_bytes)[0]
            vec = struct.unpack("<" + "f" * dim, f.read(4 * dim))
            vectors.append(list(vec))
    return vectors


def read_ivecs(filename):
    vectors = []
    with open(filename, "rb") as f:
        while True:
            dim_bytes = f.read(4)
            if len(dim_bytes) < 4:
                break
            dim = struct.unpack("<i", dim_bytes)[0]
            vec = struct.unpack("<" + "i" * dim, f.read(4 * dim))
            vectors.append(list(vec))
    return vectors


def vecs_to_arrow_table(ids, vectors):
    return pa.table({
        'id': pa.array(ids, type=pa.int32()),
        'vector': pa.array(vectors, type=pa.list_(pa.float32())),
    })


def write_to_paimon(table, ids, vectors, label):
    print(f"\nWriting {label} ({len(ids)} rows) ...")
    write_builder = table.new_batch_write_builder()
    table_write = write_builder.new_write()
    table_commit = write_builder.new_commit()

    for start in range(0, len(ids), BATCH_SIZE):
        end = min(start + BATCH_SIZE, len(ids))
        batch = vecs_to_arrow_table(ids[start:end], vectors[start:end])
        table_write.write_arrow(batch)
        print(f"  Written {end}/{len(ids)}")

    commit_messages = table_write.prepare_commit()
    table_commit.commit(commit_messages)
    table_write.close()
    table_commit.close()
    print(f"  {label} committed.")


def main():
    ak = os.environ.get("DLF_ACCESS_KEY_ID")
    sk = os.environ.get("DLF_ACCESS_KEY_SECRET")
    if not ak or not sk:
        print("ERROR: DLF_ACCESS_KEY_ID and DLF_ACCESS_KEY_SECRET must be set in environment")
        sys.exit(1)

    sift_dir = download_dataset()
    if not os.path.isdir(sift_dir):
        print(f"Cannot find siftsmall directory at {sift_dir}")
        sys.exit(1)

    print("\nReading base vectors...")
    base_vecs = read_fvecs(os.path.join(sift_dir, "siftsmall_base.fvecs"))
    print(f"  Base vectors: {len(base_vecs)} x {len(base_vecs[0])}d")

    print("Reading query vectors...")
    query_vecs = read_fvecs(os.path.join(sift_dir, "siftsmall_query.fvecs"))
    print(f"  Query vectors: {len(query_vecs)} x {len(query_vecs[0])}d")

    print("Reading ground truth...")
    gt = read_ivecs(os.path.join(sift_dir, "siftsmall_groundtruth.ivecs"))
    print(f"  Ground truth: {len(gt)} queries x {len(gt[0])} neighbors")

    print(f"\nConnecting to DLF: {DLF_URI}, warehouse: {DLF_WAREHOUSE}")
    catalog = CatalogFactory.create({
        'metastore': 'rest',
        'uri': DLF_URI,
        'warehouse': DLF_WAREHOUSE,
        'token.provider': 'dlf',
        'dlf.access-key-id': ak,
        'dlf.access-key-secret': sk,
    })

    try:
        catalog.create_database(DATABASE, True)
        print(f"Database '{DATABASE}' ready.")
    except Exception as e:
        print(f"Database create/check: {e}")

    base_schema = Schema.from_pyarrow_schema(
        pa.schema([
            ('id', pa.int32()),
            ('vector', pa.list_(pa.float32())),
        ]),
        primary_keys=['id'],
        options={'bucket': '1'},
    )
    base_table_id = f"{DATABASE}.{BASE_TABLE}"
    try:
        catalog.create_table(base_table_id, base_schema, True)
        print(f"Table '{base_table_id}' ready.")
    except Exception as e:
        print(f"Table create/check: {e}")

    query_schema = Schema.from_pyarrow_schema(
        pa.schema([
            ('id', pa.int32()),
            ('vector', pa.list_(pa.float32())),
        ]),
        primary_keys=['id'],
        options={'bucket': '1'},
    )
    query_table_id = f"{DATABASE}.{QUERY_TABLE}"
    try:
        catalog.create_table(query_table_id, query_schema, True)
        print(f"Table '{query_table_id}' ready.")
    except Exception as e:
        print(f"Table create/check: {e}")

    base_table = catalog.get_table(base_table_id)
    base_ids = list(range(len(base_vecs)))
    write_to_paimon(base_table, base_ids, base_vecs, "base vectors")

    query_table = catalog.get_table(query_table_id)
    query_ids = list(range(len(query_vecs)))
    write_to_paimon(query_table, query_ids, query_vecs, "query vectors")

    print("\n=== Verification ===")
    for table_id in [base_table_id, query_table_id]:
        t = catalog.get_table(table_id)
        rb = t.new_read_builder()
        scan = rb.new_scan()
        read = rb.new_read()
        splits = scan.plan().splits()
        df = read.to_pandas(splits)
        print(f"  {table_id}: {len(df)} rows")

    print("\n=== Ground truth (for recall evaluation) ===")
    print(f"  100 queries, each with 100 ground truth neighbors")
    print(f"  Query 0 top-10 neighbors: {gt[0][:10]}")

    print("\nDone!")


if __name__ == "__main__":
    main()