from __future__ import annotations

import json
import os
import platform
import random
import re
import subprocess
import tarfile
import threading
import time
import urllib.request
from functools import partial
from pathlib import Path
from typing import Any

import gradio as gr
from huggingface_hub import HfApi, hf_hub_download

try:
    import spaces
except Exception:  # pragma: no cover - the package exists on HF ZeroGPU runtimes
    spaces = None  # type: ignore[assignment]

# ---------------------------------------------------------------------------
# Model / runtime configuration (proven llama.cpp llama-server backend)
# ---------------------------------------------------------------------------
MODEL_REPO = os.getenv("PHASE3_MODEL_REPO", "build-small-hackathon/phase-3-gguf")
MODEL_FILE = os.getenv("PHASE3_MODEL_FILE", "model-Q8_0.gguf")
MODEL_LABEL = "First-Principle AI"
LOCAL_MODEL_PATH = Path("/Users/user/.lmstudio/models/owenisas/Phase-3-GGUF/model-Q8_0.gguf")
LLAMA_RELEASE = os.getenv("PHASE3_LLAMA_RELEASE", "b9360")
LLAMA_URL = os.getenv(
    "PHASE3_LLAMA_URL",
    f"https://github.com/ggml-org/llama.cpp/releases/download/{LLAMA_RELEASE}/llama-{LLAMA_RELEASE}-bin-ubuntu-x64.tar.gz",
)
MAX_CONTEXT = int(os.getenv("PHASE3_MAX_CONTEXT", "2048"))
MIN_RAM_GB = float(os.getenv("PHASE3_MIN_RAM_GB", "38"))
DISABLE_MODEL = os.getenv("PHASE3_DISABLE_MODEL", "").lower() in {"1", "true", "yes"}
USE_ZEROGPU_DECORATOR = os.getenv("PHASE3_USE_ZEROGPU", "").lower() in {"1", "true", "yes"}
N_BATCH = int(os.getenv("PHASE3_N_BATCH", "256"))
N_UBATCH = int(os.getenv("PHASE3_N_UBATCH", "64"))
N_THREADS = int(os.getenv("PHASE3_THREADS", str(max(1, min(16, os.cpu_count() or 2)))))
N_THREADS_BATCH = int(os.getenv("PHASE3_THREADS_BATCH", str(N_THREADS)))
USE_MMAP = os.getenv("PHASE3_USE_MMAP", "1").lower() not in {"0", "false", "no"}
USE_MLOCK = os.getenv("PHASE3_USE_MLOCK", "").lower() in {"1", "true", "yes"}
FLASH_ATTN = os.getenv("PHASE3_FLASH_ATTN", "").lower() in {"1", "true", "yes"}
OFFLOAD_KQV = os.getenv("PHASE3_OFFLOAD_KQV", "1").lower() not in {"0", "false", "no"}
INFER_TIMEOUT = int(os.getenv("PHASE3_INFER_TIMEOUT", "900"))
SERVER_HOST = "127.0.0.1"
SERVER_PORT = int(os.getenv("PHASE3_SERVER_PORT", "8088"))
NO_WARMUP = os.getenv("PHASE3_NO_WARMUP", "1").lower() not in {"0", "false", "no"}

MODEL_LOCK = threading.Lock()
MODEL_PATH: Path | None = None
LLAMA_CLI_PATH: Path | None = None
LLAMA_SERVER_PATH: Path | None = None
LLAMA_SERVER_PROCESS: subprocess.Popen[str] | None = None
MODEL_ERROR: str | None = None
MODEL_SETTINGS: dict[str, Any] = {}

DEFAULT_SYSTEM = (
    "You are First-Principle AI, a small reasoning model. For every question, first "
    "identify the goal and the physical object or person that must be PRESENT at the "
    "destination for the goal to succeed. If the goal needs a specific object (e.g. a "
    "vehicle that must be washed, charged, fueled, or repaired) to be at the location, "
    "then that object has to travel there and the distance is irrelevant. Also watch "
    "for false premises, swapped names, and altered classic riddles. Think briefly, "
    "then give a clear final answer."
)
BLINDSPOT_SYSTEM = (
    "You are First-Principle AI acting as a blind-spot checker. The user describes "
    "a plan, decision, or question. Identify the single most important hidden "
    "assumption or causal constraint they may be missing, state it in one line, "
    "then give a clear recommendation. Be concise and concrete."
)


def _gpu_decorator(fn):
    if not USE_ZEROGPU_DECORATOR:
        return fn
    if spaces is None:
        return fn
    try:
        return spaces.GPU(duration=120)(fn)
    except Exception:
        return fn


if spaces is not None:
    try:
        @spaces.GPU(duration=1)
        def _zerogpu_startup_probe() -> str:
            return "ZeroGPU configured"
    except Exception:
        def _zerogpu_startup_probe() -> str:
            return "ZeroGPU helper importable"
else:
    def _zerogpu_startup_probe() -> str:
        return "ZeroGPU helper unavailable"


def _meminfo_gb() -> tuple[float | None, float | None]:
    meminfo = Path("/proc/meminfo")
    if not meminfo.exists():
        return None, None
    data: dict[str, int] = {}
    for line in meminfo.read_text(encoding="utf-8", errors="ignore").splitlines():
        match = re.match(r"^(\w+):\s+(\d+)\s+kB", line)
        if match:
            data[match.group(1)] = int(match.group(2))
    total = data.get("MemTotal")
    available = data.get("MemAvailable")
    gb = 1024 * 1024
    return (total / gb if total else None, available / gb if available else None)


def _safe_env_summary() -> dict[str, str]:
    keys = [
        "SPACE_ID",
        "SPACE_HOST",
        "CUDA_VISIBLE_DEVICES",
        "PHASE3_MODEL_REPO",
        "PHASE3_MODEL_FILE",
        "PHASE3_THREADS",
    ]
    return {key: os.environ[key] for key in keys if key in os.environ}


def _repo_file_size() -> int | None:
    try:
        info = HfApi().model_info(MODEL_REPO, files_metadata=True)
    except Exception:
        return None
    for sibling in info.siblings or []:
        if sibling.rfilename == MODEL_FILE:
            return getattr(sibling, "size", None)
    return None


def _find_model_path() -> Path:
    if DISABLE_MODEL:
        raise RuntimeError("Model loading is disabled with PHASE3_DISABLE_MODEL=1.")

    explicit = os.getenv("PHASE3_MODEL_PATH")
    if explicit:
        path = Path(explicit)
        if path.exists():
            return path
        raise RuntimeError(f"PHASE3_MODEL_PATH does not exist: {explicit}")

    if LOCAL_MODEL_PATH.exists():
        return LOCAL_MODEL_PATH

    data_dir = Path(os.getenv("PHASE3_MODEL_DIR", "/data/phase-3-gguf"))
    if data_dir.parent.exists() and os.access(data_dir.parent, os.W_OK):
        data_dir.mkdir(parents=True, exist_ok=True)
        downloaded = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, local_dir=data_dir)
    else:
        downloaded = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    return Path(downloaded)


def _gpu_layers() -> int:
    if "PHASE3_N_GPU_LAYERS" in os.environ:
        return int(os.environ["PHASE3_N_GPU_LAYERS"])
    if os.getenv("CUDA_VISIBLE_DEVICES") and os.getenv("PHASE3_AUTO_GPU", "1").lower() not in {"0", "false", "no"}:
        return -1
    return 0


def _ensure_llama_binary(name: str) -> Path:
    global LLAMA_CLI_PATH, LLAMA_SERVER_PATH

    if name == "llama-cli" and LLAMA_CLI_PATH is not None and LLAMA_CLI_PATH.exists():
        return LLAMA_CLI_PATH
    if name == "llama-server" and LLAMA_SERVER_PATH is not None and LLAMA_SERVER_PATH.exists():
        return LLAMA_SERVER_PATH

    root = Path(os.getenv("PHASE3_LLAMA_DIR", "/tmp/phase3-llama.cpp"))
    release_dir = root / f"llama-{LLAMA_RELEASE}"
    binary = release_dir / name
    if binary.exists():
        binary.chmod(0o755)
        if name == "llama-cli":
            LLAMA_CLI_PATH = binary
        if name == "llama-server":
            LLAMA_SERVER_PATH = binary
        return binary

    root.mkdir(parents=True, exist_ok=True)
    archive = root / f"llama-{LLAMA_RELEASE}-bin-ubuntu-x64.tar.gz"
    if not archive.exists():
        urllib.request.urlretrieve(LLAMA_URL, archive)
    with tarfile.open(archive, "r:gz") as tar:
        tar.extractall(root)
    if not binary.exists():
        raise RuntimeError(f"{name} was not found after extracting {LLAMA_URL}")
    binary.chmod(0o755)
    if name == "llama-cli":
        LLAMA_CLI_PATH = binary
    if name == "llama-server":
        LLAMA_SERVER_PATH = binary
    return binary


def _prepare_runtime() -> tuple[Path, Path]:
    global MODEL_PATH, MODEL_ERROR, MODEL_SETTINGS

    if MODEL_ERROR is not None:
        raise RuntimeError(MODEL_ERROR)

    with MODEL_LOCK:
        if MODEL_ERROR is not None:
            raise RuntimeError(MODEL_ERROR)

        total_gb, available_gb = _meminfo_gb()
        if total_gb is not None and total_gb < MIN_RAM_GB:
            MODEL_ERROR = (
                f"Runtime has {total_gb:.1f} GB RAM, below the configured load threshold "
                f"of {MIN_RAM_GB:.1f} GB for the 31 GB Q8 GGUF."
            )
            raise RuntimeError(MODEL_ERROR)

        path = _find_model_path()
        server = _ensure_llama_binary("llama-server")
        MODEL_PATH = path
        n_gpu_layers = _gpu_layers()
        MODEL_SETTINGS = {
            "path": str(path),
            "llama_server": str(server),
            "n_ctx": MAX_CONTEXT,
            "n_batch": N_BATCH,
            "n_ubatch": N_UBATCH,
            "n_threads": N_THREADS,
            "n_threads_batch": N_THREADS_BATCH,
            "n_gpu_layers": n_gpu_layers,
            "use_mmap": USE_MMAP,
            "use_mlock": USE_MLOCK,
            "flash_attn": FLASH_ATTN,
            "offload_kqv": OFFLOAD_KQV,
            "no_warmup": NO_WARMUP,
        }
        return path, server


def _server_log_path() -> Path:
    return Path(os.getenv("PHASE3_SERVER_LOG", "/tmp/phase3-llama-server.log"))


def _tail_server_log(limit: int = 4000) -> str:
    path = _server_log_path()
    if not path.exists():
        return ""
    data = path.read_text(encoding="utf-8", errors="ignore")
    return data[-limit:]


def _server_url(path: str) -> str:
    return f"http://{SERVER_HOST}:{SERVER_PORT}{path}"


def _server_is_ready() -> bool:
    try:
        with urllib.request.urlopen(_server_url("/health"), timeout=5) as resp:
            return 200 <= resp.status < 500
    except Exception:
        return False


def _start_server() -> None:
    global LLAMA_SERVER_PROCESS

    model_path, server = _prepare_runtime()
    if LLAMA_SERVER_PROCESS is not None and LLAMA_SERVER_PROCESS.poll() is None and _server_is_ready():
        return

    cmd = [
        str(server),
        "-m",
        str(model_path),
        "--host",
        SERVER_HOST,
        "--port",
        str(SERVER_PORT),
        "-c",
        str(MAX_CONTEXT),
        "-t",
        str(N_THREADS),
        "-b",
        str(N_BATCH),
        "-ub",
        str(N_UBATCH),
    ]
    if _gpu_layers() != 0:
        cmd.extend(["-ngl", str(_gpu_layers())])
    if USE_MLOCK:
        cmd.append("--mlock")
    if not USE_MMAP:
        cmd.append("--no-mmap")
    if FLASH_ATTN:
        cmd.append("-fa")
    if NO_WARMUP:
        cmd.append("--no-warmup")

    env = os.environ.copy()
    binary_dir = str(server.parent)
    env["LD_LIBRARY_PATH"] = f"{binary_dir}:{env.get('LD_LIBRARY_PATH', '')}"
    log_path = _server_log_path()
    log_file = log_path.open("a", encoding="utf-8")
    log_file.write(f"\n--- starting llama-server: {' '.join(cmd)} ---\n")
    log_file.flush()
    LLAMA_SERVER_PROCESS = subprocess.Popen(
        cmd,
        cwd=binary_dir,
        env=env,
        stdout=log_file,
        stderr=subprocess.STDOUT,
        text=True,
    )

    deadline = time.time() + INFER_TIMEOUT
    while time.time() < deadline:
        if LLAMA_SERVER_PROCESS.poll() is not None:
            raise RuntimeError(f"llama-server exited early.\n{_tail_server_log()}")
        if _server_is_ready():
            return
        time.sleep(2)
    raise RuntimeError(f"llama-server did not become ready within {INFER_TIMEOUT}s.\n{_tail_server_log()}")


def _format_prompt(system_prompt: str, history: list[dict[str, str]], message: str) -> str:
    system = system_prompt.strip() or DEFAULT_SYSTEM
    turns = [f"<|im_start|>system\n{system}<|im_end|>"]
    for item in history[-10:]:
        role = item.get("role", "user")
        content = item.get("content", "")
        if role in {"user", "assistant"} and content:
            turns.append(f"<|im_start|>{role}\n{content}<|im_end|>")
    turns.append(f"<|im_start|>user\n{message}<|im_end|>")
    turns.append("<|im_start|>assistant\n")
    return "\n".join(turns)


@_gpu_decorator
def _complete(
    prompt: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    repeat_penalty: float,
) -> tuple[str, dict[str, Any]]:
    started = time.time()
    _start_server()
    payload = {
        "prompt": prompt,
        "n_predict": int(max_tokens),
        "temperature": float(temperature),
        "top_p": float(top_p),
        "repeat_penalty": float(repeat_penalty),
        "stop": ["<|im_end|>", "<|endoftext|>"],
    }
    req = urllib.request.Request(
        _server_url("/completion"),
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=INFER_TIMEOUT) as resp:
            output = json.loads(resp.read().decode("utf-8"))
    except Exception as exc:
        raise RuntimeError(f"llama-server completion failed: {exc}\n{_tail_server_log()}") from exc
    elapsed = max(time.time() - started, 0.001)
    text = (output.get("content") or "").strip()
    text = text.split("<|im_end|>", 1)[0].strip()
    completion_tokens = max(1, len(text.split()))
    return text, {
        "elapsed": elapsed,
        "completion_tokens": completion_tokens,
        "tokens_per_second": completion_tokens / elapsed,
    }


def _generate(
    message: str,
    system_prompt: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    repeat_penalty: float,
) -> tuple[str, dict[str, Any]]:
    prompt = _format_prompt(system_prompt or DEFAULT_SYSTEM, [], message)
    return _complete(prompt, max_tokens, temperature, top_p, repeat_penalty)


# ---------------------------------------------------------------------------
# The trap deck: trick questions where pattern-matching models commonly fail.
# `catch` = phrases that show the model reasoned correctly.
# `miss`  = phrases that show it fell for the familiar/decoy answer.
# ---------------------------------------------------------------------------
TRAPS: list[dict[str, Any]] = [
    {
        "id": "carwash",
        "emoji": "🚗",
        "title": "The 50-Meter Car Wash",
        "cat": "Goal-binding",
        "prompt": "I want to wash my car at a car wash that is 50 meters away. Should I walk there or drive there?",
        "why": "The car is the thing being washed, so it has to be AT the wash. The short distance is a decoy that nudges models toward 'just walk'.",
        "wrong": "Walk — it's only 50 meters.",
        "catch": ["drive", "driving", "drives", "bring the car", "take the car", "car must", "car needs to be", "car has to be", "the car itself", "need the car", "car to the wash", "car has to go", "car needs to go"],
        "miss": ["just walk", "you should walk", "better to walk", "i'd walk", "recommend walking", "walking introduces", "walk over", "walk there", "walking is", "proximity matters"],
    },
    {
        "id": "ev",
        "emoji": "🔌",
        "title": "The Almost-Empty EV",
        "cat": "Goal-binding",
        "prompt": "My electric car's battery is almost empty and the charging station is 60 meters away. Should I walk to the charger or drive there?",
        "why": "You can only charge the car if the car is at the charger. Walking solves nothing — the object that must move is the car.",
        "wrong": "Walk, the charger is close.",
        "catch": ["drive", "driving", "drives", "bring the car", "take the car", "car must", "car needs to be", "car has to be", "the car itself", "need the car", "car to the charger", "car has to go", "car needs to go"],
        "miss": ["just walk", "you should walk", "better to walk", "i'd walk", "recommend walking", "walk to the charger", "walk over", "walk there", "walking"],
    },
    {
        "id": "feathers",
        "emoji": "🪶",
        "title": "Feathers vs. Steel",
        "cat": "Unit misdirection",
        "prompt": "Which weighs more: a kilogram of feathers or a pound of steel?",
        "why": "The famous riddle compares a kilogram to a kilogram (equal). This version sneaks in a POUND of steel — a kilogram is ~2.2 lb, so the feathers win.",
        "wrong": "They weigh the same.",
        "catch": ["kilogram of feathers", "feathers weigh more", "feathers are heavier", "kilogram is heavier", "kilogram weighs more", "2.2", "the feathers"],
        "miss": ["weigh the same", "same weight", "they are equal", "are equal", "neither", "weigh exactly the same"],
    },
    {
        "id": "moses",
        "emoji": "🛶",
        "title": "Moses & the Ark",
        "cat": "Moses illusion",
        "prompt": "How many animals of each kind did Moses take on the Ark?",
        "why": "It was Noah, not Moses. The sentence reads so naturally that most readers (and models) answer 'two' without noticing the swapped name.",
        "wrong": "Two of each.",
        "catch": ["noah", "not moses", "wasn't moses", "was noah", "moses didn't", "moses did not"],
        "miss": ["two of each", "2 of each", "moses took", "moses brought"],
    },
    {
        "id": "batball",
        "emoji": "⚾",
        "title": "The Bat and the Ball",
        "cat": "Cognitive reflection",
        "prompt": "A bat and a ball cost $1.10 in total. The bat costs $1.00 more than the ball. How much does the ball cost?",
        "why": "The intuitive answer is 10 cents, but then the bat ($1.10) would be only $1.00 more total. The ball is actually 5 cents.",
        "wrong": "10 cents.",
        "catch": ["0.05", "$.05", "5 cents", "five cents", ".05", "5¢"],
        "miss": ["0.10", "10 cents", "ten cents", "$0.10", "a dime"],
    },
    {
        "id": "survivors",
        "emoji": "✈️",
        "title": "Where to Bury the Survivors",
        "cat": "Hidden presupposition",
        "prompt": "A plane crashes exactly on the border between the United States and Canada. Where should the survivors be buried?",
        "why": "You don't bury survivors — they're alive. The question smuggles in a false premise behind a geography distraction.",
        "wrong": "On the border / split between countries.",
        "catch": ["don't bury", "do not bury", "not buried", "aren't buried", "are alive", "still alive", "survivors are not", "you don't bury", "no one is buried", "survivors live", "living people", "are living", "do not need to be buried", "don't need to be buried", "not need to be buried", "no need to bury", "not dead", "you don't bury survivors"],
        "miss": ["buried on the border", "should be buried in", "bury them in", "buried in canada", "buried in the united states"],
    },
    {
        "id": "sally",
        "emoji": "👧",
        "title": "Sally's Sisters",
        "cat": "Relational logic",
        "prompt": "Sally has 3 brothers. Each of her brothers has 2 sisters. How many sisters does Sally have?",
        "why": "The 2 sisters every brother has are Sally plus one other girl. So Sally has just 1 sister — not 2, and not 6.",
        "wrong": "2 sisters (or 6).",
        "catch": ["1 sister", "one sister", "just one", "only one", "a single sister", "has 1", "she has one"],
        "miss": ["2 sisters", "two sisters", "6 sisters", "six sisters", "has 2 sister", "has six"],
    },
    {
        "id": "brain",
        "emoji": "🧠",
        "title": "The 10% Brain Myth",
        "cat": "False premise",
        "prompt": "Since humans only use 10% of their brains, what amazing things could a person do if they unlocked the other 90%?",
        "why": "The premise is a myth — people use virtually all of their brain. A first-principle answer rejects the premise instead of inventing superpowers.",
        "wrong": "Telekinesis, perfect memory, etc.",
        "catch": ["myth", "false", "not true", "misconception", "we use", "use virtually all", "use all", "use most", "entire brain", "whole brain", "premise is", "actually use"],
        "miss": ["telekinesis", "superpower", "could do anything", "unlock", "psychic"],
    },
    {
        "id": "sheep",
        "emoji": "🐑",
        "title": "All But Nine",
        "cat": "Language parsing",
        "prompt": "A farmer has 17 sheep. All but 9 run away. How many sheep does the farmer have left?",
        "why": "'All but 9 run away' means 9 stay. The trap is to subtract and answer 8.",
        "wrong": "8.",
        "catch": ["9", "nine"],
        "miss": ["8 ", "eight", " 8.", " 8,"],
    },
]
TRAP_BY_TITLE = {t["title"]: t for t in TRAPS}
TRAP_TITLES = [t["title"] for t in TRAPS]


def judge(trap: dict[str, Any], answer: str) -> str:
    a = " " + answer.lower() + " "
    caught = any(k.lower() in a for k in trap.get("catch", []))
    missed = any(k.lower() in a for k in trap.get("miss", []))
    if caught:
        return "caught"
    if missed:
        return "missed"
    return "unclear"


def _new_score() -> dict[str, int]:
    return {"caught": 0, "missed": 0, "unclear": 0}


def _bump(score: dict[str, int] | None, verdict: str) -> dict[str, int]:
    score = dict(score or _new_score())
    if verdict in score:
        score[verdict] += 1
    return score


def _rank(caught: int, total: int) -> str:
    if total == 0:
        return "No traps run yet"
    rate = caught / total
    if total >= 4 and rate >= 0.8:
        return "🧠 First-Principle Thinker"
    if rate >= 0.6:
        return "🤔 Sharp Reasoner"
    if rate >= 0.4:
        return "🙂 Sometimes Fooled"
    return "🫠 Pattern-Matcher"


def scoreboard_html(score: dict[str, int] | None) -> str:
    score = score or _new_score()
    caught = score.get("caught", 0)
    missed = score.get("missed", 0)
    unclear = score.get("unclear", 0)
    total = caught + missed + unclear
    rank = _rank(caught, total)
    return f"""
    <div class="fp-score">
      <div class="fp-stat fp-stat-good"><div class="fp-stat-n">{caught}</div><div class="fp-stat-l">✅ Caught</div></div>
      <div class="fp-stat fp-stat-bad"><div class="fp-stat-n">{missed}</div><div class="fp-stat-l">❌ Fooled</div></div>
      <div class="fp-stat fp-stat-meh"><div class="fp-stat-n">{unclear}</div><div class="fp-stat-l">🤔 Unclear</div></div>
      <div class="fp-stat fp-stat-rank"><div class="fp-stat-rank-n">{rank}</div><div class="fp-stat-l">{total} trap(s) run</div></div>
    </div>
    """


def trap_card_html(trap: dict[str, Any], revealed: bool = False) -> str:
    reveal = ""
    if revealed:
        reveal = f"""
        <div class="fp-reveal">
          <div><strong>Why it's a trap:</strong> {trap['why']}</div>
          <div class="fp-wrong"><strong>Common wrong answer:</strong> {trap['wrong']}</div>
        </div>
        """
    return f"""
    <div class="fp-qcard">
      <div class="fp-qcat">{trap['emoji']} {trap['cat']}</div>
      <div class="fp-qtitle">{trap['title']}</div>
      <div class="fp-qtext">&ldquo;{trap['prompt']}&rdquo;</div>
      {reveal}
    </div>
    """


_BANNER = {
    "idle": ("🎯", "Pick a trap", "Tap a card to challenge the model.", "fp-bn-idle"),
    "thinking": ("⏳", "Reasoning…", "First run loads the model (~1–2 min); then it's quick.", "fp-bn-think"),
    "caught": ("✅", "Caught the trap!", "First-Principle AI reasoned past the decoy.", "fp-bn-good"),
    "missed": ("❌", "Fell for it", "The familiar decoy answer slipped through.", "fp-bn-bad"),
    "unclear": ("🤔", "You be the judge", "Auto-check couldn't tell — read the answer.", "fp-bn-meh"),
    "error": ("⚠️", "Runtime issue", "Generation didn't complete — details below.", "fp-bn-meh"),
}


def verdict_banner_html(kind: str) -> str:
    icon, label, sub, cls = _BANNER.get(kind, _BANNER["idle"])
    return f"""
    <div class="fp-banner {cls}">
      <div class="fp-banner-ic">{icon}</div>
      <div class="fp-banner-tx">
        <div class="fp-banner-label">{label}</div>
        <div class="fp-banner-sub">{sub}</div>
      </div>
    </div>
    """


def _split_think(text: str) -> tuple[str, str]:
    """Split a reasoning model's output into (reasoning, final answer)."""
    text = (text or "").strip()
    if "<think>" in text:
        if "</think>" in text:
            pre, rest = text.split("</think>", 1)
            reasoning = pre.split("<think>", 1)[-1].strip()
            answer = rest.strip()
        else:
            # generation was cut off before the model finished thinking
            reasoning = text.split("<think>", 1)[-1].strip()
            answer = ""
    else:
        reasoning, answer = "", text
    # some chat templates prefix the final line with "Answer:"
    return reasoning, answer


def _render_reasoning(reasoning: str, answer: str) -> str:
    parts: list[str] = []
    if answer:
        parts.append(answer)
    else:
        parts.append(
            "_The model was still reasoning when it hit the token limit. "
            "Raise **Max tokens** in the Engine Room for the full answer._"
        )
    if reasoning:
        parts.append(
            "\n\n<details><summary>🧠 Show first-principle reasoning</summary>\n\n"
            f"{reasoning}\n\n</details>"
        )
    return "".join(parts)


def _answer_md(reasoning: str, answer: str) -> str:
    return f"**First-Principle AI says:**\n\n{_render_reasoning(reasoning, answer)}"


def _thinking_md() -> str:
    return (
        "⏳ **First-Principle AI is reasoning…**\n\n"
        "The first run loads the 31 GB model through llama.cpp (about 1–2 minutes). "
        "After that, answers come back quickly."
    )


def _status_markdown() -> str:
    total_gb, available_gb = _meminfo_gb()
    size = _repo_file_size()
    size_text = f"{size / (1024 ** 3):.1f} GB" if size else "unknown"
    spaces_state = "importable" if spaces is not None else "not importable"
    model_state = "Ready" if MODEL_PATH is not None else ("Error" if MODEL_ERROR else "Loads on first run")
    available_text = f"{available_gb:.1f} GB" if available_gb is not None else "unknown"
    server_state = "running" if LLAMA_SERVER_PROCESS is not None and LLAMA_SERVER_PROCESS.poll() is None else "not started"
    settings = MODEL_SETTINGS or {
        "n_ctx": MAX_CONTEXT,
        "n_batch": N_BATCH,
        "n_ubatch": N_UBATCH,
        "n_threads": N_THREADS,
        "n_gpu_layers": _gpu_layers(),
    }
    env = _safe_env_summary()
    cuda_text = env.get("CUDA_VISIBLE_DEVICES", "not visible")
    return f"""**Runtime status:** {model_state}

| Check | Value |
| --- | --- |
| Model | `{MODEL_REPO}` |
| File | `{MODEL_FILE}` ({size_text}) |
| Architecture | Nemotron-H MoE · ~31.6B params · ChatML |
| Runtime | `llama.cpp` `llama-server` `{LLAMA_RELEASE}` |
| ZeroGPU helper | {spaces_state} |
| Available RAM | {available_text} |
| CUDA devices | `{cuda_text}` |
| llama-server | {server_state} |
| Settings | `ctx={settings.get('n_ctx')}`, `batch={settings.get('n_batch')}`, `ubatch={settings.get('n_ubatch')}`, `threads={settings.get('n_threads')}`, `gpu_layers={settings.get('n_gpu_layers')}` |

The first prompt starts `llama-server` and loads the Q8 GGUF; later prompts reuse it.
"""


def _metrics_markdown(meta: dict[str, Any] | None = None) -> str:
    if not meta:
        return "Generation metrics appear after a run."
    return (
        f"Elapsed: `{meta['elapsed']:.2f}s` &middot; "
        f"Tokens: `{meta['completion_tokens']}` &middot; "
        f"Approx `{meta['tokens_per_second']:.2f}` tok/s"
    )


# ---------------------------------------------------------------------------
# Interaction handlers
# ---------------------------------------------------------------------------
def preview_trap(title: str) -> str:
    trap = TRAP_BY_TITLE.get(title)
    if not trap:
        return ""
    return trap_card_html(trap, revealed=False)


def random_trap() -> str:
    return random.choice(TRAP_TITLES)


def _pick_title(title: str) -> str:
    """Zero-input helper: each trap button binds its own title via functools.partial."""
    return title


def run_trap(
    title: str,
    score: dict[str, int] | None,
    max_tokens: int,
    temperature: float,
    top_p: float,
    repeat_penalty: float,
    system_prompt: str,
):
    trap = TRAP_BY_TITLE.get(title)
    if not trap:
        yield "", verdict_banner_html("idle"), "Pick a trap to begin.", scoreboard_html(score), _metrics_markdown(), score
        return

    yield (
        trap_card_html(trap, revealed=False),
        verdict_banner_html("thinking"),
        _thinking_md(),
        scoreboard_html(score),
        "Queued.",
        score,
    )

    try:
        text, meta = _generate(trap["prompt"], system_prompt, max_tokens, temperature, top_p, repeat_penalty)
        reasoning, answer = _split_think(text)
        verdict = judge(trap, answer or reasoning)
        score = _bump(score, verdict)
    except Exception as exc:  # noqa: BLE001 - surface runtime issues in the UI
        reasoning = ""
        answer = (
            "Model load or inference failed.\n\n"
            f"```\n{exc}\n```\n\n"
            "The UI is live and the model artifact is published, but the runtime could not "
            "complete a llama.cpp generation pass. Check the Engine Room status and Space logs."
        )
        meta = None
        verdict = "error"

    yield (
        trap_card_html(trap, revealed=True),
        verdict_banner_html(verdict),
        _answer_md(reasoning, answer),
        scoreboard_html(score),
        _metrics_markdown(meta),
        score,
    )


def run_random(
    score: dict[str, int] | None,
    max_tokens: int,
    temperature: float,
    top_p: float,
    repeat_penalty: float,
    system_prompt: str,
):
    title = random.choice(TRAP_TITLES)
    yield from run_trap(title, score, max_tokens, temperature, top_p, repeat_penalty, system_prompt)


def reset_score():
    s = _new_score()
    return scoreboard_html(s), s


def ask_own(
    message: str,
    mode: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    repeat_penalty: float,
):
    message = (message or "").strip()
    if not message:
        yield "Type a question or a plan first.", _metrics_markdown()
        return
    system = BLINDSPOT_SYSTEM if mode == "Blind-Spot Check" else DEFAULT_SYSTEM
    yield _thinking_md(), "Queued."
    try:
        text, meta = _generate(message, system, max_tokens, temperature, top_p, repeat_penalty)
    except Exception as exc:  # noqa: BLE001
        yield (
            f"**Generation failed.**\n\n```\n{exc}\n```",
            _metrics_markdown(),
        )
        return
    reasoning, answer = _split_think(text)
    label = "🔎 Blind-spot check" if mode == "Blind-Spot Check" else "🧠 First-Principle AI"
    yield f"**{label}:**\n\n{_render_reasoning(reasoning, answer)}", _metrics_markdown(meta)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
:root {
  --fp-bg: #f3f4fb; --fp-panel: #ffffff; --fp-ink: #0f1222; --fp-muted: #5b6478;
  --fp-line: #e3e7f0; --fp-accent: #5b4be6; --fp-accent2: #7c3aed; --fp-accent-d: #4338ca;
  --fp-good: #15a34a; --fp-good-bg: #e7f9ee; --fp-bad: #dc2626; --fp-bad-bg: #fdecec;
  --fp-meh: #c2700a; --fp-meh-bg: #fef3c7;
}
.gradio-container {
  background:
    radial-gradient(900px 360px at 100% -10%, #ece9ff 0%, transparent 55%),
    radial-gradient(800px 300px at -10% 0%, #e6fbf1 0%, transparent 50%),
    var(--fp-bg) !important;
  color: var(--fp-ink) !important;
  font-family: Inter, ui-sans-serif, system-ui, -apple-system, "Segoe UI", sans-serif !important;
}
.fp-shell { max-width: 1180px; margin: 0 auto; padding: 14px 14px 48px; }

/* hero */
.fp-hero {
  position: relative; overflow: hidden; border: 1px solid var(--fp-line);
  background: linear-gradient(180deg, #ffffff, #f7f7ff);
  border-radius: 20px; padding: 30px 30px 26px; margin-bottom: 16px;
  box-shadow: 0 20px 50px rgba(20,20,60,.10);
}
.fp-hero h1 { margin: 0 0 8px; font-size: 36px; line-height: 1.05; letter-spacing: -1px; font-weight: 850; color: var(--fp-ink) !important; }
.fp-hero h1 .fp-grad { color: var(--fp-accent) !important; }
.fp-hero .fp-tag { color: var(--fp-muted); font-size: 16.5px; line-height: 1.5; margin: 0 0 16px; max-width: 800px; }
.fp-hero .fp-tag strong { color: var(--fp-ink); font-weight: 800; }
.fp-badges { display: flex; flex-wrap: wrap; gap: 8px; }
.fp-badge { border: 1px solid var(--fp-line); background: rgba(255,255,255,.85); color: var(--fp-muted); border-radius: 999px; padding: 6px 13px; font-size: 12.5px; }
.fp-badge strong { color: var(--fp-ink); font-weight: 700; }

/* tabs (fix invisible labels) */
.gradio-container .tab-nav { border-bottom: 1px solid var(--fp-line) !important; gap: 2px; }
.gradio-container .tab-nav button { color: var(--fp-muted) !important; font-weight: 650 !important; font-size: 15px !important; background: transparent !important; border: none !important; }
.gradio-container .tab-nav button.selected { color: var(--fp-accent) !important; border-bottom: 2px solid var(--fp-accent) !important; }

/* generic cleanup */
.gradio-container .block { border-color: var(--fp-line) !important; border-radius: 14px !important; box-shadow: none !important; background: transparent !important; }
.gradio-container .gr-group, .gradio-container .form { background: transparent !important; border: none !important; }
textarea, input { background: #fff !important; color: var(--fp-ink) !important; border-color: var(--fp-line) !important; }
.gradio-container code { background: #eef0f8 !important; color: #1c2030 !important; border-radius: 5px; padding: 1px 5px; }
.fp-note { border: 1px solid #d6dbff; background: #eef0ff; color: #3b3a8c; border-radius: 12px; padding: 11px 15px; font-size: 13.5px; line-height: 1.55; margin: 2px 0 14px; }
button.primary { background: linear-gradient(90deg, var(--fp-accent), var(--fp-accent2)) !important; border: none !important; color: #fff !important; font-weight: 700 !important; border-radius: 12px !important; }
button.primary:hover { filter: brightness(1.06); }

/* scoreboard stat tiles */
.fp-score { display: grid; grid-template-columns: repeat(4, 1fr); gap: 10px; margin: 2px 0 14px; }
.fp-stat { border: 1px solid var(--fp-line); background: var(--fp-panel); border-radius: 14px; padding: 12px 14px; box-shadow: 0 6px 16px rgba(20,20,60,.05); }
.fp-stat-n { font-size: 26px; font-weight: 850; line-height: 1; }
.fp-stat-l { font-size: 12px; color: var(--fp-muted); margin-top: 6px; font-weight: 600; }
.fp-stat-good .fp-stat-n { color: var(--fp-good); }
.fp-stat-bad .fp-stat-n { color: var(--fp-bad); }
.fp-stat-meh .fp-stat-n { color: var(--fp-meh); }
.fp-stat-rank-n { font-size: 15.5px; font-weight: 800; color: var(--fp-ink); line-height: 1.18; }

/* trap card grid (buttons) */
.fp-deck-title { font-size: 12.5px; font-weight: 700; color: var(--fp-muted); text-transform: uppercase; letter-spacing: .6px; margin: 4px 0 6px; }
.fp-trapcard { min-height: 58px !important; white-space: normal !important; text-align: left !important; line-height: 1.25 !important;
  border: 1px solid var(--fp-line) !important; background: var(--fp-panel) !important; color: var(--fp-ink) !important;
  border-radius: 14px !important; font-weight: 650 !important; font-size: 14px !important; padding: 12px 14px !important;
  box-shadow: 0 4px 12px rgba(20,20,60,.05) !important;
  transition: transform .12s ease, border-color .12s ease, box-shadow .12s ease !important; }
.fp-trapcard:hover { transform: translateY(-2px) !important; border-color: var(--fp-accent) !important; box-shadow: 0 12px 24px rgba(91,75,230,.18) !important; }
.fp-ghost { background: #fff !important; border: 1px solid var(--fp-line) !important; color: var(--fp-ink) !important; border-radius: 12px !important; font-weight: 650 !important; }
.fp-ghost:hover { border-color: var(--fp-accent) !important; color: var(--fp-accent) !important; }

/* verdict banner */
.fp-banner { display: flex; align-items: center; gap: 14px; border-radius: 16px; padding: 16px 18px; margin: 2px 0 12px;
  border: 1px solid var(--fp-line); animation: fpPop .28s cubic-bezier(.2,.8,.2,1); }
@keyframes fpPop { from { opacity: 0; transform: translateY(8px) scale(.98); } to { opacity: 1; transform: none; } }
.fp-banner-ic { font-size: 30px; line-height: 1; }
.fp-banner-label { font-size: 19px; font-weight: 800; line-height: 1.15; }
.fp-banner-sub { font-size: 13.5px; color: var(--fp-muted); margin-top: 2px; }
.fp-bn-idle { background: #f3f4fb; }
.fp-bn-think { background: #eef0ff; border-color: #d6dbff; }
.fp-bn-good { background: var(--fp-good-bg); border-color: #bbedce; }
.fp-bn-good .fp-banner-label { color: var(--fp-good); }
.fp-bn-bad { background: var(--fp-bad-bg); border-color: #f6c9c9; }
.fp-bn-bad .fp-banner-label { color: var(--fp-bad); }
.fp-bn-meh { background: var(--fp-meh-bg); border-color: #f5dfa0; }
.fp-bn-meh .fp-banner-label { color: var(--fp-meh); }

/* active question card */
.fp-qcard { border: 1px solid var(--fp-line); background: var(--fp-panel); border-radius: 16px; padding: 18px 20px; box-shadow: 0 8px 22px rgba(20,20,60,.06); }
.fp-qcat { font-size: 12px; font-weight: 700; color: var(--fp-accent); text-transform: uppercase; letter-spacing: .5px; }
.fp-qtitle { font-size: 22px; font-weight: 800; margin: 5px 0 10px; letter-spacing: -.3px; }
.fp-qtext { font-size: 17px; line-height: 1.5; color: var(--fp-ink); font-style: italic; }
.fp-reveal { margin-top: 14px; padding-top: 12px; border-top: 1px dashed var(--fp-line); font-size: 13.5px; color: var(--fp-muted); line-height: 1.55; }
.fp-reveal strong { color: var(--fp-accent); }
.fp-reveal .fp-wrong { margin-top: 7px; }
.fp-reveal .fp-wrong strong { color: var(--fp-bad); }

/* answer panel */
.fp-answer { border: 1px solid var(--fp-line); background: var(--fp-panel); border-radius: 16px; padding: 4px 18px; box-shadow: 0 8px 22px rgba(20,20,60,.06); }
.fp-answer details summary { cursor: pointer; color: var(--fp-accent); font-weight: 650; margin: 8px 0; }
.fp-answer details { border-top: 1px dashed var(--fp-line); margin-top: 10px; padding-top: 6px; }

@media (max-width: 920px) { .fp-hero h1 { font-size: 27px; } .fp-score { grid-template-columns: repeat(2, 1fr); } }
"""

HERO = """
<div class="fp-hero">
  <h1>🎯 First-Principle AI <span class="fp-grad">· Gotcha Arena</span></h1>
  <p class="fp-tag">A <strong>small</strong> reasoning model that reads the question, not the pattern.
  Throw it the trick questions that fool the big models — and watch it catch the catch.</p>
  <div class="fp-badges">
    <span class="fp-badge"><strong>~31.6B</strong> Q8 MoE · Nemotron-H</span>
    <span class="fp-badge"><strong>Runs local</strong> · llama.cpp on ZeroGPU</span>
    <span class="fp-badge"><strong>Reasons</strong> from first principles</span>
    <span class="fp-badge">🌲 Thousand Token Wood</span>
  </div>
</div>
"""

SHORT_LABELS = {
    "carwash": "🚗 The 50-Meter Car Wash",
    "ev": "🔌 The Almost-Empty EV",
    "feathers": "🪶 Feathers vs. Steel",
    "moses": "🛶 Moses & the Ark",
    "batball": "⚾ The Bat and the Ball",
    "survivors": "✈️ Bury the Survivors",
    "sally": "👧 Sally's Sisters",
    "brain": "🧠 The 10% Brain Myth",
    "sheep": "🐑 All But Nine",
}

with gr.Blocks(title="First-Principle AI · Gotcha Arena", fill_width=True) as demo:
    with gr.Column(elem_classes=["fp-shell"]):
        gr.HTML(HERO)

        score_state = gr.State(_new_score())

        with gr.Tabs():
            # ---------------- Tab 1: The Gauntlet ----------------
            with gr.Tab("🎯 The Gauntlet"):
                gr.HTML(
                    '<div class="fp-note">Tap a trap card to throw it at the model — or hit '
                    '<strong>🎲 Surprise me</strong>. The banner shows a quick auto-verdict; '
                    'the real proof is the answer underneath.</div>'
                )
                scoreboard = gr.HTML(scoreboard_html(_new_score()))
                with gr.Row(equal_height=False):
                    with gr.Column(scale=5, min_width=300):
                        gr.HTML('<div class="fp-deck-title">The trap deck — tap to play</div>')
                        trap_events = []
                        for _row in range(0, len(TRAPS), 2):
                            with gr.Row():
                                for _trap in TRAPS[_row:_row + 2]:
                                    _b = gr.Button(SHORT_LABELS[_trap["id"]], elem_classes=["fp-trapcard"])
                                    _ts = gr.State(_trap["title"])
                                    trap_events.append((_ts, _b))
                        with gr.Row():
                            surprise = gr.Button("🎲 Surprise me", elem_classes=["fp-ghost"])
                            reset_btn = gr.Button("↺ Reset score", elem_classes=["fp-ghost"])
                    with gr.Column(scale=7, min_width=360):
                        verdict_banner = gr.HTML(verdict_banner_html("idle"))
                        trap_card = gr.HTML(trap_card_html(TRAPS[0], revealed=False))
                        answer_md = gr.Markdown(
                            "Tap a trap card to see how First-Principle AI handles it.",
                            elem_classes=["fp-answer"],
                        )
                        gauntlet_metrics = gr.Markdown(_metrics_markdown())

            # ---------------- Tab 2: Trap It Yourself ----------------
            with gr.Tab("🧠 Trap It Yourself"):
                gr.HTML(
                    '<div class="fp-note"><strong>Two ways to play:</strong> '
                    '<em>Trap the AI</em> — write your own trick question and try to fool it. '
                    '<em>Blind-Spot Check</em> — paste a real plan or "should I X or Y?" decision and the '
                    'model surfaces the hidden assumption you might be missing.</div>'
                )
                mode = gr.Radio(
                    choices=["Trap the AI", "Blind-Spot Check"],
                    value="Trap the AI",
                    label="Mode",
                )
                own_input = gr.Textbox(
                    label="Your question or plan",
                    placeholder="e.g. I need to mail a package. The post box is 30 m away. Should I walk or drive?",
                    lines=3,
                    autofocus=True,
                )
                ask_btn = gr.Button("▶ Ask First-Principle AI", variant="primary")
                gr.Examples(
                    examples=[
                        ["A snail climbs 3 m up a wall each day and slips 2 m each night. The wall is 5 m tall. How many days to reach the top?", "Trap the AI"],
                        ["I have a meeting downtown and a flat tire on my car. The tire shop is 200 m away. Should I walk or drive to the shop?", "Trap the AI"],
                        ["Is it safe to look directly at a solar eclipse if I squint?", "Trap the AI"],
                        ["We're about to launch our app only on iOS to move fast. Should we?", "Blind-Spot Check"],
                        ["I'm going to quit my job next week to go full-time on my side project.", "Blind-Spot Check"],
                    ],
                    inputs=[own_input, mode],
                    label="Try one",
                )
                own_answer = gr.Markdown("Your answer will appear here.", elem_classes=["fp-answer"])
                own_metrics = gr.Markdown(_metrics_markdown())

            # ---------------- Tab 3: Engine Room ----------------
            with gr.Tab("⚙️ Engine Room"):
                gr.HTML(
                    '<div class="fp-note">Sampling controls and live runtime status for the llama.cpp backend. '
                    'These settings apply to every tab.</div>'
                )
                with gr.Row():
                    max_tokens = gr.Slider(64, 2048, value=1024, step=64, label="Max tokens")
                    temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="Temperature")
                with gr.Row():
                    top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
                    repeat_penalty = gr.Slider(1.0, 1.4, value=1.08, step=0.01, label="Repeat penalty")
                system_prompt = gr.Textbox(
                    label="System prompt",
                    value=DEFAULT_SYSTEM,
                    lines=4,
                    max_lines=10,
                )
                status = gr.Markdown(_status_markdown())

    # ----- wiring -----
    run_outputs = [trap_card, verdict_banner, answer_md, scoreboard, gauntlet_metrics, score_state]
    _shared = [score_state, max_tokens, temperature, top_p, repeat_penalty, system_prompt]
    for _ts, _btn in trap_events:
        _btn.click(fn=run_trap, inputs=[_ts, *_shared], outputs=run_outputs, show_progress="minimal")
    surprise.click(fn=run_random, inputs=_shared, outputs=run_outputs, show_progress="minimal")
    reset_btn.click(fn=reset_score, inputs=None, outputs=[scoreboard, score_state])
    ask_btn.click(
        fn=ask_own,
        inputs=[own_input, mode, max_tokens, temperature, top_p, repeat_penalty],
        outputs=[own_answer, own_metrics],
        show_progress="minimal",
    )
    own_input.submit(
        fn=ask_own,
        inputs=[own_input, mode, max_tokens, temperature, top_p, repeat_penalty],
        outputs=[own_answer, own_metrics],
        show_progress="minimal",
    )
    demo.load(fn=_status_markdown, outputs=status)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(css=CSS)