import re
import faiss
import numpy as np
import pandas as pd
import torch
import gradio as gr

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM

# -----------------------------
# Embedding Model
# -----------------------------
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# -----------------------------
# Load Dataset
# -----------------------------
docs_df = pd.read_pickle("docs.pkl")

embeddings = np.array(docs_df["embeddings"].tolist()).astype("float32")
faiss.normalize_L2(embeddings)

dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
index.add(embeddings)

# -----------------------------
# FORCE SAFE MODEL (IMPORTANT)
# -----------------------------
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

# -----------------------------
# Preprocessing
# -----------------------------
def preprocess_text(text):
    text = text.lower()
    text = text.replace("\n", " ").replace("\t", " ")
    text = re.sub(r"[^\w\s.,;:>-]", " ", text)
    return " ".join(text.split()).strip()

# -----------------------------
# Retrieval
# -----------------------------
def retrieve_docs(query, k=3):
    query_embedding = embedding_model.encode([query])[0].astype("float32")
    faiss.normalize_L2(query_embedding.reshape(1, -1))

    distances, indices = index.search(np.array([query_embedding]), k)

    results = docs_df.iloc[indices[0]].copy()
    results["score"] = distances[0]

    return results

# -----------------------------
# Prompt Builder
# -----------------------------
def build_prompt(query, history, context):
    history_text = ""

    for user, bot in history[-3:]:
        history_text += f"User: {user}\nAssistant: {bot}\n"

    return f"""
You are a helpful medical assistant.

Rules:
- Only use the provided context.
- If answer is not in context, say "Insufficient information."

Conversation:
{history_text}

Context:
{context}

Question:
{query}

Answer:
"""

# -----------------------------
# Chat Function
# -----------------------------
def chat_fn(message, history):

    query = preprocess_text(message)

    retrieved_docs = retrieve_docs(query)
    context = "\n".join(retrieved_docs["text"].tolist())

    prompt = build_prompt(query, history, context)

    inputs = tokenizer(prompt, return_tensors="pt")

    outputs = model.generate(
        **inputs,
        max_new_tokens=180,
        temperature=0.2,
        do_sample=True,
        top_p=0.9
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    if "Answer:" in response:
        response = response.split("Answer:")[-1].strip()

    return response

# -----------------------------
# UI
# -----------------------------
demo = gr.ChatInterface(
    fn=chat_fn,
    title="🧠 Medical RAG Assistant",
    description="RAG system using FAISS + TinyLlama (deployment safe)",
    examples=[
        "What are symptoms of diabetes?",
        "What causes kidney stones?",
        "Treatment for fever?",
        "Is vitamin D deficiency dangerous?"
    ]
)

demo.queue()   # 🔥 IMPORTANT for HF Spaces
demo.launch()  # 🔥 THIS WAS MISSING