import gradio as gr
import torch
import requests
import json
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen2.5-0.5B"
PEFT_REPO  = "coderofpears/gimkit-copilot"
SB_URL     = "https://unztspukhwdzrllavnef.supabase.co"
SB_ANON    = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InVuenRzcHVraHdkenJsbGF2bmVmIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NzM2MDMzNTAsImV4cCI6MjA4OTE3OTM1MH0.eUWx9O9uX-BqI7oUikVvSAwljgt4Gw3Fl6hb-Xz_YWo"

ALL_TOOLS = [
    "gkc_create_build","gkc_undo_last_build","gkc_get_map_devices",
    "gkc_look_around","gkc_capture_screenshot","gkc_remove_elements",
    "gkc_vc_checkpoint","gkc_list_devices","gkc_search_terrain",
    "gkc_search_props","gkc_get_device_info","gkc_skill_search",
    "gkc_web_search","gkc_wiki_search",
]

# ── Load model once ───────────────────────────────────────────────────────────
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

_base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(_base, PEFT_REPO)
model.eval()
DEVICE = next(model.parameters()).device
print("Model ready.")

# ── Supabase helpers ──────────────────────────────────────────────────────────
def base_h():
    return {"Content-Type": "application/json", "apikey": SB_ANON}

def auth_h(token):
    return {**base_h(), "Authorization": "Bearer " + token}

def queue_tool(token, name, args, timeout=30):
    r = requests.post(
        SB_URL + "/rest/v1/tool_calls",
        headers={**auth_h(token), "Prefer": "return=representation"},
        json={"name": name, "args": args, "status": "pending"},
        timeout=10,
    )
    r.raise_for_status()
    row_id = r.json()[0]["id"]
    deadline = time.time() + timeout
    while time.time() < deadline:
        time.sleep(1.2)
        pr = requests.get(
            SB_URL + f"/rest/v1/tool_calls?id=eq.{row_id}&select=status,result",
            headers=auth_h(token), timeout=10,
        )
        pr.raise_for_status()
        rows = pr.json()
        if rows and rows[0]["status"] in ("complete", "error"):
            return rows[0]["result"], rows[0]["status"]
    return {"error": "Timed out — is the Gimloader plugin open and connected?"}, "error"

# ── Model inference ───────────────────────────────────────────────────────────
SYSTEM = (
    "You are Gimkit Copilot, an AI that helps build Gimkit Creative maps. "
    "Call tools by responding with JSON: {\"tool\":\"<name>\",\"args\":{...}} "
    "Otherwise reply normally. Tools: " + ", ".join(ALL_TOOLS)
)

def run_model(history, user_msg):
    prompt = SYSTEM + "\n\n"
    for u, a in history:
        prompt += f"User: {u}\nAssistant: {a}\n"
    prompt += f"User: {user_msg}\nAssistant:"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(DEVICE)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    new_tokens = out[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

def try_parse_tool(text):
    start, end = text.find("{"), text.rfind("}") + 1
    if start == -1 or end == 0:
        return None
    try:
        obj = json.loads(text[start:end])
        if "tool" in obj and obj["tool"] in ALL_TOOLS:
            return obj["tool"], obj.get("args", {})
    except Exception:
        pass
    return None

def respond(message, history, token):
    if not token or not token.strip():
        history.append((message, "❌ Enter your GimkitCopilot token in the sidebar first."))
        return history, ""
    model_out  = run_model(history, message)
    tool_call  = try_parse_tool(model_out)
    if tool_call:
        name, args = tool_call
        history.append((message, f"🔧 Calling `{name}`..."))
        result, status = queue_tool(token.strip(), name, args)
        icon = "✅" if status == "complete" else "❌"
        result_str = json.dumps(result, indent=2) if isinstance(result, dict) else str(result)
        followup = run_model(
            history + [(message, model_out)],
            f"Tool result for {name}: {result_str}. Summarise what happened."
        )
        history[-1] = (message, f"🔧 `{name}`\n\n{icon} ```\n{result_str}\n```\n\n{followup}")
    else:
        history.append((message, model_out))
    return history, ""

# ── UI ────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=DM+Sans:wght@300;400;600&display=swap');
:root {
    --bg:#070b07; --surface:#0c120c; --border:#162016;
    --green:#4CAF50; --green-dim:#1a3a1a; --glow:rgba(76,175,80,.12);
    --text:#cce0cc; --muted:#4a664a; --accent:#a3e635;
}
body, .gradio-container { background:var(--bg)!important; font-family:'DM Sans',sans-serif; color:var(--text); }
#logo { font-family:'Space Mono',monospace; color:var(--green); font-size:20px; font-weight:700; padding:16px 0 2px; }
#sub  { color:var(--muted); font-size:11px; font-family:'Space Mono',monospace; margin-bottom:20px; }
.token-box textarea {
    background:#040804!important; border:1px solid var(--border)!important;
    color:var(--accent)!important; font-family:'Space Mono',monospace!important; font-size:11px!important;
    border-radius:8px!important;
}
.token-box textarea:focus { border-color:var(--green)!important; box-shadow:0 0 0 2px var(--glow)!important; }
#chatbot { background:var(--bg)!important; border-color:var(--border)!important; }
.send-btn button {
    background:var(--green-dim)!important; border:1px solid var(--green)!important;
    color:var(--green)!important; font-family:'Space Mono',monospace!important; font-weight:700!important;
}
.send-btn button:hover { background:#223a22!important; }
.clear-btn button {
    background:transparent!important; border:1px solid var(--border)!important;
    color:var(--muted)!important; font-size:11px!important;
}
.msg-box textarea {
    background:var(--surface)!important; border:1px solid var(--border)!important;
    color:var(--text)!important; font-family:'DM Sans',sans-serif!important; border-radius:8px!important;
}
.msg-box textarea:focus { border-color:var(--green)!important; }
"""

with gr.Blocks(title="Gimkit Copilot") as demo:
    token_state = gr.State("")

    with gr.Row():
        # ── Left sidebar: token input ─────────────────────────────────────────
        with gr.Column(scale=1, min_width=220):
            gr.HTML('<div id="logo">◆ Gimkit<br>Copilot</div>')
            gr.HTML('<div id="sub">Qwen2.5-0.5B + MCP</div>')
            token_box = gr.Textbox(
                label="GimkitCopilot Token",
                placeholder="Paste Bearer token…",
                lines=4,
                type="password",
                elem_classes=["token-box"],
            )
            gr.HTML(
                '''<div style="font-size:10px;color:#2a4a2a;line-height:1.6;margin-top:6px;">
                Get token from the<br>Gimloader plugin:<br>
                <b style="color:#3a6a3a">Shift+X → Account<br>→ Copy Token</b><br><br>
                Plugin must show<br><span style="color:#4CAF50">● Connected</span> for<br>map tools to work.
                </div>'''
            )
            clear_btn = gr.Button("Clear chat", size="sm", elem_classes=["clear-btn"])

        # ── Right: chat ───────────────────────────────────────────────────────
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                elem_id="chatbot",
                height=520,
                show_label=False,
                placeholder="<div style=\'text-align:center;color:#2a4a2a;font-family:Space Mono,monospace;margin-top:40px\'>Sign in via token and start building →</div>",
            )
            with gr.Row():
                msg_box = gr.Textbox(
                    placeholder="Ask me to build something, inspect your map, search the wiki…",
                    show_label=False,
                    lines=2,
                    scale=8,
                    elem_classes=["msg-box"],
                )
                send_btn = gr.Button("Send ▶", variant="primary", scale=1, elem_classes=["send-btn"])

    # wire up token box → state
    token_box.change(lambda t: t, inputs=[token_box], outputs=[token_state])

    send_btn.click(respond, inputs=[msg_box, chatbot, token_state], outputs=[chatbot, msg_box])
    msg_box.submit(respond, inputs=[msg_box, chatbot, token_state], outputs=[chatbot, msg_box])
    clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg_box])

demo.launch(css=CSS)