import gradio as gr
import requests
import os
import faiss
import numpy as np
import json
from fastapi import FastAPI, Request
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

# ✅ Load vector data
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_KEY = os.environ.get("OPENROUTER_API_KEY")
MODEL = "qwen/qwen-2.5-coder-32b-instruct:free"

# ✅ Semantic search
def get_context(query, top_k=5):
    query_vec = embed_model.encode([query])
    D, I = index.search(np.array(query_vec), top_k)
    return "\n".join([texts[i] for i in I[0]])

# ✅ Chatbot response
def chat_fn(message, history):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    context = get_context(message)
    messages = [{"role": "system", "content": f"You are CODEX Assistant by Mirxa Kamran. Use this context:\n{context}"}]

    for user, assistant in history:
        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 reply

# ✅ Gradio UI
demo = gr.ChatInterface(
    fn=chat_fn,
    title="💻 CODEX Assistant by Mirxa Kamran",
    description="Chat with a context-aware AI code assistant.",
    theme="soft"
)

# ✅ FastAPI app
app = FastAPI()

# ✅ Mount Gradio on root path
app = gr.mount_gradio_app(app, demo, path="/")

# ✅ FastAPI POST API endpoint
class ChatRequest(BaseModel):
    message: str
    history: list = []

@app.post("/chat")
def api_chat(req: ChatRequest):
    reply = chat_fn(req.message, req.history)
    return {"response": reply}

# ✅ Run manually in local dev or on Spaces
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)