import os
import time
import gradio as gr
import numpy as np
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
from huggingface_hub import InferenceClient
from huggingface_hub.errors import HfHubHTTPError

# ============================================================
# CONFIG
# ============================================================

HF_TOKEN = os.getenv("HF_TOKEN", None)

# Free model compatible with new HF API
#LLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
#LLM_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
# LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.3" Does not work
# LLM_MODEL ="HuggingFaceH4/zephyr-7b-beta"
# LLM_MODEL = "NousResearch/Hermes-2-Pro-Mistral-7B" Does not work

# Were Great LLM
# LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
# LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"

# LLM model
# LLM_MODEL = "deepseek-ai/DeepSeek-V3"
LLM_MODEL = 'Qwen/Qwen2.5-14B-Instruct'

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
embedding_model = SentenceTransformer(EMBEDDING_MODEL)

# IMPORTANT: do NOT pass model=... here
llm = InferenceClient(token=HF_TOKEN)


# ============================================================
# SAFE LLM CALL (prevents 429 errors)
# ============================================================

def safe_chat_completion(llm, **kwargs):
    for attempt in range(5):
        try:
            return llm.chat.completions.create(**kwargs)
        except HfHubHTTPError as e:
            if "429" in str(e):
                sleep_time = 2 ** attempt
                print(f"[Rate limit] Retrying in {sleep_time}s...")
                time.sleep(sleep_time)
            else:
                raise
    raise RuntimeError("Too many retries after rate limits.")


# ============================================================
# PDF PROCESSING
# ============================================================

def read_pdf(file_obj):
    reader = PdfReader(file_obj)
    pages = []
    for page in reader.pages:
        text = page.extract_text()
        if text:
            pages.append(text)
    return "\n".join(pages)


def chunk_text(text, chunk_size=800, overlap=200):
    words = text.split()
    chunks = []
    start = 0

    while start < len(words):
        end = start + chunk_size
        chunk = " ".join(words[start:end])
        chunks.append(chunk)
        start += chunk_size - overlap

    return chunks


# ============================================================
# PER-FILE RETRIEVAL
# ============================================================

def retrieve_context_per_file(question, state, k=10):
    """
    For each file, retrieve top-k chunks separately,
    then merge all selected chunks into a single context.
    """
    if not state["files"]:
        return ""

    q_emb = embedding_model.encode([question], convert_to_numpy=True)[0]
    all_selected = []

    for filename, data in state["files"].items():
        chunks = data["chunks"]
        embeddings = data["embeddings"]

        if len(chunks) == 0:
            continue

        # Compute L2 distances to the query embedding
        distances = np.linalg.norm(embeddings - q_emb, axis=1)

        # Get top-k indices for this file
        top_k_idx = distances.argsort()[: min(k, len(chunks))]

        # Tag chunks with filename for transparency
        selected = [f"[{filename}]\n{chunks[i]}" for i in top_k_idx]
        all_selected.extend(selected)

    if not all_selected:
        return ""

    return "\n\n---\n\n".join(all_selected)


# ============================================================
# LLM ANSWERING
# ============================================================

def generate_answer(question, context):
    if context:
        prompt = (
            "You are a helpful assistant answering questions strictly from the provided context.\n"
            "If the answer is not in the context, say you don't know.\n\n"
            f"Context:\n{context}\n\n"
            f"Question: {question}\nAnswer:"
        )
    else:
        prompt = (
            "You are a helpful assistant. The user asked a question but no document context is available.\n\n"
            f"Question: {question}\nAnswer:"
        )

    response = safe_chat_completion(
        llm,
        model=LLM_MODEL,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1024,
        temperature=0.25,
        top_p=0.9,
        frequency_penalty=0.1,
        presence_penalty=0.0,
    )

    return response.choices[0].message["content"].strip()


# ============================================================
# GRADIO CALLBACKS
# ============================================================

def load_pdf(files, state):
    """
    files: current list from the UI (after add/remove).
    We rebuild state from scratch to stay in sync.

    state = {
        "files": {
            "filename.pdf": {
                "chunks": [...],
                "embeddings": np.array([...])
            },
            ...
        },
        "pdf_count": int,
        "total_chunks": int,
        "history": [...],
        "facts": [...]
    }
    """
    # If user cleared all files
    if not files:
        state["files"] = {}
        state["pdf_count"] = 0
        state["total_chunks"] = 0

        return (
            "No file uploaded.",
            state,
            gr.update(choices=[], value=None),
        )

    try:
        # Reset and rebuild from current files only
        state["files"] = {}
        state["pdf_count"] = 0
        state["total_chunks"] = 0

        # Deduplicate while preserving order
        unique_files = list(dict.fromkeys(files))

        for file in unique_files:
            filename = os.path.basename(str(file))

            text = read_pdf(file)
            chunks = chunk_text(text)

            if not chunks:
                continue

            embeddings = embedding_model.encode(chunks, convert_to_numpy=True)

            state["files"][filename] = {
                "chunks": chunks,
                "embeddings": embeddings,
            }

            state["pdf_count"] += 1
            state["total_chunks"] += len(chunks)

        status_msg = (
            f"{state['pdf_count']} PDF(s) loaded. "
            f"Total chunks: {state['total_chunks']}"
        )

        return (
            status_msg,
            state,
            gr.update(choices=list(state["files"].keys()), value=None),
        )

    except Exception as e:
        return (
            f"Error processing PDF: {e}",
            state,
            gr.update(choices=list(state["files"].keys()), value=None),
        )


def remove_pdf(filename, state):
    if filename not in state["files"]:
        return (
            f"{filename} not found.",
            state,
            gr.update(choices=list(state["files"].keys()), value=None),
        )

    # Remove the file entry
    state["files"].pop(filename, None)

    # Recompute counts
    pdf_count = len(state["files"])
    total_chunks = sum(len(data["chunks"]) for data in state["files"].values())

    state["pdf_count"] = pdf_count
    state["total_chunks"] = total_chunks

    status_msg = (
        f"Removed {filename}. "
        f"Total PDFs: {pdf_count}. Total chunks: {total_chunks}"
    )

    return (
        status_msg,
        state,
        gr.update(choices=list(state["files"].keys()), value=None),
    )


def answer_question(question, state):
    print("HF_TOKEN loaded:", HF_TOKEN is not None)

    if HF_TOKEN:
        print("HF_TOKEN length:", len(HF_TOKEN))
        print("HF_TOKEN preview:", HF_TOKEN[:6] + "..." + HF_TOKEN[-4:])
    else:
        print("HF_TOKEN is None or empty")

    question = question.strip()
    if not question:
        return "Please enter a question."

    if not state["files"]:
        return "Please upload at least one PDF first."

    # Simple persistent memory hook: "Remember that ..."
    lower_q = question.lower()
    if lower_q.startswith("remember that"):
        fact = question[len("remember that"):].strip(" .")
        if fact:
            state["facts"].append(fact)
            return f"Okay, I will remember that {fact}."
        else:
            return "What should I remember?"

    # RAG context (per-file balanced retrieval)
    rag_context = retrieve_context_per_file(question, state, k=10)

    # Conversation memory (last 5 turns)
    memory_context = ""
    if state["history"]:
        memory_context = "\n\n".join(
            [
                f"User: {h['user']}\nAssistant: {h['assistant']}"
                for h in state["history"][-5:]
            ]
        )

    # Persistent facts
    facts_context = ""
    if state["facts"]:
        facts_context = "User Facts:\n" + "\n".join(state["facts"]) + "\n\n"

    # Merge all contexts
    full_context_parts = []
    if facts_context:
        full_context_parts.append(facts_context)
    if memory_context:
        full_context_parts.append("Conversation Memory:\n" + memory_context)
    if rag_context:
        full_context_parts.append("Document Context:\n" + rag_context)

    full_context = "\n\n".join(full_context_parts)

    # Generate answer
    answer = generate_answer(question, full_context)

    # Save conversational memory
    state["history"].append({"user": question, "assistant": answer})

    return answer


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(title="PDF Q&A Chatbot") as demo:
    gr.Markdown("## 📄 PDF Q&A Chatbot\nUpload PDFs and ask questions about their content.")

    state = gr.State(
        {
            "files": {},          # filename -> {chunks, embeddings}
            "pdf_count": 0,
            "total_chunks": 0,
            "history": [],        # conversational memory
            "facts": [],          # persistent user facts
        }
    )

    with gr.Row():
        # LEFT COLUMN — PDF upload + removal
        with gr.Column(scale=1):
            pdf_input = gr.File(
                label="Upload PDF(s)",
                file_types=[".pdf"],
                file_count="multiple",  # proper multi-upload in Gradio 6
            )
            pdf_status = gr.Textbox(label="Status", interactive=False)

            file_selector = gr.Dropdown(
                label="Remove a PDF",
                choices=[],
                interactive=True,
            )
            remove_button = gr.Button("Remove Selected PDF")

        # RIGHT COLUMN — question + clear + answer
        with gr.Column(scale=2):
            with gr.Row():
                question = gr.Textbox(
                    label="Your question",
                    lines=2,
                    scale=3,
                )
                clear_button = gr.Button(
                    "Clear",
                    scale=1,
                    min_width=80,
                    elem_classes="clear-btn",
                )

            answer = gr.Textbox(label="Answer", lines=8)
            ask_button = gr.Button("Ask")

    # Bind events AFTER components exist
    pdf_input.change(
        load_pdf,
        inputs=[pdf_input, state],
        outputs=[pdf_status, state, file_selector],
    )

    remove_button.click(
        remove_pdf,
        inputs=[file_selector, state],
        outputs=[pdf_status, state, file_selector],
    )

    ask_button.click(
        answer_question,
        inputs=[question, state],
        outputs=answer,
    )



    clear_button.click(lambda: "", outputs=question)


if __name__ == "__main__":
    demo.launch(
        css=".clear-btn button { height: 100%; min-height: 52px; }"
    )