import gradio as gr
import requests
import os
from typing import Generator

REGRAPH_API_KEY = os.getenv("REGRAPH_API_KEY", "")
REGRAPH_BASE_URL = "https://api.regraph.tech/v1"

# Gradio 5: history is list[dict] with keys "role" and "content"
def chat(message: str, history: list[dict]) -> Generator[str, None, None]:
    messages = [
        {"role": "system", "content": "You are a helpful AI assistant powered by ReGraph LLM — a decentralized, continuously-trained language model running on distributed GPU/NPU nodes worldwide."}
    ]
    # Gradio 5 passes history as list of {"role": ..., "content": ...} dicts
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": message})

    try:
        resp = requests.post(
            f"{REGRAPH_BASE_URL}/chat/completions",
            headers={"Authorization": f"Bearer {REGRAPH_API_KEY}", "Content-Type": "application/json"},
            json={"model": "regraph-llm-latest", "messages": messages, "max_tokens": 1024},
            timeout=60,
        )
        resp.raise_for_status()
        yield resp.json()["choices"][0]["message"]["content"]
    except Exception as e:
        yield f"⚠️ {e}"

demo = gr.ChatInterface(
    fn=chat,
    type="messages",
    title="⚡ ReGraph LLM",
    description="""Interact with **ReGraph LLM** — a continuously-trained language model powered by decentralized GPU/NPU nodes worldwide.

[Platform](https://regraph.tech) · [Docs](https://regraph.tech/docs) · [GitHub](https://github.com/ildu00/ReGraph)""",
    examples=[
        "What is decentralized AI compute?",
        "Explain the ReGraph network in simple terms",
        "Write a Python script to call the ReGraph API",
        "Compare centralized vs decentralized AI inference",
    ],
    theme=gr.themes.Soft(primary_hue="violet"),
)

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