import os
import gradio as gr
import faiss
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# -----------------------------
# Load Documents
# -----------------------------

def load_documents():
    docs = []

    for file in os.listdir("knowledge_base"):
        with open(f"knowledge_base/{file}", "r", encoding="utf-8") as f:
            docs.append(f.read())

    for file in os.listdir("synthetic_data"):
        with open(f"synthetic_data/{file}", "r", encoding="utf-8") as f:
            docs.append(f.read())

    return docs


documents = load_documents()

# -----------------------------
# Embeddings + FAISS
# -----------------------------

embed_model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = embed_model.encode(documents)

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

# -----------------------------
# Load FLAN-T5 properly
# -----------------------------

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base").to(device)

# -----------------------------
# Retrieve Context
# -----------------------------

def retrieve(query, k=3):
    query_embedding = embed_model.encode([query])
    distances, indices = index.search(query_embedding, k)
    return "\n\n".join([documents[i] for i in indices[0]])

# -----------------------------
# Generate Answer (Proper RAG)
# -----------------------------

def generate_answer(query):
    context = retrieve(query)

    prompt = f"""
Answer the question using ONLY the context below.
If the answer is not in the context, say "Information not found in profile."

Context:
{context}

Question:
{query}

Answer:
"""

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        do_sample=False
    )

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

    return answer

# -----------------------------
# Gradio Interface
# -----------------------------

interface = gr.Interface(
    fn=generate_answer,
    inputs=gr.Textbox(label="Ask a question"),
    outputs=gr.Textbox(label="Answer"),
    title="Hari's AI Twin",
    description="Ask me anything about my professional journey."
)

interface.launch()