import os
import uuid
import requests
import gradio as gr

# Models hosted by cognitivecomputations
MODELS = {
    "Dolphin Llama3 8B": "cognitivecomputations/dolphin-2.9-llama3-8b",
    "Dolphin Mixtral 8x7B": "cognitivecomputations/dolphin-mixtral-8x7b",
    "Dolphin Llama3 8B 256K": "cognitivecomputations/dolphin-llama3-8b-256k"
}

HF_API_TOKEN = os.environ.get("HF_TOKEN", "")

client_memory = {}
client_tokens = {}

def ensure_client(cid):
    if cid not in client_memory:
        client_memory[cid] = []
        client_tokens[cid] = 0

def query_hf(model, prompt):
    url = f"https://api-inference.huggingface.co/models/{model}"
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    data = {"inputs": prompt, "parameters": {"max_new_tokens": 512}}

    response = requests.post(url, headers=headers, json=data)
    try:
        return response.json()[0]["generated_text"]
    except:
        return f"Error: {response.text}"

def chat_fn(message, client_id, model_name):
    if not client_id or client_id.strip() == "":
        client_id = str(uuid.uuid4())

    ensure_client(client_id)

    history = ""
    for msg in client_memory[client_id]:
        history += f"{msg['role']}: {msg['content']}\n"

    prompt = history + f"user: {message}\nassistant:"

    model = MODELS[model_name]
    output = query_hf(model, prompt)

    client_memory[client_id].append({"role": "user", "content": message})
    client_memory[client_id].append({"role": "assistant", "content": output})

    client_tokens[client_id] += len(output.split())

    return output, client_id, client_tokens[client_id]

with gr.Blocks() as app:
    gr.Markdown("# 🔥 Dolphin Chat — HF API Powered")

    with gr.Row():
        client_id = gr.Textbox(label="Client ID")
        model_selector = gr.Dropdown(
            choices=list(MODELS.keys()),
            value="Dolphin Llama3 8B",
            label="Model"
        )

    chatbox = gr.Chatbot()
    message_box = gr.Textbox(label="Message")
    send_btn = gr.Button("Send")
    tokens_used = gr.Number(label="Tokens Used")

    send_btn.click(
        chat_fn,
        inputs=[message_box, client_id, model_selector],
        outputs=[chatbox, client_id, tokens_used]
    )

app.launch()