import os
import gradio as gr
import pandas as pd
import torch
import numpy as np
from sentence_transformers import util
import google.generativeai as genai
import chromadb
from langchain_chroma import Chroma

# === Configuration ===
genai.configure(api_key=os.environ["GEMINI_API_KEY"])
embedding_model = "models/embedding-001"
llm_model_name = "models/gemma-3-4b-it"
collection_name = "xeno_collection"

# === Load and Clean Knowledge Base ===
df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
df_kb.dropna(subset=['Content'], inplace=True)

def prepare_documents(data):
    documents, metadatas, ids = [], [], []
    for item in data:
        documents.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
        metadatas.append({
            "question": item["Question"],
            "content": item["Content"],
            "section": item.get("Section", ""),
            "source": item.get("Source", ""),
            "owner": item.get("Owner", ""),
            "tag": item.get("Tag", "")
        })
        ids.append(item["ID"])
    return documents, metadatas, ids

xeno_data_list = df_kb.to_dict('records')
documents, metadatas, ids = prepare_documents(xeno_data_list)

# === Setup ChromaDB ===
client = chromadb.PersistentClient(path="./xeno_db")
try:
    collection = client.get_collection(name=collection_name)
except:
    collection = client.create_collection(name=collection_name)
    collection.add(documents=documents, metadatas=metadatas, ids=ids)

vector_store = Chroma(client=client, collection_name=collection_name)
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})

# === Prompt System ===
SYSTEM_PROMPT = """You are a friendly XENO Support Assistant, an AI-powered helpful and professional customer service representative.
Use only the information provided in the knowledge base context to answer user queries.
Do not hallucinate. If context doesn't contain relevant info, say so in a calm polite manner.
Only use context that is clearly relevant to the user's question.
For greetings like “hi” or “hello”, respond politely without using the context.
remember previous conversations.
"""

# === Context Processing ===
def process_context(results, cosine_scores, max_results=2):
    sorted_indices = np.argsort(cosine_scores)[::-1][:max_results]
    formatted_context = ""
    for i, idx in enumerate(sorted_indices, 1):
        result = results[idx]
        score = cosine_scores[idx]
        formatted_context += f"Knowledge Entry {i}:\n"
        formatted_context += f"Q: {result.metadata.get('question', 'N/A')}\n"
        formatted_context += f"A: {result.metadata.get('content', 'N/A')}\n"
        formatted_context += "-" * 40 + "\n"
    return formatted_context

# === LLM Generation ===
def generate_xeno_response(context, question):
    model = genai.GenerativeModel(llm_model_name)
    prompt = f"""{SYSTEM_PROMPT}

### CONTEXT ###
{context}

### QUESTION ###
{question}"""
    response = model.generate_content(prompt)
    return response.text.strip()

# === Main Interface Logic ===
def get_context_and_answer(message, history):
    if message.lower().strip() in {"hi", "hello", "hey"}:
        return "Hello! How can I assist you with XENO services today?"

    queried_results = retriever.invoke(message)
    query_embedding = genai.embed_content(model=embedding_model,
                                          content=message,
                                          task_type="retrieval_query")['embedding']
    cosine_scores = []
    for doc in queried_results:
        doc_embedding = genai.embed_content(model=embedding_model,
                                            content=doc.page_content,
                                            task_type="retrieval_document")['embedding']
        cos_sim = util.cos_sim(torch.tensor(query_embedding).float(), torch.tensor(doc_embedding).float())[0][0].item()
        cosine_scores.append(cos_sim)

    # If none of the results have sufficient similarity, fallback
    if max(cosine_scores) < 0.4:
        return "I'm sorry, I couldn't find the specific information you're looking for in my knowledge base."

    context = process_context(queried_results, cosine_scores)
    return generate_xeno_response(context, message)

# === Gradio UI ===
iface = gr.ChatInterface(
    fn=get_context_and_answer,
    title="ASKXENO",
    description="Ask anything about XENO's financial services.",
    theme="soft"
)

if __name__ == "__main__":
    iface.launch()