#!/usr/bin/env python3
from __future__ import annotations

import base64
import difflib
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.request
from decimal import Decimal, InvalidOperation
from pathlib import Path
from typing import Any

PORT = 8080
SEEDS = [20260816, 20260817, 20260818]


def dump(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def post(path: str, payload: dict, timeout: int = 900) -> dict:
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}{path}",
        data=json.dumps(payload, ensure_ascii=False).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as res:
            return json.load(res)
    except urllib.error.HTTPError as exc:
        body = exc.read().decode(errors="replace")
        raise RuntimeError(f"HTTP {exc.code}: {body[:4000]}") from exc


def get(path: str, timeout: int = 30) -> dict:
    with urllib.request.urlopen(f"http://127.0.0.1:{PORT}{path}", timeout=timeout) as res:
        return json.load(res)


def peak_gpu(path: Path) -> int | None:
    values = []
    if path.exists():
        for line in path.read_text(errors="replace").splitlines():
            bits = [x.strip() for x in line.split(",")]
            if len(bits) > 1:
                try:
                    values.append(int(float(bits[1])))
                except ValueError:
                    pass
    return max(values) if values else None


class Server:
    def __init__(self, model_id: str, model: Path, binary: Path, out: Path, *, ctx: int, reasoning: str, effort: str = "default", mmproj: Path | None = None, cache_k: str = "f16", cache_v: str = "f16", label: str = "server"):
        self.model_id, self.model, self.binary, self.out = model_id, model, binary, out
        self.ctx, self.reasoning, self.effort, self.mmproj = ctx, reasoning, effort, mmproj
        self.cache_k, self.cache_v, self.label = cache_k, cache_v, label
        self.proc = self.monitor = self.log_handle = self.gpu_handle = None
        self.ready_seconds = None
        self.log_path = out / f"{label}.log"
        self.gpu_path = out / f"{label}-gpu.csv"

    def start(self):
        self.out.mkdir(parents=True, exist_ok=True)
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = f"/usr/local/cuda-11.8/lib64:{self.binary.parent}:" + env.get("LD_LIBRARY_PATH", "")
        cmd = [
            str(self.binary), "--model", str(self.model), "--alias", self.model_id,
            "--host", "127.0.0.1", "--port", str(PORT), "--ctx-size", str(self.ctx),
            "--n-gpu-layers", "all", "--flash-attn", "on",
            "--cache-type-k", self.cache_k, "--cache-type-v", self.cache_v,
            "--batch-size", "2048", "--ubatch-size", "512", "--parallel", "1",
            "--threads", "16", "--threads-batch", "16", "--metrics", "--slots",
            "--reasoning", self.reasoning, "--reasoning-effort", self.effort,
            "--reasoning-format", "deepseek", "--jinja",
        ]
        if self.mmproj:
            cmd += ["--mmproj", str(self.mmproj)]
        dump(self.out / f"{self.label}-command.json", {"argv": cmd, "env": {"LD_LIBRARY_PATH": env["LD_LIBRARY_PATH"]}})
        self.log_handle = self.log_path.open("w")
        self.gpu_handle = self.gpu_path.open("w")
        self.monitor = subprocess.Popen(
            ["nvidia-smi", "--query-gpu=timestamp,memory.used,power.draw,utilization.gpu", "--format=csv,noheader,nounits", "-lms", "200"],
            stdout=self.gpu_handle, stderr=subprocess.DEVNULL, text=True,
        )
        began = time.perf_counter()
        self.proc = subprocess.Popen(cmd, stdout=self.log_handle, stderr=subprocess.STDOUT, text=True, env=env)
        deadline = time.time() + 900
        while time.time() < deadline:
            if self.proc.poll() is not None:
                self.stop()
                raise RuntimeError(self.log_path.read_text(errors="replace")[-8000:])
            try:
                if get("/health", 3).get("status") == "ok":
                    self.ready_seconds = time.perf_counter() - began
                    dump(self.out / f"{self.label}-ready.json", {"seconds": self.ready_seconds})
                    return self
            except Exception:
                pass
            time.sleep(1)
        self.stop()
        raise TimeoutError("server load timeout")

    def stop(self) -> dict:
        if self.proc and self.proc.poll() is None:
            self.proc.terminate()
            try:
                self.proc.wait(20)
            except subprocess.TimeoutExpired:
                self.proc.kill()
        if self.monitor and self.monitor.poll() is None:
            self.monitor.terminate()
            try:
                self.monitor.wait(5)
            except subprocess.TimeoutExpired:
                self.monitor.kill()
        if self.log_handle:
            self.log_handle.close()
        if self.gpu_handle:
            self.gpu_handle.close()
        return {"peak_gpu_mib": peak_gpu(self.gpu_path), "ready_seconds": self.ready_seconds}

    def complete(self, messages: list[dict], *, seed: int, temperature: float, max_tokens: int, tools: list[dict] | None = None) -> dict:
        payload = {"model": self.model_id, "messages": messages, "seed": seed, "temperature": temperature, "top_p": .95, "max_tokens": max_tokens, "stream": False}
        if tools is not None:
            payload.update({"tools": tools, "tool_choice": "auto"})
        began = time.perf_counter()
        response = post("/v1/chat/completions", payload)
        return {"request": payload, "response": response, "elapsed_seconds": time.perf_counter() - began}

    def tokenize(self, text: str) -> int:
        return len(post("/tokenize", {"content": text, "add_special": True}).get("tokens", []))


def message(run: dict) -> dict:
    return run["response"]["choices"][0]["message"]


def content(run: dict) -> str:
    value = message(run).get("content")
    return value if isinstance(value, str) else ""


def norm(value: Any) -> str:
    return re.sub(r"\s+", " ", str(value).strip()).strip(" .").casefold()


def answer_matches(actual: Any, expected: Any) -> bool:
    """Strict semantic final-answer match with narrow support for labels/units on one number."""
    if norm(actual) == norm(expected):
        return True
    expected_text = str(expected).strip()
    if re.fullmatch(r"-?\d+(?:\.\d+)?", expected_text):
        numbers = re.findall(r"(?<![\w.])-?\d+(?:\.\d+)?(?![\w.])", str(actual))
        if len(numbers) != 1:
            return False
        try:
            return Decimal(numbers[0]) == Decimal(expected_text)
        except InvalidOperation:
            return False
    return False


def has(text: str, values: list[str]) -> bool:
    low = text.casefold()
    return all(str(value).casefold() in low for value in values)


def score_instruction(case: dict, text: str) -> tuple[bool, dict]:
    check = case["check"]
    kind, detail = check["type"], {}
    if kind == "exact":
        return text.strip() == check["value"], detail
    if kind == "json_exact":
        cleaned = text.strip()
        if cleaned.startswith(chr(96) * 3):
            cleaned = re.sub(r"^" + re.escape(chr(96) * 3) + r"(?:json)?\s*", "", cleaned, flags=re.I)
            cleaned = re.sub(r"\s*" + re.escape(chr(96) * 3) + r"$", "", cleaned)
        try:
            parsed = json.loads(cleaned)
        except Exception as exc:
            detail["error"], parsed = str(exc), None
        detail["parsed"] = parsed
        return parsed == check["value"], detail
    if kind == "word_limit_contains":
        count = len(re.findall(r"\b[\w.-]+\b", text))
        detail["word_count"] = count
        return count <= check["max_words"] and has(text, check["contains"]), detail
    if kind == "bullet_count_contains":
        count = len([line for line in text.splitlines() if line.strip().startswith("- ")])
        detail["bullet_count"] = count
        return count == check["count"] and has(text, check["contains"]), detail
    raise ValueError(kind)


def score_writing(case: dict, text: str) -> tuple[bool, dict]:
    words = len(re.findall(r"\b[\w.-]+\b", text))
    bullets = len([line for line in text.splitlines() if line.strip().startswith("- ")])
    sentences = len([x for x in re.split(r"(?<=[.!?])\s+", text.strip()) if x])
    passed = has(text, case["contains"]) and not any(x.casefold() in text.casefold() for x in case["forbidden"])
    if case["max_words"]:
        passed &= words <= case["max_words"]
    if case["sentence_count"]:
        passed &= sentences == case["sentence_count"]
    if case["bullet_count"]:
        passed &= bullets == case["bullet_count"]
    return bool(passed), {"word_count": words, "sentence_count": sentences, "bullet_count": bullets}


def final_answer(text: str) -> str:
    hits = re.findall(r"(?:^|\n)\s*FINAL\s*:\s*(.+?)\s*$", text, flags=re.I | re.M)
    return hits[-1].strip() if hits else (text.strip().splitlines()[-1] if text.strip() else "")


def args_equal(actual: dict, expected: dict) -> bool:
    if set(actual) != set(expected):
        return False
    for key, exp in expected.items():
        got = actual[key]
        if isinstance(exp, (int, float)) and isinstance(got, (int, float)):
            if abs(float(exp) - float(got)) > 1e-6:
                return False
        elif got != exp:
            return False
    return True


def tool_calls(msg: dict) -> list[dict]:
    rows = []
    for call in msg.get("tool_calls") or []:
        fn = call.get("function") or {}
        raw = fn.get("arguments", "{}")
        try:
            parsed = json.loads(raw) if isinstance(raw, str) else raw
        except Exception:
            parsed = {"parse_error": raw}
        rows.append({"id": call.get("id", ""), "name": fn.get("name"), "args": parsed, "raw": call})
    return rows


DB = {
    "prices": {"AX-17": 12.5, "BR-04": 8.25, "CZ-91": 19.0, "DN-33": 5.75, "EP-20": 31.4},
    "inventory": {"AX-17": 12, "BR-04": 7, "CZ-91": 19, "DN-33": 5, "EP-20": 3},
    "distances": {("Vancouver", "Victoria"): 115, ("Calgary", "Banff"): 128, ("Seattle", "Portland"): 280, ("Toronto", "Ottawa"): 450, ("Kelowna", "Kamloops"): 165},
    "rates": {("USD", "CAD"): 1.36, ("CAD", "USD"): .7353, ("EUR", "CAD"): 1.49, ("GBP", "USD"): 1.31, ("JPY", "CAD"): .0091},
}


def execute_tool(name: str, args: dict) -> dict:
    if name == "lookup_price":
        return {"unit_price": DB["prices"][args["sku"]]}
    if name == "calculate_total":
        return {"total": round(args["unit_price"] * args["quantity"] * (1 + args["tax_rate"]), 2)}
    if name == "lookup_distance":
        return {"km": DB["distances"][(args["origin"], args["destination"])]}
    if name == "convert_distance":
        return {"miles": round(args["km"] * .621371, 2)}
    if name == "get_inventory":
        return {"available": DB["inventory"][args["sku"]]}
    if name == "reserve_item":
        return {"status": "reserved", "sku": args["sku"], "quantity": args["quantity"]}
    if name == "lookup_rate":
        return {"rate": DB["rates"][(args["base"], args["quote"])]}
    if name == "convert_currency":
        return {"converted": round(args["amount"] * args["rate"], 2)}
    raise KeyError(name)


CODING_TOOLS = [
    {"type": "function", "function": {"name": "read_file", "description": "Read a repository text file.", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"], "additionalProperties": False}}},
    {"type": "function", "function": {"name": "write_file", "description": "Replace app.py with complete corrected content.", "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"], "additionalProperties": False}}},
    {"type": "function", "function": {"name": "run_tests", "description": "Run unit tests.", "parameters": {"type": "object", "properties": {}, "additionalProperties": False}}},
]


def tests(workdir: Path) -> dict:
    try:
        proc = subprocess.run([sys.executable, "-m", "unittest", "-q"], cwd=workdir, capture_output=True, text=True, timeout=30)
        return {"exit_code": proc.returncode, "stdout": proc.stdout[-5000:], "stderr": proc.stderr[-5000:]}
    except subprocess.TimeoutExpired:
        return {"exit_code": 124, "stdout": "", "stderr": "timeout"}


def coding_tool(workdir: Path, name: str, args: dict) -> dict:
    if name == "read_file":
        rel = args.get("path")
        if rel not in {"app.py", "test_app.py", "TASK.md"}:
            return {"error": "path not allowed"}
        return {"path": rel, "content": (workdir / rel).read_text()}
    if name == "write_file":
        if args.get("path") != "app.py":
            return {"error": "only app.py may be edited"}
        (workdir / "app.py").write_text(args.get("content", ""))
        return {"status": "written", "bytes": len(args.get("content", "").encode())}
    if name == "run_tests":
        return tests(workdir)
    return {"error": "unknown tool"}


def image_message(image_path: Path, prompt_text: str) -> list[dict]:
    encoded = base64.b64encode(image_path.read_bytes()).decode()
    return [{"role": "user", "content": [{"type": "text", "text": prompt_text}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded}"}}]}]
