"""
Index pre-parsed documents (JSON files from parse_pdf.py / parse_wikipedia.py)
into PostgreSQL, Elasticsearch, FAISS, and MinHash LSH.

Usage:
    python scripts/index_documents.py <dir_with_json_files> [--limit N]
"""
import sys
import os
import json
import argparse
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from app.config import settings
from app.models.db import SessionLocal, Document, Fingerprint, engine, Base
from app.core.fingerprint import WinNowing, MinHashLSHIndex
from app.core.embeddings import EmbeddingIndex
from app.core.fulltext import FulltextSearch


Base.metadata.create_all(bind=engine)


def load_json_docs(directory: str, limit: int | None = None):
    docs = []
    for path in sorted(Path(directory).glob("*.json")):
        try:
            with open(path, encoding="utf-8") as f:
                doc = json.load(f)
            if doc.get("text") and doc.get("title"):
                docs.append(doc)
        except Exception as e:
            print(f"Skip {path.name}: {e}")
        if limit and len(docs) >= limit:
            break
    return docs


def index_batch(docs: list, batch_size: int = 32):
    db = SessionLocal()
    winnowing = WinNowing()
    minhash_lsh = MinHashLSHIndex()
    embeddings = EmbeddingIndex()
    fulltext = FulltextSearch()

    print(f"Indexing {len(docs)} documents …")

    embed_batch: list[tuple[str, str]] = []

    for i, doc in enumerate(docs):
        source = doc.get("source", "unknown")
        ext_id = doc.get("id", "")
        title = doc["title"]
        text = doc["text"]
        url = doc.get("url", "")
        lang = doc.get("lang", "ru")

        # 1. PostgreSQL
        db_doc = Document(source=source, ext_id=ext_id, title=title, url=url, lang=lang)
        db.add(db_doc)
        db.flush()

        # 2. Fingerprints → PostgreSQL
        fp = winnowing.fingerprint(text)
        for pos, h in enumerate(list(fp)[:500]):
            db.add(Fingerprint(doc_id=db_doc.id, hash=format(h, '016x'), position=pos))

        # 3. MinHash LSH (in-memory; call save separately if needed)
        try:
            minhash_lsh.add_document(str(db_doc.id), text)
        except ValueError:
            pass

        # 4. Elasticsearch
        fulltext.index_document(str(db_doc.id), title, text, source, lang, url)

        # 5. Accumulate for batch embedding
        embed_batch.append((str(db_doc.id), text))

        if len(embed_batch) >= batch_size:
            embeddings.batch_add(embed_batch, batch_size=batch_size)
            embed_batch.clear()

        if (i + 1) % 100 == 0:
            db.commit()
            print(f"  {i + 1}/{len(docs)}")

    # Flush remaining embeddings
    if embed_batch:
        embeddings.batch_add(embed_batch, batch_size=batch_size)

    db.commit()
    db.close()

    embeddings.save_index()
    print("Done. FAISS index saved.")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("directory", help="Directory with parsed JSON files")
    parser.add_argument("--limit", type=int, default=None, help="Max documents to index")
    parser.add_argument("--batch-size", type=int, default=32)
    args = parser.parse_args()

    docs = load_json_docs(args.directory, args.limit)
    print(f"Loaded {len(docs)} documents from {args.directory}")

    if not docs:
        print("No documents found, exiting.")
        sys.exit(0)

    index_batch(docs, batch_size=args.batch_size)


if __name__ == "__main__":
    main()