"""
Build turbovec IdMapIndex from Qdrant tipitaka_chunks collection.

ID encoding: volume * 100000 + page  (uint64, fully reversible)
Output: tipitaka_chunks.tvim  (turbovec index file)

Usage:
    python build_turbovec_index.py
    python build_turbovec_index.py --collection tipitaka_chunks --output tipitaka_chunks.tvim
"""
import argparse
import logging
import time
import numpy as np
from qdrant_client import QdrantClient
from turbovec import IdMapIndex

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

QDRANT_URL   = "http://localhost:6333"
COLLECTION   = "tipitaka_chunks"
OUTPUT_PATH  = "F:/_Ai/Tipitaka-AI-Expert/RAG/webapp/tipitaka-api/data/tipitaka_chunks.tvim"
EMBED_DIM    = 1024
BATCH_SIZE   = 500


def encode_id(volume: str, page: str) -> int:
    """Encode (volume, page) as a single uint64. Reversible via decode_id()."""
    return int(volume) * 100_000 + int(page)


def decode_id(uid: int) -> tuple[int, int]:
    """Decode uint64 back to (volume, page)."""
    return uid // 100_000, uid % 100_000


def main(collection: str, output: str):
    logger.info(f"Connecting to Qdrant at {QDRANT_URL}...")
    client = QdrantClient(url=QDRANT_URL)

    info = client.get_collection(collection)
    total = info.points_count
    logger.info(f"Collection '{collection}': {total} points, dim={EMBED_DIM}")

    # Build turbovec index
    index = IdMapIndex(dim=EMBED_DIM, bit_width=4)

    vectors_batch = []
    ids_batch     = []
    processed     = 0
    skipped       = 0
    offset        = None
    start         = time.time()

    while True:
        pts, next_offset = client.scroll(
            collection,
            offset=offset,
            limit=BATCH_SIZE,
            with_payload=True,
            with_vectors=True,
        )
        if not pts:
            break

        for pt in pts:
            payload = pt.payload or {}
            volume  = payload.get("volume", payload.get("volume_number", ""))
            page    = payload.get("page",   payload.get("page_number",   ""))

            # Get dense vector only
            vec = pt.vector
            if isinstance(vec, dict):
                vec = vec.get("dense")
            if vec is None or not volume or not page:
                skipped += 1
                continue

            try:
                uid = encode_id(volume, page)
                vectors_batch.append(vec)
                ids_batch.append(uid)
            except (ValueError, TypeError):
                skipped += 1
                continue

        # Add batch to index
        if vectors_batch:
            vecs_np = np.array(vectors_batch, dtype=np.float32)
            ids_np  = np.array(ids_batch,     dtype=np.uint64)
            index.add_with_ids(vecs_np, ids_np)
            processed += len(vecs_np)
            vectors_batch.clear()
            ids_batch.clear()

        elapsed = time.time() - start
        rate    = processed / elapsed if elapsed > 0 else 0
        eta     = (total - processed) / rate if rate > 0 else 0
        logger.info(f"  {processed}/{total} ({100*processed//total}%) | {rate:.0f} pts/s | ETA {eta:.0f}s")

        if next_offset is None:
            break
        offset = next_offset

    # Save index
    logger.info(f"\nSaving index to {output}...")
    index.write(output)

    elapsed_total = time.time() - start
    logger.info("=" * 60)
    logger.info(f"Done in {elapsed_total:.1f}s")
    logger.info(f"Indexed: {processed} | Skipped: {skipped}")
    logger.info(f"Saved: {output}")
    logger.info(f"ID encoding: volume * 100000 + page  (decode: v=id//100000, p=id%100000)")
    logger.info("=" * 60)

    # Quick sanity check
    logger.info("\nSanity check — reload and search with random query...")
    from turbovec import IdMapIndex as TVI
    loaded = TVI.load(output)
    query  = np.random.randn(1, EMBED_DIM).astype(np.float32)
    scores, ids = loaded.search(query, k=3)
    for s, uid in zip(scores[0], ids[0]):
        vol, pg = decode_id(int(uid))
        logger.info(f"  score={s:.4f}  vol={vol}  page={pg}")

    client.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--collection", default=COLLECTION)
    parser.add_argument("--output",     default=OUTPUT_PATH)
    args = parser.parse_args()

    import os
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    main(args.collection, args.output)