"""Embedding + BM25 cache helpers, as used by Notebooks 2-6.

Cache key = md5(domain | config_name | embedding_model_name | chunk_size |
chunk_overlap)[:12] -- this is the same scheme the notebooks use, so a
combo built by a notebook (and synced into `data/cache/indices/<key>/`) is
picked up here unchanged, and a combo built here is picked up by the
notebooks unchanged.
"""

import hashlib
import json
import pickle
from pathlib import Path

import faiss
import numpy as np
import pandas as pd
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

from src import settings


def make_cache_key(*parts) -> str:
    raw = "|".join(str(p) for p in parts)
    return hashlib.md5(raw.encode("utf-8")).hexdigest()[:12]


def get_or_build_dense_index(
    chunks_df: pd.DataFrame,
    model: SentenceTransformer,
    model_name: str,
    domain: str,
    config_name: str,
    chunk_size: int,
    chunk_overlap: int,
    cache_root: Path = None,
    docs_df: pd.DataFrame = None,
):
    cache_root = cache_root or settings.INDICES_DIR
    cache_key = make_cache_key(domain, config_name, model_name, chunk_size, chunk_overlap)
    cache_dir = cache_root / cache_key

    emb_path = cache_dir / "embeddings.npy"
    index_path = cache_dir / "faiss.index"

    if emb_path.exists() and index_path.exists():
        embeddings = np.load(emb_path)
        index = faiss.read_index(str(index_path))
        if docs_df is not None and not (cache_dir / "docs.parquet").exists():
            docs_df.to_parquet(cache_dir / "docs.parquet")
        return embeddings, index, cache_dir

    cache_dir.mkdir(parents=True, exist_ok=True)

    texts = chunks_df["text"].tolist()
    embeddings = model.encode(
        texts, batch_size=64, show_progress_bar=True, convert_to_numpy=True, normalize_embeddings=True
    ).astype("float32")

    index = faiss.IndexFlatIP(embeddings.shape[1])
    index.add(embeddings)

    np.save(emb_path, embeddings)
    faiss.write_index(index, str(index_path))
    chunks_df.to_parquet(cache_dir / "chunks.parquet")
    if docs_df is not None:
        docs_df.to_parquet(cache_dir / "docs.parquet")
    with open(cache_dir / "meta.json", "w") as f:
        json.dump({
            "model_name": model_name, "chunk_size": chunk_size, "chunk_overlap": chunk_overlap,
            "domain": domain, "config": config_name, "n_chunks": len(chunks_df),
        }, f, indent=2)

    return embeddings, index, cache_dir


def get_or_build_bm25(chunks_df: pd.DataFrame, cache_dir: Path) -> BM25Okapi:
    bm25_path = cache_dir / "bm25.pkl"
    if bm25_path.exists():
        with open(bm25_path, "rb") as f:
            return pickle.load(f)

    tokenized = [text.lower().split() for text in chunks_df["text"]]
    bm25 = BM25Okapi(tokenized)
    with open(bm25_path, "wb") as f:
        pickle.dump(bm25, f)
    return bm25


def load_dense_index(cache_dir: Path):
    """Load a precomputed FAISS index + chunks.parquet (no rebuilding)."""
    index = faiss.read_index(str(cache_dir / "faiss.index"))
    chunks_df = pd.read_parquet(cache_dir / "chunks.parquet")
    return index, chunks_df


def load_bm25(cache_dir: Path) -> BM25Okapi:
    with open(cache_dir / "bm25.pkl", "rb") as f:
        return pickle.load(f)


def load_docs(cache_dir: Path) -> pd.DataFrame:
    """Full document texts (doc_id, text) -- used to build sentence-window context."""
    return pd.read_parquet(cache_dir / "docs.parquet")