import os
import requests
import gradio as gr

BASE_URL = "https://router.huggingface.co/v1"

CSS = """
:root{
  --bg:#0b0f19;
  --panel:rgba(255,255,255,.06);
  --panel2:rgba(255,255,255,.08);
  --text:rgba(255,255,255,.92);
  --muted:rgba(255,255,255,.65);
  --border:rgba(255,255,255,.10);
  --shadow: 0 20px 60px rgba(0,0,0,.45);
}
body.light{
  --bg:#f6f7fb;
  --panel:rgba(0,0,0,.04);
  --panel2:rgba(0,0,0,.06);
  --text:rgba(0,0,0,.88);
  --muted:rgba(0,0,0,.55);
  --border:rgba(0,0,0,.10);
  --shadow: 0 20px 60px rgba(0,0,0,.12);
}
.gradio-container{
  background: radial-gradient(1200px 700px at 20% 10%, rgba(98, 123, 255, .22), transparent 60%),
              radial-gradient(1000px 700px at 80% 30%, rgba(255, 90, 150, .18), transparent 60%),
              var(--bg) !important;
  color: var(--text) !important;
}
#wrap{ max-width: 980px; margin: 22px auto; padding: 0 14px; }
.header{
  display:flex; justify-content:space-between; align-items:center;
  padding: 18px 14px; border: 1px solid var(--border);
  border-radius: 18px; background: var(--panel); box-shadow: var(--shadow);
}
.block{
  border: 1px solid var(--border); border-radius: 18px;
  background: var(--panel); box-shadow: var(--shadow);
  padding: 14px; margin-top: 14px;
}
.subtle{ color: var(--muted); }
.chat_shell{
  border: 1px solid var(--border); border-radius: 18px;
  background: var(--panel); box-shadow: var(--shadow);
  padding: 10px; margin-top: 14px;
}
.composer{ margin-top: 10px; }
.primary button, .secondary button{
  border-radius: 14px !important;
  border: 1px solid var(--border) !important;
}
.primary button{ background: rgba(98, 123, 255, .85) !important; color: white !important; }
.secondary button{ background: var(--panel2) !important; color: var(--text) !important; }
footer{ display:none !important; }
"""

TOGGLE = """
<div style="display:flex;align-items:center;gap:10px;user-select:none;">
  <span class="subtle">Theme</span>
  <label style="display:flex;align-items:center;gap:10px;">
    <span class="subtle">Light</span>
    <input id="themeToggle" type="checkbox" style="transform:scale(1.1);"/>
    <span class="subtle">Dark</span>
  </label>
</div>
<script>
  const saved = localStorage.getItem("theme") || "dark";
  if(saved==="light") document.body.classList.add("light");
  const t = document.getElementById("themeToggle");
  t.checked = (saved !== "light");
  t.addEventListener("change", ()=>{
    if(t.checked){
      document.body.classList.remove("light");
      localStorage.setItem("theme","dark");
    } else {
      document.body.classList.add("light");
      localStorage.setItem("theme","light");
    }
  });
</script>
"""

# ----------------------------
# HF Router helpers
# ----------------------------
def get_headers():
    key = os.environ.get("HF_API_KEY")
    if not key:
        return None, "HF_API_KEY missing (Space Settings → Variables)"
    return {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}, None


def list_models(headers):
    r = requests.get(f"{BASE_URL}/models", headers=headers, timeout=20)
    if r.status_code != 200:
        return None, f"HTTP {r.status_code}: {r.text}"
    data = r.json()
    models = [m.get("id") for m in data.get("data", []) if m.get("id")]
    if not models:
        return None, "No models returned"
    return models, None


def chat_call(headers, model, messages, temperature, max_tokens):
    payload = {
        "model": model,
        "messages": messages,
        "temperature": float(temperature),
        "max_tokens": int(max_tokens),
    }
    r = requests.post(
        f"{BASE_URL}/chat/completions",
        headers=headers,
        json=payload,
        timeout=60,
    )
    if r.status_code != 200:
        return f"HTTP {r.status_code}: {r.text}"

    data = r.json()
    try:
        return data["choices"][0]["message"]["content"]
    except Exception:
        return f"Unexpected response shape: {data}"


# ----------------------------
# Message utilities
# ----------------------------
MAX_TURNS = 14  # last N user+assistant pairs

def trimmed_messages(msgs):
    return (msgs or [])[-MAX_TURNS * 2 :]

def build_messages(system_prompt, msgs, user_msg):
    return (
        [{"role": "system", "content": system_prompt}]
        + trimmed_messages(msgs)
        + [{"role": "user", "content": user_msg}]
    )

def ui_messages_from_state(msgs):
    # Chatbot (messages format) expects list of dicts with role/content.
    # Hide system messages from UI.
    out = []
    for m in msgs or []:
        role = m.get("role")
        if role in ("user", "assistant"):
            out.append({"role": role, "content": m.get("content", "")})
    return out


# ----------------------------
# Gradio app
# ----------------------------
with gr.Blocks(title="Chat") as demo:
    with gr.Column(elem_id="wrap"):
        with gr.Row(elem_classes=["header"]):
            gr.Markdown("## My AI Chatbot")
            gr.HTML(TOGGLE)

        status = gr.Markdown("Starting…", elem_classes=["subtle"])

        with gr.Column(elem_classes=["block"]):
            with gr.Row():
                model_dd = gr.Dropdown(
                    label="Model (from /v1/models)",
                    choices=[],
                    value=None,
                    interactive=True,
                )
                refresh_btn = gr.Button("Refresh models", elem_classes=["secondary"])

            model_override = gr.Textbox(
                label="Model override (optional)",
                placeholder="Paste a chat-capable model id here (takes precedence).",
            )

            system_prompt = gr.Textbox(label="System prompt", value="You are a helpful assistant.")
            with gr.Row():
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
                max_tokens = gr.Slider(64, 2048, value=800, step=32, label="Max tokens")

        with gr.Column(elem_classes=["chat_shell"]):
            # IMPORTANT: no type=... (your Gradio build doesn't support it)
            chatbot = gr.Chatbot(height=560)

        msg_state = gr.State([])  # OpenAI-style list[dict]

        with gr.Column(elem_classes=["composer"]):
            msg = gr.Textbox(placeholder="Message…", show_label=False, lines=1)
            with gr.Row():
                send = gr.Button("Send", elem_classes=["primary"])
                clear = gr.Button("Clear", elem_classes=["secondary"])

    def do_refresh_models():
        headers, err = get_headers()
        if err:
            return f"❌ {err}", gr.update(choices=[], value=None)

        models, err = list_models(headers)
        if err:
            return f"❌ {err}", gr.update(choices=[], value=None)

        default = models[0] if models else None
        return f"✅ Models loaded ({len(models)}).", gr.update(choices=models, value=default)

    def init():
        return do_refresh_models()

    def respond(user_msg, msgs, model_dd_value, model_text, sys_prompt, temp, mx):
        user_msg = (user_msg or "").strip()
        msgs = msgs or []
        if not user_msg:
            return "", ui_messages_from_state(msgs), msgs

        headers, err = get_headers()
        if err:
            msgs = msgs + [
                {"role": "user", "content": user_msg},
                {"role": "assistant", "content": f"Setup error: {err}"},
            ]
            return "", ui_messages_from_state(msgs), msgs

        model = (model_text or "").strip() or (model_dd_value or "").strip()
        if not model:
            msgs = msgs + [
                {"role": "user", "content": user_msg},
                {"role": "assistant", "content": "No model selected. Choose one or set Model override."},
            ]
            return "", ui_messages_from_state(msgs), msgs

        req_messages = build_messages(sys_prompt or "You are a helpful assistant.", msgs, user_msg)
        bot = chat_call(headers, model, req_messages, temp, mx)

        msgs = msgs + [
            {"role": "user", "content": user_msg},
            {"role": "assistant", "content": bot},
        ]
        return "", ui_messages_from_state(msgs), msgs

    def do_clear():
        return [], []

    demo.load(init, outputs=[status, model_dd])
    refresh_btn.click(do_refresh_models, outputs=[status, model_dd])

    send.click(
        respond,
        inputs=[msg, msg_state, model_dd, model_override, system_prompt, temperature, max_tokens],
        outputs=[msg, chatbot, msg_state],
    )
    msg.submit(
        respond,
        inputs=[msg, msg_state, model_dd, model_override, system_prompt, temperature, max_tokens],
        outputs=[msg, chatbot, msg_state],
    )
    clear.click(do_clear, outputs=[chatbot, msg_state])

# Keep this if your Gradio build supports css= on launch()
demo.launch(css=CSS)