import os
import gradio as gr
import requests
from sentence_transformers import SentenceTransformer
import chromadb

HF_API_URL = os.getenv("HF_API_URL")
HF_TOKEN = os.getenv("HF_TOKEN")
EMB_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/all-MiniLM-L6-v2")

client = chromadb.Client()
collection = client.get_or_create_collection("medini_memory")

embedding_model = SentenceTransformer(EMB_MODEL, use_auth_token=HF_TOKEN)

ROLES = [
    "Tutor", "Product Manager", "Project Manager", "Designer/Developer",
    "Technical Writer", "Analyst", "Researcher", "Innovator", "Friend", "Video Generator"
]

def retrieve_memory(user_input):
    results = collection.query(query_texts=[user_input], n_results=3)
    return " ".join(results.get("documents", [[]])[0])

def upsert_memory(user_input, response):
    emb = embedding_model.encode(user_input)
    collection.add(documents=[user_input], embeddings=[emb], ids=[str(hash(user_input))])

def medini_chat(user_input, role):
    context = retrieve_memory(user_input)
    payload = {"inputs": f"Role: {role}\nContext: {context}\nUser: {user_input}"}
    headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
    response = requests.post(HF_API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        result = response.json()
        generated = result.get("generated_text", str(result))
        upsert_memory(user_input, generated)
        return generated
    else:
        return f"Error: {response.text}"

iface = gr.Interface(
    fn=medini_chat,
    inputs=[gr.Textbox(label="Your Message"), gr.Dropdown(ROLES, label="Select Role")],
    outputs="text",
    title="Medini - Multi-role AI Assistant",
    description="Chat with Medini as Tutor, PM, Designer, Analyst, Innovator, Friend, and more."
)

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