import os
import gc
from pypdf import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings # LOKAL!
from langchain_community.vectorstores import FAISS
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_community.llms import LlamaCpp
from huggingface_hub import hf_hub_download
# WICHTIG: Import für die Metadaten-Struktur
from langchain_core.documents import Document

# --- Konfiguration des lokalen Modells ---
# Wechsel zurück auf 8B für höhere Präzision bei komplexen Analysen
MODEL_REPO_ID = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF"
MODEL_FILENAME = "Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf"
# Lokaler Cache-Pfad im Container
MODEL_CACHE_DIR = os.environ.get("MODEL_CACHE_DIR", "/app/models")

# Globaler LLM-Slot - wird nur einmal geladen und dann wiederverwendet
_llm_instance: LlamaCpp | None = None


def _get_local_model_path() -> str:
    """
    Gibt den lokalen Pfad zur .gguf-Datei zurück.
    Ist die Datei noch nicht vorhanden, wird sie automatisch von Hugging Face
    in MODEL_CACHE_DIR heruntergeladen.
    """
    os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
    local_path = os.path.join(MODEL_CACHE_DIR, MODEL_FILENAME)
    
    # HF_TOKEN aus den Hugging Face Secrets laden (wichtig für Llama 3.2)
    hf_token = os.getenv("HF_TOKEN")

    if not os.path.exists(local_path):
        print(f"[RAG] Modell nicht gefunden – lade '{MODEL_FILENAME}' herunter …")
        local_path = hf_hub_download(
            repo_id=MODEL_REPO_ID,
            filename=MODEL_FILENAME,
            local_dir=MODEL_CACHE_DIR,
            local_dir_use_symlinks=False,   # echte Kopie statt Symlink → sicherer im Container
            token=hf_token                  # Auth-Token für gated Models übergeben
        )
        print(f"[RAG] Modell erfolgreich gespeichert unter: {local_path}")
    else:
        print(f"[RAG] Modell bereits im Cache: {local_path}")

    return local_path


def _get_llm() -> LlamaCpp:
    """
    Gibt die singleton-LlamaCpp-Instanz zurück.
    Beim ersten Aufruf wird das Modell geladen; danach wird dieselbe Instanz
    wiederverwendet (spart RAM und Ladezeit).
    """
    global _llm_instance
    if _llm_instance is None:
        model_path = _get_local_model_path()
        print("[RAG] Lade Modell in den Arbeitsspeicher …")
        _llm_instance = LlamaCpp(
            model_path=model_path,
            temperature=0.0,
            max_tokens=1024,     # ERHÖHT: Damit Antworten nicht mehr abgeschnitten werden
            n_ctx=16384,        # Kontextfenster (in Token) - 16384 für Llama 3.2 <-- Komma entfernt!
            n_batch=512,         # Token-Batch-Größe beim Prompt-Verarbeiten
            n_gpu_layers=-1,     # -1 = ALLES in den VRAM (CUDA); 0 = reine CPU <--UPDATE!
            streaming=True,      # Live-Streaming aktivieren für sofortige UI-Updates!
            stop=["<|eot_id|>", "<|end_of_text|>", "Human:", "user:"], # <--- DIE NOTBREMSE!
            verbose=False,
        )
        print("[RAG] Modell geladen und bereit.")
    return _llm_instance


def unload_llm() -> None:
    """
    Gibt das LLM explizit aus dem RAM frei (nützlich in ressourcenbeschränkten
    Containern, wenn das Modell nach einer Session nicht mehr gebraucht wird).
    """
    global _llm_instance
    if _llm_instance is not None:
        del _llm_instance
        _llm_instance = None
        gc.collect()
        print("[RAG] LLM-Instanz aus dem Arbeitsspeicher entfernt.")


# ---------------------------------------------------------------------------
# 
# ---------------------------------------------------------------------------

def process_pdf_to_vectorstore(uploaded_file):
    """
    Nimmt das PDF, liest den Text, zerteilt ihn und speichert ihn als Embeddings.
    Gibt den fertigen Vectorstore zurück.
    """
    # A. Text aus dem PDF extrahieren (NEU: Mit Metadaten für Seitenzahlen)
    pdf_reader = PdfReader(uploaded_file)
    docs = []
    
    # Nutzt enumerate, um die Seiten zu zählen (i + 1 für echte Seitenzahlen)
    for i, page in enumerate(pdf_reader.pages):
        page_text = page.extract_text()
        if page_text:
            # Erstellt Dokument-Objekt mit Metadaten-Tag für die Seite
            docs.append(Document(page_content=page_text, metadata={"page": i + 1}))
        
    # B. Text in kleine "Häppchen" (Chunks) zerteilen
    # LLMs können nicht endlos viel Text auf einmal lesen, daher zerlegen wir ihn.
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,      # Jeder Chunk hat ca. 1000 Zeichen
        chunk_overlap=200     # 200 Zeichen Überlappung, damit kein Satz in der Mitte abreißt
    )
    # Nutzt split_documents statt split_text, um Seitenzahlen beizubehalten
    text_chunks = text_splitter.split_documents(docs)
    
    # C. EIGENE INFRASTRUKTUR: Lokale Embeddings
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    # Nutzt from_documents statt from_texts
    vectorstore = FAISS.from_documents(documents=text_chunks, embedding=embeddings)
    
    return vectorstore

def get_rag_answer(vectorstore, user_question):
    """
    Nimmt die Frage des Users, sucht im Vectorstore nach passendem Kontext 
    und generiert mit dem LLM eine Antwort.
    """
    # A. LOKALES LLM: GGUF-Modell via llama-cpp-python (Singleton)
    llm = _get_llm()
    
    # B. Llama 3.2 Instruct Prompt Format
    # Die Steuerzeichen nutzen, auf die Llama 3.2 trainiert wurde!
    # HINWEIS: <|begin_of_text|> entfernt, um doppelte Tokens zu vermeiden
    llama3_template = """<|start_header_id|>system<|end_header_id|>

Du bist PDF'i, eine schlaue, hilfsbereite und moderne Technik-Eule. Du hast das hochgeladene Dokument mit deinen scharfen Eulenaugen analysiert. 
Beantworte die Fragen des Nutzers sympathisch und präzise, aber nutze NUR den folgenden Kontext. Wenn die Antwort nicht im Text steht, gib ehrlich zu, 
dass du sie in den vorliegenden Seiten nicht erblicken kannst. Erfinde niemals Informationen dazu.

Kontext:
{context}<|eot_id|><|start_header_id|>user<|end_header_id|>

{input}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""
    prompt_template = PromptTemplate.from_template(llama3_template)
    # C. Die RAG-Kette zusammenbauen
    question_answer_chain = create_stuff_documents_chain(llm, prompt_template)
    
    # Holt die relevantesten Chunks aus der Datenbank
    retriever = vectorstore.as_retriever(search_kwargs={"k": 30})
    
    # Verbindet Retriever und LLM zur finalen Kette
    rag_chain = create_retrieval_chain(retriever, question_answer_chain)
    
    # D. DIE NEUE STREAMING LOGIK
    # Statt invoke() nutzen wir stream() und feuern jedes Wort per 'yield' an die UI
    for chunk in rag_chain.stream({"input": user_question}):
        if "answer" in chunk:
            yield chunk["answer"]