from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import List, Optional

from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS

logger = logging.getLogger(__name__)


@dataclass
class RAGConfig:
    chunk_size: int = 1000
    chunk_overlap: int = 100
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
    embedding_device: str = "cpu"
    top_k: int = 4
    normalize_embeddings: bool = True


@dataclass
class RetrievedContext:
    query: str
    chunks: List[Document]
    combined_text: str

    @property
    def source_files(self) -> List[str]:
        seen, out = set(), []
        for doc in self.chunks:
            s = doc.metadata.get("source", "unknown")
            if s not in seen:
                seen.add(s)
                out.append(s)
        return out

    @property
    def chunk_count(self) -> int:
        return len(self.chunks)


class RAGEngine:
    def __init__(self, config: Optional[RAGConfig] = None) -> None:
        self.config = config or RAGConfig()
        self._vector_store: Optional[FAISS] = None
        self._embeddings: Optional[HuggingFaceEmbeddings] = None
        self._splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.config.chunk_size,
            chunk_overlap=self.config.chunk_overlap,
            separators=["\n\n", "\n", ". ", " ", ""],
            add_start_index=True,
        )

    def build_index(self, documents: List[Document]) -> int:
        if not documents:
            raise ValueError("Cannot build index from an empty document list.")
        chunks = self._chunk(documents)
        self._vector_store = FAISS.from_documents(chunks, self._get_embeddings())
        return len(chunks)

    def retrieve(self, query: str, top_k: Optional[int] = None) -> RetrievedContext:
        if self._vector_store is None:
            raise RuntimeError("Call build_index() before retrieve().")
        docs = self._vector_store.similarity_search(query, k=top_k or self.config.top_k)
        return RetrievedContext(query=query, chunks=docs, combined_text=self._format_context(docs))

    def add_documents(self, documents: List[Document]) -> int:
        if self._vector_store is None:
            return self.build_index(documents)
        chunks = self._chunk(documents)
        self._vector_store.add_documents(chunks, embedding=self._get_embeddings())
        return len(chunks)

    def reset(self) -> None:
        self._vector_store = None

    def _chunk(self, documents: List[Document]) -> List[Document]:
        chunks = self._splitter.split_documents(documents)
        for idx, chunk in enumerate(chunks):
            chunk.metadata["chunk_id"] = idx
        return chunks

    def _get_embeddings(self) -> HuggingFaceEmbeddings:
        if self._embeddings is None:
            self._embeddings = HuggingFaceEmbeddings(
                model_name=self.config.embedding_model,
                model_kwargs={"device": self.config.embedding_device},
                encode_kwargs={"normalize_embeddings": self.config.normalize_embeddings},
            )
        return self._embeddings

    @staticmethod
    def _format_context(docs: List[Document]) -> str:
        separator = "─" * 60
        parts = []
        for i, doc in enumerate(docs, start=1):
            meta = doc.metadata
            header = (
                f"[CHUNK {i} | source: {meta.get('source', 'unknown')}, "
                f"page: {meta.get('page', 'N/A')}, "
                f"chunk_id: {meta.get('chunk_id', 'N/A')}]"
            )
            parts.append(f"{header}\n{doc.page_content.strip()}")
        return f"\n{separator}\n".join(parts)