import json
import asyncio
import os
import numpy as np
import gradio as gr
import plotly.graph_objects as go

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse

from huggingface_hub import InferenceClient

# ─────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────

MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
HF_TOKEN = os.environ.get("HF_TOKEN")

if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN missing.")

client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

# ─────────────────────────────────────────────
# FASTAPI
# ─────────────────────────────────────────────

api = FastAPI()

api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─────────────────────────────────────────────
# CHAT ENDPOINT (UNCHANGED CORE)
# ─────────────────────────────────────────────

@api.post("/api/chat")
async def chat(request: Request):
    body = await request.json()
    messages = body.get("messages", [])

    async def event_stream():
        try:
            stream = client.chat.completions.create(
                model=MODEL_ID,
                messages=messages,
                max_tokens=512,
                temperature=0.7,
                stream=True,
            )

            full_text = ""

            for chunk in stream:
                try:
                    delta = chunk.choices[0].delta
                    if delta and delta.content:
                        full_text += delta.content

                        yield json.dumps({
                            "content": delta.content
                        }) + "\n"

                        await asyncio.sleep(0.01)

                except Exception:
                    continue

            yield json.dumps({
                "done": True,
                "full": full_text
            }) + "\n"

        except Exception as e:
            yield json.dumps({
                "error": str(e),
                "done": True
            }) + "\n"

    return StreamingResponse(event_stream(), media_type="application/x-ndjson")

# ─────────────────────────────────────────────
# CODETTE UI (GRADIO)
# ─────────────────────────────────────────────

CUSTOM_CSS = """
body {
    background: radial-gradient(circle at top, #14142b, #0b0b17);
    color: #e5e7eb;
}

.metric-box {
    background: rgba(20,20,40,0.7);
    border: 1px solid rgba(168,85,247,0.3);
    padding: 10px;
    border-radius: 10px;
    font-family: monospace;
    margin-bottom: 10px;
}

button {
    background: linear-gradient(135deg,#a855f7,#06b6d4) !important;
    border: none !important;
}
"""

def call_backend(message):

    import requests

    url = "http://localhost:7860/api/chat"

    response = requests.post(
        url,
        json={"messages": [{"role": "user", "content": message}]},
        stream=True,
    )

    full = ""

    for line in response.iter_lines():
        if not line:
            continue
        data = json.loads(line.decode())

        if "content" in data:
            full += data["content"]

    return full

def process(msg, history):

    if not msg.strip():
        return history, "", "", None

    history.append({"role": "user", "content": msg})

    response = call_backend(msg)

    history.append({"role": "assistant", "content": response})

    # simple metrics (can upgrade later)
    coherence = min(0.99, 0.6 + len(msg)/200)
    eta = 0.7

    metrics_html = f"""
    <div class="metric-box">
    Γ Phase Coherence: {coherence:.4f}<br>
    η Ethical Alignment: {eta:.4f}<br>
    Risk: LOW
    </div>
    """

    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=[0,1,0],
        y=[0,1,1],
        mode='markers+text',
        text=["newton","empathy","quantum"]
    ))

    return history, "", metrics_html, fig

def create_ui():

    with gr.Blocks(title="Codette-Demo not the actual codette model") as demo:

        gr.Markdown("# Codette")

        with gr.Row():

            with gr.Column(scale=3):
                chat = gr.Chatbot(height=520)

                msg = gr.Textbox(
                    lines=2,
                    placeholder="Ask Codette..."
                )

                send = gr.Button("▶")

            with gr.Column(scale=2):
                metrics = gr.HTML()
                graph = gr.Plot()

        def run(m, h):
            return process(m, h)

        send.click(
            run,
            [msg, chat],
            [chat, msg, metrics, graph]
        )

        msg.submit(
            run,
            [msg, chat],
            [chat, msg, metrics, graph]
        )

    return demo

# ─────────────────────────────────────────────
# COMBINE (IMPORTANT PART)
# ─────────────────────────────────────────────

app = gr.mount_gradio_app(api, create_ui(), path="/")

# ─────────────────────────────────────────────
# RUN
# ─────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)