"""
RAG Retriever - Ricerca semantica nelle Meditazioni di Marco Aurelio

Questo modulo gestisce il retrieval semantico dai documenti indicizzati,
isolando la complessità di LangChain e fornendo un'API semplice.
"""
import os
import pickle
from typing import List, Optional

from langchain_chroma import Chroma
from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import LocalFileStore, EncoderBackedStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings


class RAGRetriever:
    """Gestisce il retrieval semantico dai documenti indicizzati"""

    def __init__(self, vectorstore_path: str, docstore_path: str, embedding_model: str):
        """
        Inizializza RAGRetriever.

        Args:
            vectorstore_path: Path al database Chroma
            docstore_path: Path al docstore locale
            embedding_model: Nome del modello HuggingFace per embeddings
        """
        self.vectorstore_path = vectorstore_path
        self.docstore_path = docstore_path
        self.embedding_model = embedding_model
        self.retriever = None
        self.available = False

    def initialize(self) -> bool:
        """
        Inizializza il retriever caricando il database RAG se esiste.

        Returns:
            True se inizializzazione ha successo, False altrimenti
        """
        if not os.path.exists(self.vectorstore_path) or not os.path.exists(self.docstore_path):
            print(f"[RAG] Database non trovato in: {self.vectorstore_path}")
            print("[RAG] Esegui indicizzazione con: python scripts/indicizza_documenti.py data/marco_aurelio.pdf")
            self.available = False
            return False

        try:
            # Inizializza embeddings (HuggingFace locale)
            print(f"[RAG] Caricamento embeddings: {self.embedding_model}")
            embeddings = HuggingFaceEmbeddings(model_name=self.embedding_model)

            # Carica vectorstore
            print(f"[RAG] Caricamento vectorstore da: {self.vectorstore_path}")
            vectorstore = Chroma(
                collection_name="split_parents",
                persist_directory=self.vectorstore_path,
                embedding_function=embeddings
            )

            # Carica docstore
            print(f"[RAG] Caricamento docstore da: {self.docstore_path}")
            fs = LocalFileStore(self.docstore_path)
            store = EncoderBackedStore(fs, lambda key: key, pickle.dumps, pickle.loads)

            # Crea ParentDocumentRetriever
            self.retriever = ParentDocumentRetriever(
                vectorstore=vectorstore,
                docstore=store,
                child_splitter=RecursiveCharacterTextSplitter(
                    chunk_size=400,
                    chunk_overlap=40
                ),
            )

            self.available = True
            print("[RAG] ✅ Retriever inizializzato con successo")
            return True

        except Exception as e:
            print(f"[RAG] ❌ Errore durante inizializzazione: {e}")
            self.available = False
            return False

    def search(self, query: str, max_results: int = 3) -> str:
        """
        Cerca nei documenti indicizzati.

        Args:
            query: Query di ricerca
            max_results: Numero massimo di risultati da ritornare

        Returns:
            Testo formattato con i chunk rilevanti trovati, oppure
            messaggio di errore se RAG non disponibile o nessun risultato
        """
        if not self.available or not self.retriever:
            return (
                "La memoria stoica non è disponibile. "
                "Esegui: python scripts/indicizza_documenti.py data/marco_aurelio.pdf"
            )

        try:
            # Invoke retriever
            docs = self.retriever.invoke(query)

            if not docs:
                return (
                    "Nessuna informazione trovata nelle Meditazioni per questa domanda. "
                    "Prova a riformulare la query."
                )

            # Deduplica contenuti
            seen = set()
            unique_contents = []
            for doc in docs[:max_results]:  # Limita a max_results
                content = doc.page_content or ""
                content_norm = " ".join(content.split()).lower()
                if content_norm and content_norm not in seen:
                    seen.add(content_norm)
                    unique_contents.append(content)

            if not unique_contents:
                return "Nessun contenuto rilevante trovato dopo deduplicazione."

            # Formatta output pulito
            result = "\n\n---\n\n".join(unique_contents)
            return result

        except Exception as e:
            print(f"[RAG] ❌ Errore durante search: {e}")
            return f"Errore durante la ricerca nelle Meditazioni: {e}"

    def is_available(self) -> bool:
        """Ritorna True se il retriever è disponibile e pronto all'uso"""
        return self.available