import spaces  # MUST precede torch — monkey-patches torch.cuda for ZeroGPU
import json
import re
import sys
from pathlib import Path

import torch
import gradio as gr
from huggingface_hub import snapshot_download, hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer

# ── Environment source (downloaded once at startup; CPU-only, fine at module scope) ──
ENV_REPO = "Nanthasit/hermes-tool-use-rl-env"
env_dir = Path(snapshot_download(ENV_REPO, repo_type="dataset"))
sys.path.insert(0, str(env_dir))
sys.path.insert(0, str(env_dir / "server"))

from models import HermesToolAction  # noqa: E402
from tasks import TASKS_BY_ID  # noqa: E402
from hermes_tool_env import HermesToolEnvironment  # noqa: E402

TASK_IDS = sorted(TASKS_BY_ID)
MAX_STEPS = 12

SYSTEM = ("You are a coding agent working in a sandboxed workspace. Inspect the "
          "files, make the change the task asks for, then call submit to have it "
          "graded. Call exactly one tool per turn.")

TOOLS = [
    {"type": "function", "function": {"name": "terminal",
        "description": "Run a shell command in the task workspace; returns combined stdout/stderr.",
        "parameters": {"type": "object", "properties": {
            "command": {"type": "string", "description": "Shell command to run."}}, "required": ["command"]}}},
    {"type": "function", "function": {"name": "read_file",
        "description": "Read a file from the task workspace.",
        "parameters": {"type": "object", "properties": {
            "path": {"type": "string", "description": "Relative path of the file to read."}}, "required": ["path"]}}},
    {"type": "function", "function": {"name": "write_file",
        "description": "Write or overwrite a file in the task workspace.",
        "parameters": {"type": "object", "properties": {
            "path": {"type": "string", "description": "Relative path of the file to write."},
            "content": {"type": "string", "description": "Full content to write to the file."}},
            "required": ["path", "content"]}}},
    {"type": "function", "function": {"name": "patch",
        "description": "Find-and-replace a unique substring in a file.",
        "parameters": {"type": "object", "properties": {
            "path": {"type": "string", "description": "Relative path of the file to patch."},
            "old_string": {"type": "string", "description": "Exact unique substring to replace."},
            "new_string": {"type": "string", "description": "Replacement text."}},
            "required": ["path", "old_string", "new_string"]}}},
    {"type": "function", "function": {"name": "submit",
        "description": "End the episode and grade the task. Call this only when you believe the task is solved.",
        "parameters": {"type": "object", "properties": {}}}},
]

_TC = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)


# ── Hand-rolled renderer (bench-v2 format; matches the original -tools models) ──
def _text(c):
    return "" if c is None else (c if isinstance(c, str) else json.dumps(c, ensure_ascii=False))


def _tools_block(tools):
    if not tools:
        return ""
    sigs = "\n".join(json.dumps(t, ensure_ascii=False) for t in tools)
    return ("\n\n# Tools\n\nYou may call one or more functions. Signatures are within "
            "<tools></tools>:\n<tools>\n" + sigs + "\n</tools>\n\nFor each call return:\n"
            "<tool_call>\n{\"name\": <name>, \"arguments\": <json>}\n</tool_call>")


def _render_msg(m, tools_sys):
    r = m.get("role")
    if r == "system":
        return "<|im_start|>system\n" + _text(m.get("content")) + _tools_block(tools_sys) + "<|im_end|>\n"
    if r == "user":
        return "<|im_start|>user\n" + _text(m.get("content")) + "<|im_end|>\n"
    if r == "tool":
        return "<|im_start|>user\n<tool_response>\n" + _text(m.get("content")) + "\n</tool_response><|im_end|>\n"
    if r == "assistant":
        return "<|im_start|>assistant\n" + _text(m.get("content")) + "<|im_end|>\n"
    return ""


def render_handrolled(messages):
    out = []
    for i, m in enumerate(messages):
        out.append(_render_msg(m, TOOLS if (i == 0 and m.get("role") == "system") else None))
    out.append("<|im_start|>assistant\n")
    return "".join(out)


def parse_tool_call(text):
    m = _TC.search(text)
    if not m:
        return None
    try:
        return json.loads(m.group(1))
    except Exception:
        return None


def _peft_base_model(repo_id):
    try:
        cfg_path = hf_hub_download(repo_id, "adapter_config.json")
    except Exception:
        return None
    with open(cfg_path) as f:
        return json.load(f).get("base_model_name_or_path")


# Cache loaded (model, tokenizer) per repo_id, populated inside @spaces.GPU where
# the GPU is actually attached. Read-mostly; fine for a personal eval tool.
_CACHE = {}


def load_model(repo_id):
    if repo_id in _CACHE:
        return _CACHE[repo_id]
    base_id = _peft_base_model(repo_id)
    load_id = base_id or repo_id
    try:
        tok = AutoTokenizer.from_pretrained(repo_id)
    except Exception:
        tok = AutoTokenizer.from_pretrained(base_id)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    m = AutoModelForCausalLM.from_pretrained(load_id, torch_dtype=torch.bfloat16, device_map="cuda")
    if base_id:
        from peft import PeftModel
        m = PeftModel.from_pretrained(m, repo_id)
    m.eval()
    _CACHE[repo_id] = (m, tok)
    return m, tok


def step_fn(model, tok, messages, render_mode):
    if render_mode == "native":
        enc = tok.apply_chat_template(messages, tools=TOOLS, add_generation_prompt=True,
                                      return_dict=True, return_tensors="pt").to(model.device)
        input_len = enc["input_ids"].shape[-1]
    else:
        prompt = render_handrolled(messages)
        enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
        input_len = enc["input_ids"].shape[-1]
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=512, do_sample=False, pad_token_id=tok.pad_token_id)
    return tok.decode(out[0, input_len:], skip_special_tokens=True)


def run_episode(model, tok, task_id, render_mode):
    env = HermesToolEnvironment()
    env.reset()
    obs = env.step(HermesToolAction(tool="select_task", task_id=task_id))
    messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": obs.result}]
    lines = [f"[task] {obs.result[:200]}"]
    for _ in range(MAX_STEPS):
        completion = step_fn(model, tok, messages, render_mode)
        messages.append({"role": "assistant", "content": completion})
        tc = parse_tool_call(completion)
        if tc is None:
            lines.append(f"[model] (no tool call) {completion[:120]}")
            messages.append({"role": "tool", "content": "No <tool_call> found; you must call a tool."})
            continue
        lines.append(f"[model] {tc.get('name')}({json.dumps(tc.get('arguments', {}))[:100]})")
        try:
            obs = env.step(HermesToolAction(tool=tc["name"], **(tc.get("arguments") or {})))
        except Exception as e:
            messages.append({"role": "tool", "content": f"Invalid tool call: {e}"})
            lines.append(f"[env] invalid: {e}")
            continue
        messages.append({"role": "tool", "content": obs.result})
        lines.append(f"[env] {obs.result[:100]}")
        if obs.done:
            return float(obs.reward or 0.0), lines
    return 0.0, lines


@spaces.GPU(duration=300)
def evaluate(repo_id, render_mode, show_transcript):
    repo_id = repo_id.strip()
    if not repo_id:
        return "Enter a model repo id.", ""
    model, tok = load_model(repo_id)
    rows, transcripts, passed = [], [], 0
    for task_id in TASK_IDS:
        reward, lines = run_episode(model, tok, task_id, render_mode)
        passed += int(reward >= 1.0)
        mark = "✅" if reward >= 1.0 else "❌"
        rows.append(f"| {task_id} | {mark} {reward:.0f} |")
        if show_transcript:
            transcripts.append(f"### {task_id} — reward {reward:.0f}\n```\n" + "\n".join(lines) + "\n```")
    summary = (f"## `{repo_id}` — **{passed}/{len(TASK_IDS)}** ({render_mode} render)\n\n"
               "| task | reward |\n|---|---|\n" + "\n".join(rows))
    return summary, ("\n\n".join(transcripts) if show_transcript else "")


with gr.Blocks(title="SakThai Agentic Eval") as demo:
    gr.Markdown(
        "# 🧪 SakThai Agentic Eval (free, ZeroGPU)\n"
        "Runs the 6-task [hermes-tool-use-rl-env](https://huggingface.co/datasets/Nanthasit/hermes-tool-use-rl-env) "
        "agentic coding benchmark against any SakThai model — binary pass/fail per task, graded by the "
        "environment's real checker. No HF Jobs / paid compute needed."
    )
    with gr.Row():
        repo = gr.Textbox(value="Nanthasit/sakthai-context-7b-tools", label="Model repo id", scale=3)
        render = gr.Dropdown(choices=["native", "handrolled"], value="native", label="Prompt render",
                             info="native = apply_chat_template (use for SFT/GRPO-trained models); handrolled = bench-v2 renderer", scale=2)
    show_tr = gr.Checkbox(value=True, label="Show per-task transcript")
    btn = gr.Button("Run agentic eval (6 tasks)", variant="primary")
    summary = gr.Markdown()
    transcript = gr.Markdown()
    btn.click(evaluate, inputs=[repo, render, show_tr], outputs=[summary, transcript])
    gr.Examples(
        examples=[
            ["Nanthasit/sakthai-context-7b-tools", "native", True],
            ["Nanthasit/sakthai-context-0.5b-tools-sft", "native", True],
            ["Nanthasit/sakthai-context-0.5b-tools", "native", True],
        ],
        inputs=[repo, render, show_tr],
    )

demo.launch()