"""
indexer.py
──────────
Reads articles.jsonl, embeds each article using a multilingual
sentence-transformers model, and saves a FAISS index to disk.

What this script produces:
  data/faiss.index      — the vector index (for similarity search)
  data/metadata.jsonl   — article metadata in the same order as the index
                          (needed to recover the text after a search)

Usage:
  python indexer.py

Install:
  pip install faiss-cpu sentence-transformers
"""

import json
import faiss
import numpy as np
from pathlib import Path
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

# ── Config ────────────────────────────────────────────────────────────────────

MODEL_NAME   = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
ARTICLES_FILE = Path("data/articles.jsonl")
INDEX_FILE    = Path("data/faiss.index")
METADATA_FILE = Path("data/metadata.jsonl")
BATCH_SIZE    = 32  # embed this many articles at a time


# ── Load articles ─────────────────────────────────────────────────────────────

def load_articles(path: Path) -> list[dict]:
    articles = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                articles.append(json.loads(line))
    print(f"Loaded {len(articles)} articles from {path}")
    return articles


# ── Embed ─────────────────────────────────────────────────────────────────────

def embed_articles(articles: list[dict], model: SentenceTransformer) -> np.ndarray:
    """
    Convert each article's 'chunk' field into a vector.

    The 'chunk' field contains:
        "LEGE 53 2003\nArticolul 142\n\nSalariații au dreptul..."

    Including the law title and article number in the chunk means the
    embedding captures source context, not just the article text alone.
    This improves retrieval accuracy.

    Returns a 2D numpy array of shape (n_articles, embedding_dim).
    """
    chunks = [a["chunk"] for a in articles]
    all_embeddings = []

    print(f"Embedding {len(chunks)} articles in batches of {BATCH_SIZE}...")
    for i in tqdm(range(0, len(chunks), BATCH_SIZE)):
        batch = chunks[i : i + BATCH_SIZE]
        embeddings = model.encode(
            batch,
            convert_to_numpy=True,
            normalize_embeddings=True,  # normalize for cosine similarity
            show_progress_bar=False,
        )
        all_embeddings.append(embeddings)

    return np.vstack(all_embeddings).astype("float32")


# ── Build FAISS index ─────────────────────────────────────────────────────────

def build_index(embeddings: np.ndarray) -> faiss.Index:
    """
    Build a FAISS index from the embeddings.

    We use IndexFlatIP (Inner Product) because our embeddings are
    normalized — inner product == cosine similarity when vectors are
    unit length. This gives us semantic similarity search for free.

    For larger datasets (100k+ articles) you'd switch to IndexIVFFlat
    which is approximate but much faster. For our size, exact search is fine.
    """
    dim = embeddings.shape[1]  # embedding dimension (384 for MiniLM)
    print(f"Building FAISS index — {len(embeddings)} vectors, dim={dim}")

    index = faiss.IndexFlatIP(dim)  # IP = Inner Product (cosine similarity)
    index.add(embeddings)

    print(f"Index contains {index.ntotal} vectors")
    return index


# ── Save ──────────────────────────────────────────────────────────────────────

def save_metadata(articles: list[dict], path: Path):
    """
    Save article metadata in the same order as the FAISS index.

    FAISS only stores vectors — it doesn't store text. When we search
    and get back index positions [42, 17, 8], we need this file to
    look up what article #42 actually says.
    """
    with open(path, "w", encoding="utf-8") as f:
        for article in articles:
            # Save only what we need for display — not the full chunk
            meta = {
                "law_id":         article.get("law_id"),
                "law_title":      article.get("law_title"),
                "article_number": article.get("article_number"),
                "text":           article.get("text"),
                "chunk":          article.get("chunk"),
            }
            f.write(json.dumps(meta, ensure_ascii=False) + "\n")
    print(f"Saved metadata to {path}")


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    # 1. Load articles
    articles = load_articles(ARTICLES_FILE)
    if not articles:
        print("No articles found. Run pipeline.py first.")
        return

    # 2. Load embedding model
    # First run downloads ~120MB — subsequent runs use the cache
    print(f"Loading model: {MODEL_NAME}")
    model = SentenceTransformer(MODEL_NAME)

    # 3. Embed all articles
    embeddings = embed_articles(articles, model)

    # 4. Build FAISS index
    index = build_index(embeddings)

    # 5. Save index and metadata
    faiss.write_index(index, str(INDEX_FILE))
    print(f"Saved FAISS index to {INDEX_FILE}")

    save_metadata(articles, METADATA_FILE)

    print("\nDone! You can now run retriever.py to search the index.")
    print(f"  Index:    {INDEX_FILE} ({INDEX_FILE.stat().st_size // 1024} KB)")
    print(f"  Metadata: {METADATA_FILE} ({METADATA_FILE.stat().st_size // 1024} KB)")


if __name__ == "__main__":
    main()