import gradio as gr
import requests
import os
import faiss
import numpy as np
import json
from sentence_transformers import SentenceTransformer

# ✅ Load context
with open("texts.json", "r", encoding="utf-8") as f:
    texts = json.load(f)

index = faiss.read_index("faiss_index.bin")
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

# ✅ API Setup
API_KEY = os.environ.get("OPENROUTER_API_KEY")
MODEL = "meta-llama/llama-3.3-8b-instruct:free"

# ✅ Context search
def get_context(query, top_k=5, min_score=0.05):
    query_vec = embed_model.encode([query])
    D, I = index.search(np.array(query_vec), top_k)

    # If top score is too low, reject
    if D[0][0] < min_score:
        return None

    return "\n".join([texts[i] for i in I[0]])

# ✅ Chat function
def chat_fn(message, history):
    context = get_context(message)
    if context is None:
        return history[:-1] + [(message, "❌ Sorry! I cannot answer that.")]

    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    messages = [
        {
            "role": "system",
            "content": f"You are a helpful assistant. Use ONLY this context to answer:\n{context}\nIf not found in context, reply: 'Sorry! I cannot answer that.'"
        }
    ]

    for user, assistant in history[:-1]:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": assistant})

    messages.append({"role": "user", "content": message})

    payload = {
        "model": MODEL,
        "messages": messages
    }

    try:
        response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload)
        response.raise_for_status()
        reply = response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        reply = f"❌ Error: {e}"

    return history[:-1] + [(message, reply)]

# ✅ UI
with gr.Blocks(css="""
  .footer {display: none !important;}
  #chat-window {height: 500px !important; overflow-y: auto;}
""") as demo:
    chatbot = gr.Chatbot(elem_id="chat-window")
    state = gr.State([])

    with gr.Row():
        msg = gr.Textbox(placeholder="Type your message and press enter...", scale=9)

    # Typing sim
    def user_send(message, history):
        return "", history + [(message, "⏳ ...")]

    def complete_chat(message, history):
        return chat_fn(message, history)

    msg.submit(user_send, [msg, state], [msg, chatbot]).then(
        complete_chat, [msg, state], [chatbot]
    )

demo.launch()