"""
Local RAG setup with LangChain, Ollama/OpenAI, and FAISS
Minimal dependencies, simple code
"""
import json
import os
import re
import time
from pathlib import Path

import numpy as np

from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_ollama import ChatOllama
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter

import docling_loader

# Pipeline tuning: used when adding documents (rebuild index after changing).
# Characters, not tokens: 2000 characters of German prose are roughly 550 tokens.
# With the reranker, 2000/400 beat 1500/300 on both Kafka goldsets (eval/matrix_2000.out).
CHUNK_SIZE = 2000
CHUNK_OVERLAP = 400

# Multilingual embedding model with a long input window (8192 tokens), so chunks
# are never truncated. Changing it requires rebuilding the index.
EMBEDDING_MODEL = "BAAI/bge-m3"

# How PDFs are read: "pypdf" (PyPDFLoader, one unit per page) or "docling"
# (structured output from prepare_corpus.py, cleaned by docling_loader.py).
PDF_LOADER = "pypdf"
# Where prepare_corpus.py put the Docling output for the PDFs being indexed.
PROCESSED_DIR = "corpus_processed/data_split"
# Prefix every chunk with author and title from the file name (Docling: plus section),
# so questions like "according to Lehmann ..." can match. See docling_loader.chunk_header.
# Needed for Docling (which drops running heads); neutral to slightly harmful for pypdf,
# whose pages carry the author in the running head anyway.
CHUNK_HEADERS = False

# Retrieval / query tuning
RETRIEVAL_K = 5

# Hybrid search: besides the vector search, a keyword search (BM25 on German word stems)
# finds names, titles and quotations that embeddings blur.
#   "dense"  vector search alone
#   "hybrid" both lists merged with Reciprocal Rank Fusion, which also sets the order
#   "union"  BM25 only adds candidates (BM25_CANDIDATES) to the vector candidates;
#            the order comes from the vector search or, if enabled, from the reranker
RETRIEVAL_MODE = "hybrid"
# 25 instead of 50 halves the reranker time and was not worse on either Kafka goldset
# (pypdf 2000/400, R@5 human .70 -> .73, synthetic .91 -> .98).
CANDIDATES_K = 25         # candidates per search before merging and reranking
BM25_CANDIDATES = 20      # extra keyword candidates in "union" mode
RRF_K = 60                # standard constant of Reciprocal Rank Fusion

# Cross-encoder that rereads question and candidate together and re-sorts the candidates.
# None switches reranking off.
RERANKER_MODEL = "BAAI/bge-reranker-v2-m3"

# Appended to the embedding query for FAISS (does not affect the answer-generation prompt).
# Left empty on purpose: embedding models do not follow instructions, and on the Kafka
# corpus the former text ("Do not return a list of references ...") lowered Recall@5
# from 0.67 to 0.53 (eval_retrieval.py --query-instructions). Measure before re-enabling.
QUERY_ADDITIONAL_INSTRUCTIONS = ""

# Written next to every index; an index built with another embedding model is not loaded.
INDEX_INFO_FILE = "index_info.json"

# The corpus search is offered to the LLM as a tool; the LLM decides per message whether
# to search (content questions) or to answer directly (greetings, thanks, questions about
# the conversation). It also writes the search query itself, using the chat history.
SEARCH_TOOL_NAME = "search_kafka_corpus"
SEARCH_TOOL_DESCRIPTION = (
    "Durchsucht die Kafka-Bibliothek (Kafkas Texte und die Forschungsliteratur) und liefert die "
    "passendsten Textauszüge mit Quelle und Seite. Die Suchanfrage sollte eine eigenständige, "
    "präzise Frage oder Stichwortfolge auf Deutsch sein."
)
MAX_TOOL_ROUNDS = 3       # searches per message before the LLM must answer

# System prompt for every chat turn (German, like the corpus and the questions).
SYSTEM_PROMPT = f"""Du beantwortest Fragen zu Franz Kafka, seinen Werken, seinem Leben und der Forschung.
Du hast ein Tool, {SEARCH_TOOL_NAME}, das eine Bibliothek aus Kafkas Texten und Forschungsliteratur durchsucht.
- Für jede inhaltliche Frage zu Kafka, seinen Werken, Figuren, seiner Biografie, Sprache oder zur Forschung: Rufe zuerst das Werkzeug auf und stütze die Antwort auf die gefundenen Auszüge. Antworte nicht aus dem Gedächtnis.
- Bei Begrüßung, Dank, Smalltalk oder Fragen zum Gespräch selbst: Antworte direkt, ohne Tool.
- Formuliere die Suchanfrage so, dass sie ohne den Gesprächsverlauf verständlich ist.
- Finden sich in den Auszügen keine Belege, sage das offen, statt zu raten.
- Nenne bei Aussagen aus den Auszügen die Quelle (Autor, Seite)."""


# Cross-encoders by model name, shared by all LocalRAG instances in a process (one copy in memory)
_RERANKERS = {}
# Embedding models, shared the same way
_EMBEDDINGS = {}
# BM25 keyword indexes by (vector store path, number of chunks), shared by instances on the same index
_BM25_INDEXES = {}


def load_embeddings(model_name):
    """Load an embedding model once per process and return the shared instance."""
    if model_name not in _EMBEDDINGS:
        print("Loading embeddings model...")
        # Normalized vectors: FAISS's L2 distance then ranks exactly like cosine similarity.
        _EMBEDDINGS[model_name] = HuggingFaceEmbeddings(
            model_name=model_name,
            encode_kwargs={"normalize_embeddings": True},
            query_encode_kwargs={"normalize_embeddings": True},
        )
    return _EMBEDDINGS[model_name]


def load_reranker(model_name):
    """Load a cross-encoder once per process and return the shared instance."""
    if model_name not in _RERANKERS:
        import torch
        from sentence_transformers import CrossEncoder
        # Half precision on the GPU (Apple MPS or CUDA): about 15-25 % faster, same rankings
        # on both Kafka goldsets. On the CPU fp16 is slow, so keep fp32 there.
        gpu = torch.backends.mps.is_available() or torch.cuda.is_available()
        model_kwargs = {"torch_dtype": torch.float16} if gpu else {}
        print(f"Loading reranker {model_name} ({'fp16' if gpu else 'fp32'})...")
        reranker = CrossEncoder(model_name, max_length=1024, model_kwargs=model_kwargs)
        reranker.predict([("Probe", "Probe")])  # warm-up, so the first real query is not slower
        _RERANKERS[model_name] = reranker
    return _RERANKERS[model_name]


class LocalRAG:
    def __init__(
        self,
        vectorstore_path="./vectorstore",
        llm_provider="ollama",
        ollama_model="gemma4:12b",
        openai_model="gpt-6-luna",
        realtime_model="gpt-realtime-2.1-mini",
        ollama_base_url="http://localhost:11434",
        pdf_loader=PDF_LOADER,
        processed_dir=PROCESSED_DIR,
        retrieval_mode=RETRIEVAL_MODE,
        reranker_model=RERANKER_MODEL,
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        chunk_headers=CHUNK_HEADERS,
    ):
        """Initialize local RAG system. llm_provider: 'ollama', 'openai' or 'realtime'
        (OpenAI Realtime API, see realtime_chat.py).
        pdf_loader: 'pypdf' or 'docling' (see PDF_LOADER)."""
        self.vectorstore_path = vectorstore_path
        self.llm_provider = llm_provider
        self.pdf_loader = pdf_loader
        self.processed_dir = Path(processed_dir)
        self.chunk_size, self.chunk_overlap, self.chunk_headers = chunk_size, chunk_overlap, chunk_headers
        self.retrieval_mode = retrieval_mode
        self.reranker_model = reranker_model
        self._bm25 = None          # built on first use from the documents in the index
        self._reranker = None      # loaded on first use, or up front via preload()
        self.last_rerank_seconds = 0.0  # time of the last reranking, for the timing trace

        # Embeddings (one shared copy per process)
        self.embeddings = load_embeddings(EMBEDDING_MODEL)

        # Text splitter (used when adding documents; tune CHUNK_SIZE / CHUNK_OVERLAP above)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )
        print(f"Text splitter: chunk_size={chunk_size}, chunk_overlap={chunk_overlap}")

        # LLM (Ollama, OpenAI or OpenAI Realtime)
        if llm_provider in ("openai", "realtime"):
            api_key = os.environ.get("OPENAI_API_KEY")
            if not api_key:
                raise ValueError(
                    f"OPENAI_API_KEY environment variable is required when llm_provider='{llm_provider}'"
                )
        if llm_provider == "realtime":
            from realtime_chat import ChatOpenAIRealtime
            print(f"Using OpenAI Realtime (model: {realtime_model})...")
            self.llm = ChatOpenAIRealtime(model=realtime_model, api_key=api_key)
        elif llm_provider == "openai":
            print(f"Using OpenAI (model: {openai_model})...")
            # Responses API: GPT-6 models accept function tools on Chat Completions only with
            # reasoning_effort="none"; with "low" the search tool needs the Responses API.
            self.llm = ChatOpenAI(model=openai_model, api_key=api_key, reasoning_effort="low",
                                  use_responses_api=True)
        else:
            print(f"Using Ollama (model: {ollama_model})...")
            self.llm = ChatOllama(
                model=ollama_model,
                base_url=ollama_base_url
            )

        # Vector store (load if exists, otherwise None)
        self.vectorstore = None
        self._load_vectorstore()
    
    def _index_info(self):
        """Settings that must match between the index on disk and this code."""
        return {
            "embedding_model": EMBEDDING_MODEL,
            "chunk_size": self.chunk_size,
            "chunk_overlap": self.chunk_overlap,
            "pdf_loader": self.pdf_loader,
            "chunk_headers": self.chunk_headers,
        }

    def _index_matches(self, index_file):
        """False if the index was built with a different embedding model than EMBEDDING_MODEL."""
        info_file = os.path.join(self.vectorstore_path, INDEX_INFO_FILE)
        if os.path.exists(info_file):
            with open(info_file, encoding="utf-8") as f:
                built_with = json.load(f).get("embedding_model")
            if built_with != EMBEDDING_MODEL:
                print(f"Index in {self.vectorstore_path} was built with {built_with}, "
                      f"but EMBEDDING_MODEL is {EMBEDDING_MODEL}. Rebuild it (add_pdfs.py).")
                return False
            return True
        # Older index without info file: at least the vector dimensions must agree
        import faiss
        index_dim = faiss.read_index(index_file).d
        model_dim = len(self.embeddings.embed_query("dimension check"))
        if index_dim != model_dim:
            print(f"Index in {self.vectorstore_path} has {index_dim} dimensions, {EMBEDDING_MODEL} "
                  f"produces {model_dim}. It was built with another model; rebuild it (add_pdfs.py).")
            return False
        return True

    def _load_vectorstore(self):
        """Load existing vector store if available"""
        index_file = os.path.join(self.vectorstore_path, "index.faiss")
        if os.path.exists(index_file) and self._index_matches(index_file):
            try:
                self.vectorstore = FAISS.load_local(
                    self.vectorstore_path,
                    self.embeddings,
                    allow_dangerous_deserialization=True
                )
                print(f"Loaded existing vector store from {self.vectorstore_path}")
            except Exception as e:
                print(f"Could not load vector store: {e}")
                self.vectorstore = None
    
    def add_documents(self, file_paths):
        """Add documents to the vector store"""
        print(f"\nLoading {len(file_paths)} document(s)...")
        all_docs = []
        
        for file_path in file_paths:
            path = Path(file_path)
            if not path.exists():
                print(f"Warning: {file_path} not found, skipping")
                continue
            
            # Docling: already structured and cleaned, the loader returns finished chunks
            if path.suffix.lower() == '.pdf' and self.pdf_loader == "docling":
                processed = self.processed_dir / f"{path.stem}.json"
                if processed.exists():
                    chunks = docling_loader.load(processed, str(path), self.text_splitter,
                                                 header=self.chunk_headers)
                    all_docs.extend(chunks)
                    print(f"  - {path.name}: {len(chunks)} chunks (docling)")
                    continue
                print(f"Warning: no Docling output for {path.name} in {self.processed_dir}, using PyPDFLoader")

            # Load document
            if path.suffix.lower() == '.pdf':
                loader = PyPDFLoader(str(path))
            elif path.suffix.lower() in ['.txt', '.md']:
                loader = TextLoader(str(path))
            else:
                print(f"Warning: Unsupported file type {path.suffix}, skipping")
                continue
            
            docs = loader.load()
            chunks = self.text_splitter.split_documents(docs)
            if self.chunk_headers:
                header = docling_loader.chunk_header(str(path))
                for chunk in chunks:
                    chunk.page_content = header + chunk.page_content
            all_docs.extend(chunks)
            print(f"  - {path.name}: {len(chunks)} chunks")
        
        if not all_docs:
            print("No documents loaded!")
            return
        
        # Create or update vector store
        print(f"\nCreating embeddings for {len(all_docs)} chunks...")
        if self.vectorstore is None:
            self.vectorstore = FAISS.from_documents(all_docs, self.embeddings)
        else:
            new_store = FAISS.from_documents(all_docs, self.embeddings)
            self.vectorstore.merge_from(new_store)
        
        # Save
        os.makedirs(self.vectorstore_path, exist_ok=True)
        self.vectorstore.save_local(self.vectorstore_path)
        with open(os.path.join(self.vectorstore_path, INDEX_INFO_FILE), "w", encoding="utf-8") as f:
            json.dump(self._index_info(), f, indent=2)
        print(f"Vector store saved to {self.vectorstore_path}")
    
    def list_documents(self):
        """List all documents in the vector store"""
        if self.vectorstore is None:
            print("No documents in vector store.")
            return []
        
        # Get all documents from the vector store
        # We'll retrieve a large number to get all documents
        all_docs = self.vectorstore.similarity_search("", k=10000)  # Large k to get all
        
        # Extract unique document sources from metadata
        documents = {}
        for doc in all_docs:
            source = doc.metadata.get('source', 'Unknown')
            if source not in documents:
                documents[source] = {
                    'source': source,
                    'chunks': 0,
                    'page': doc.metadata.get('page', None)
                }
            documents[source]['chunks'] += 1
        
        # Convert to list and sort
        doc_list = list(documents.values())
        doc_list.sort(key=lambda x: x['source'])
        
        print(f"\nDocuments in vector store ({len(doc_list)} unique documents):")
        print("-" * 60)
        for doc_info in doc_list:
            print(f"  - {doc_info['source']}")
            print(f"    Chunks: {doc_info['chunks']}")
            if doc_info['page'] is not None:
                print(f"    Page: {doc_info['page']}")
        
        return doc_list
    
    def _docs_to_retrieved(self, docs):
        """Convert document list to retrieved chunks format for API."""
        return [
            {
                "content": doc.page_content,
                "source": doc.metadata.get("source", ""),
                "page": doc.metadata.get("page"),
            }
            for doc in docs
        ]

    def _docs_scores_to_retrieved(self, docs_with_scores):
        """Convert (Document, score) list to retrieved chunks format with score. FAISS returns L2 distance (lower = more similar)."""
        return [
            {
                "content": doc.page_content,
                "source": doc.metadata.get("source", ""),
                "page": doc.metadata.get("page"),
                "score": float(score),
            }
            for doc, score in docs_with_scores
        ]

    # ---------- retrieval: vector search, keyword search, reranking ----------

    _stemmer = None
    _stems = {}  # word -> stem; the corpus has far fewer distinct words than word occurrences

    @classmethod
    def _bm25_tokens(cls, text):
        """Lowercase German word stems, so "ohnmächtigen" matches "Ohnmacht"."""
        if cls._stemmer is None:
            import snowballstemmer
            cls._stemmer = snowballstemmer.stemmer("german")
        stems = cls._stems
        tokens = []
        for word in re.findall(r"[a-zäöüß0-9]{2,}", text.lower()):
            stem = stems.get(word)
            if stem is None:
                stem = stems[word] = cls._stemmer.stemWord(word)
            tokens.append(stem)
        return tokens

    def _bm25_index(self):
        if self._bm25 is None:
            key = (self.vectorstore_path, self.vectorstore.index.ntotal)
            if key not in _BM25_INDEXES:
                from rank_bm25 import BM25Okapi
                started = time.perf_counter()
                docs = list(self.vectorstore.docstore._dict.values())
                _BM25_INDEXES[key] = (BM25Okapi([self._bm25_tokens(d.page_content) for d in docs]), docs)
                print(f"BM25 index for {self.vectorstore_path}: {len(docs)} chunks, "
                      f"{time.perf_counter() - started:.1f}s")
            self._bm25 = _BM25_INDEXES[key]
        return self._bm25

    def retrieve(self, query, k=None):
        """The k best chunks for a query as [(Document, score)].

        dense:  vector search only (score = L2 distance, lower is better)
        hybrid: vector search + BM25, merged by Reciprocal Rank Fusion (higher is better)
        union:  vector candidates plus BM25_CANDIDATES keyword candidates, vector order first
        With a reranker, the merged candidates are re-sorted by the cross-encoder
        (score = relevance, higher is better)."""
        k = k or RETRIEVAL_K
        pool = max(k, CANDIDATES_K) if (self.retrieval_mode != "dense" or self.reranker_model) else k
        dense = self.vectorstore.similarity_search_with_score(query, k=pool)
        key = lambda doc: (doc.metadata.get("source"), doc.metadata.get("page"), doc.page_content[:200])
        if self.retrieval_mode in ("hybrid", "union"):
            bm25, docs = self._bm25_index()
            scores = bm25.get_scores(self._bm25_tokens(query))
            keyword = [docs[i] for i in np.argsort(scores)[::-1][:pool]]
        if self.retrieval_mode == "union":
            seen = {key(d) for d, _ in dense}
            extra = [d for d in keyword if key(d) not in seen][:BM25_CANDIDATES]
            candidates = [[d, s] for d, s in dense] + [[d, float("nan")] for d in extra]
        elif self.retrieval_mode == "hybrid":
            fused = {}
            for ranking in ([d for d, _ in dense], keyword):
                for rank, doc in enumerate(ranking):
                    entry = fused.setdefault(key(doc), [doc, 0.0])
                    entry[1] += 1 / (RRF_K + rank + 1)
            candidates = sorted(fused.values(), key=lambda e: e[1], reverse=True)[:pool]
        else:
            candidates = [[d, s] for d, s in dense]
        if self.reranker_model and candidates:
            if self._reranker is None:
                self._reranker = load_reranker(self.reranker_model)
            started = time.perf_counter()
            relevance = self._reranker.predict([(query, d.page_content) for d, _ in candidates])
            self.last_rerank_seconds = time.perf_counter() - started
            candidates = [[d, float(r)] for (d, _), r in zip(candidates, relevance)]
            candidates.sort(key=lambda e: e[1], reverse=True)
        return [(d, s) for d, s in candidates[:k]]

    def preload(self):
        """Load the reranker and build the BM25 index now instead of on the first search
        (e.g. at server start)."""
        if self.reranker_model and self._reranker is None:
            self._reranker = load_reranker(self.reranker_model)
        if self.retrieval_mode != "dense" and self.vectorstore is not None:
            self._bm25_index()

    def query(self, question, k=None):
        """Query the RAG system (no conversation history). Returns dict with 'answer' and 'retrieved'."""
        if k is None:
            k = RETRIEVAL_K
        return self.query_with_history(question, chat_history=[], k=k)

    @staticmethod
    def _text(message):
        """Plain text of an LLM reply (some models return a list of content blocks)."""
        content = getattr(message, "content", message)
        if isinstance(content, list):
            content = "".join(part.get("text", "") if isinstance(part, dict) else str(part) for part in content)
        return (content or "").strip()

    def query_with_history(self, question, chat_history=None, k=None, system_prompt=None,
                           tool_description=None):
        """Answer a chat message. The LLM decides whether to search the corpus: the search is
        offered as a tool (SEARCH_TOOL_NAME), which the LLM calls with its own query, up to
        MAX_TOOL_ROUNDS times. system_prompt defaults to SYSTEM_PROMPT and tool_description to
        SEARCH_TOOL_DESCRIPTION (see personas.py for others).
        Returns dict with 'answer', 'retrieved' (chunks the LLM saw, with content, source, page,
        score) and 'searches' (the queries it used)."""
        from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
        from langchain_core.tools import StructuredTool

        k = k or RETRIEVAL_K
        if self.vectorstore is None:
            return {"answer": "Error: No documents loaded. Please add documents first.",
                    "retrieved": [], "searches": []}

        retrieved, searches = [], []

        def search_kafka_corpus(query: str) -> str:
            retrieval_query = query + (f"\n\n{QUERY_ADDITIONAL_INSTRUCTIONS}" if QUERY_ADDITIONAL_INSTRUCTIONS else "")
            started = time.perf_counter()
            hits = self.retrieve(retrieval_query, k=k)
            seconds = time.perf_counter() - started
            searches.append(query)
            retrieved.extend(self._docs_scores_to_retrieved(hits))
            print(f"[RAG] Tool search {query!r}: {len(hits)} chunk(s)")
            print(f"[Zeit] Suche {seconds:.2f}s (davon Reranker {self.last_rerank_seconds:.2f}s)")
            if not hits:
                return "Keine passenden Textstellen gefunden."
            return "\n\n".join(
                f"[{i}] {docling_loader.source_label(Path(d.metadata.get('source', '')).name)}"
                f"{', S. ' + str(d.metadata['page'] + 1) if d.metadata.get('page') is not None else ''}\n"
                f"{d.page_content.strip()}"
                for i, (d, _) in enumerate(hits, start=1))

        tool = StructuredTool.from_function(func=search_kafka_corpus, name=SEARCH_TOOL_NAME,
                                            description=tool_description or SEARCH_TOOL_DESCRIPTION)
        llm_with_tools = self.llm.bind_tools([tool])

        messages = [SystemMessage(system_prompt or SYSTEM_PROMPT)]
        for turn in chat_history or []:
            content = (turn.get("content") or "").strip()
            if turn.get("role") == "user":
                messages.append(HumanMessage(content))
            elif turn.get("role") == "assistant":
                messages.append(AIMessage(content))
        messages.append(HumanMessage(question))
        print(f"[RAG] User question: {question!r}")

        started_total = time.perf_counter()
        answer = None
        for round_no in range(1, MAX_TOOL_ROUNDS + 1):
            started = time.perf_counter()
            reply = llm_with_tools.invoke(messages)
            print(f"[Zeit] LLM-Aufruf {round_no}: {time.perf_counter() - started:.2f}s "
                  f"({'Suche angefordert' if reply.tool_calls else 'Antwort'})")
            messages.append(reply)
            if not reply.tool_calls:
                answer = self._text(reply)
                break
            for n, call in enumerate(reply.tool_calls):
                result = tool.invoke(call["args"])
                messages.append(ToolMessage(result, tool_call_id=call.get("id") or f"call_{len(searches)}_{n}"))
        if answer is None:
            # Search budget used up: answer with what was found, without further tool calls
            started = time.perf_counter()
            answer = self._text(self.llm.invoke(messages + [HumanMessage(
                "Beantworte jetzt die Frage mit den gefundenen Auszügen, ohne weitere Suche.")]))
            print(f"[Zeit] LLM-Aufruf (Antwort erzwungen): {time.perf_counter() - started:.2f}s")
        print(f"[Zeit] Gesamt {time.perf_counter() - started_total:.2f}s")
        print(f"[RAG] Searches: {searches or 'keine'}")
        return {"answer": answer, "retrieved": retrieved, "searches": searches}


def main():
    """Example usage"""
    print("=" * 60)
    print("Local RAG with LangChain, Ollama/OpenAI, and FAISS")
    print("=" * 60)
    
    # Initialize
    rag = LocalRAG(ollama_model="gemma4:12b")
    
    # Add documents (uncomment and add your file paths)
    # rag.add_documents([
    #     "data/dok1.pdf",
    #     "data/dok2.pdf",
    #     "data/dok3.pdf"
    # ])
    
    # List documents
    rag.list_documents()
    
    # Query
    question = "What do you knowabout modality for perceived message perception?"
    result = rag.query(question)
    print(f"\nQuestion: {question}")
    print(f"Answer: {result['answer']}")
    if result.get("retrieved"):
        print(f"Retrieved {len(result['retrieved'])} chunks")
    
    # print("\nSetup complete! Uncomment the code above to add documents and query.")


if __name__ == "__main__":
    main()