"""
rag.py — Core RAG engine: chunk → embed → index → search.
Works with Hebrew, Russian, and English out of the box.
"""

import json
import os
import pickle
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer

# ── Config ──────────────────────────────────────────────────
MODEL_NAME = "intfloat/multilingual-e5-base"  # great Hebrew support
INDEX_DIR = "index"
CHUNK_SIZE = 400       # characters per chunk
CHUNK_OVERLAP = 80     # overlap between chunks


class TextChunker:
    """Split long texts into overlapping chunks with metadata."""

    def __init__(self, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP):
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk_text(self, text: str, metadata: dict = None) -> list[dict]:
        """Split text into overlapping chunks, preserving sentence boundaries."""
        if metadata is None:
            metadata = {}

        # Split by paragraphs first, then by sentences
        paragraphs = [p.strip() for p in text.split("\n") if p.strip()]

        chunks = []
        current_chunk = ""

        for para in paragraphs:
            if len(current_chunk) + len(para) + 1 <= self.chunk_size:
                current_chunk += ("\n" + para if current_chunk else para)
            else:
                if current_chunk:
                    chunks.append({
                        "text": current_chunk,
                        **metadata,
                    })
                # Start new chunk with overlap from previous
                if self.overlap > 0 and current_chunk:
                    overlap_text = current_chunk[-self.overlap:]
                    current_chunk = overlap_text + "\n" + para
                else:
                    current_chunk = para

        if current_chunk:
            chunks.append({"text": current_chunk, **metadata})

        return chunks

    def chunk_documents(self, documents: list[dict]) -> list[dict]:
        """Chunk a list of extracted documents."""
        all_chunks = []
        for doc in documents:
            text = doc["text"]
            metadata = {k: v for k, v in doc.items() if k != "text"}
            chunks = self.chunk_text(text, metadata)
            all_chunks.extend(chunks)
        return all_chunks


class RAGIndex:
    """FAISS-based vector index with semantic search."""

    def __init__(self, model_name: str = MODEL_NAME):
        print(f"🔄 Loading embedding model: {model_name}")
        self.model = SentenceTransformer(model_name)
        self.dimension = self.model.get_sentence_embedding_dimension()
        self.index = None
        self.chunks = []

    def _embed(self, texts: list[str]) -> np.ndarray:
        """Embed texts with E5 prefix for better retrieval."""
        # multilingual-e5 uses "query: " and "passage: " prefixes
        return self.model.encode(texts, normalize_embeddings=True, show_progress_bar=True)

    def build(self, chunks: list[dict]):
        """Build FAISS index from text chunks."""
        self.chunks = chunks
        texts = [f"passage: {c['text']}" for c in chunks]

        print(f"🧮 Embedding {len(texts)} chunks...")
        embeddings = self._embed(texts)

        # Use IndexFlatIP for cosine similarity (embeddings are normalized)
        self.index = faiss.IndexFlatIP(self.dimension)
        self.index.add(embeddings.astype("float32"))
        print(f"✅ Index built: {self.index.ntotal} vectors, dim={self.dimension}")

    def search(self, query: str, top_k: int = 5) -> list[dict]:
        """Semantic search: returns top_k most relevant chunks."""
        query_vec = self._embed([f"query: {query}"])
        scores, indices = self.index.search(query_vec.astype("float32"), top_k)

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx == -1:
                continue
            chunk = self.chunks[idx].copy()
            chunk["score"] = float(score)
            results.append(chunk)
        return results

    def save(self, index_dir: str = INDEX_DIR):
        """Save index and chunks to disk."""
        os.makedirs(index_dir, exist_ok=True)
        faiss.write_index(self.index, os.path.join(index_dir, "faiss.index"))
        with open(os.path.join(index_dir, "chunks.pkl"), "wb") as f:
            pickle.dump(self.chunks, f)
        print(f"💾 Index saved to {index_dir}/")

    def load(self, index_dir: str = INDEX_DIR):
        """Load index and chunks from disk."""
        self.index = faiss.read_index(os.path.join(index_dir, "faiss.index"))
        with open(os.path.join(index_dir, "chunks.pkl"), "rb") as f:
            self.chunks = pickle.load(f)
        print(f"📂 Index loaded: {self.index.ntotal} vectors")


# ── CLI ─────────────────────────────────────────────────────
if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "build":
        # Build index from extracted texts
        with open("data/extracted_texts.json", "r", encoding="utf-8") as f:
            documents = json.load(f)

        chunker = TextChunker()
        chunks = chunker.chunk_documents(documents)
        print(f"📝 {len(chunks)} chunks created")

        rag = RAGIndex()
        rag.build(chunks)
        rag.save()

    elif len(sys.argv) > 1 and sys.argv[1] == "search":
        query = " ".join(sys.argv[2:])
        rag = RAGIndex()
        rag.load()
        results = rag.search(query)

        print(f"\n🔍 Query: {query}\n{'─' * 60}")
        for i, r in enumerate(results, 1):
            src = r.get("source", "?")
            score = r["score"]
            text = r["text"][:200]
            print(f"\n  [{i}] (score: {score:.3f}) [{src}]")
            print(f"      {text}...")

    else:
        print("Usage:")
        print("  python rag.py build           # Build index from data/extracted_texts.json")
        print("  python rag.py search <query>  # Search the index")