"""ThirtySpokes KOTH routing agent — cost-aware cascade + optional verify.

Contract required by the measured runtime:
  build_agent(weights) -> agent
  agent(prompt, call_model) -> answer

`call_model(model, messages, params)` is the only channel to the owner-pinned pool.
Answers must come from pool responses (grounding). Always prefer cheaper models and
escalate only when heuristics say the ask is hard.
"""

from __future__ import annotations

import json
import re

# Match the owner's pool-reference generation settings (see docs/MINER.md).
PARAMS = {"max_tokens": 16384, "reasoning": {"effort": "low"}}

# Signals that LiveCodeBench / hard reasoning usually needs a stronger model.
_HARD_MARKERS = (
    "prove",
    "proof",
    "leetcode",
    "write a function",
    "write a program",
    "implement",
    "algorithm",
    "complexity",
    "def ",
    "class ",
    "```",
    "stdin",
    "stdout",
    "test case",
    "constraints",
    "hard",
)

_CODE_MARKERS = (
    "python",
    "code",
    "function",
    "program",
    "algorithm",
    "compile",
    "runtime error",
)

_WEAK_ANSWER = re.compile(
    r"^\s*(i co(?:uld|n't)|sorry|as an ai|cannot|can't help|not sure|n/?a)\b",
    re.I,
)


def _load_cfg(weights: bytes) -> dict:
    try:
        cfg = json.loads(weights.decode() if isinstance(weights, (bytes, bytearray)) else weights)
    except Exception:
        cfg = {}
    if not isinstance(cfg, dict):
        cfg = {}
    # Defaults align with MockPool model ids for local smoke; override via weights.bin for mainnet.
    cheap = cfg.get("cheap") or cfg.get("model") or "cheap"
    mid = cfg.get("mid") or cheap
    strong = cfg.get("strong") or mid
    return {
        "cheap": cheap,
        "mid": mid,
        "strong": strong,
        "escalate_chars": int(cfg.get("escalate_chars", 400)),
        "verify": bool(cfg.get("verify", True)),
        "verify_on_code": bool(cfg.get("verify_on_code", True)),
        "max_calls": int(cfg.get("max_calls", 3)),
    }


def _looks_hard(prompt: str, escalate_chars: int) -> bool:
    p = prompt or ""
    pl = p.lower()
    if len(p) >= escalate_chars:
        return True
    if any(m in pl for m in _HARD_MARKERS):
        return True
    # Many coding tasks include multi-line specs / sample IO.
    if pl.count("\n") >= 8 and any(m in pl for m in _CODE_MARKERS):
        return True
    return False


def _looks_code(prompt: str) -> bool:
    pl = (prompt or "").lower()
    return any(m in pl for m in _CODE_MARKERS) or "```" in pl or "def " in pl


def _is_weak(answer: str) -> bool:
    if not answer or not str(answer).strip():
        return True
    text = str(answer).strip()
    if len(text) < 8:
        return True
    if _WEAK_ANSWER.search(text):
        return True
    return False


def _call(call_model, model: str, prompt: str) -> str:
    return call_model(model, [{"role": "user", "content": prompt}], PARAMS)


def build_agent(weights):
    cfg = _load_cfg(weights)
    cheap, mid, strong = cfg["cheap"], cfg["mid"], cfg["strong"]
    escalate_chars = cfg["escalate_chars"]
    do_verify = cfg["verify"]
    verify_on_code = cfg["verify_on_code"]
    max_calls = max(1, cfg["max_calls"])

    def agent(prompt, call_model):
        calls = 0
        hard = _looks_hard(prompt, escalate_chars)
        codey = _looks_code(prompt)

        # 1) Cheap first on easy asks; mid/strong first on hard asks to avoid wasted cheap fails.
        first = mid if hard else cheap
        ans = _call(call_model, first, prompt)
        calls += 1
        chosen = first

        # 2) Escalate if the cheap/mid answer looks weak or the ask is hard and we started cheap.
        if calls < max_calls and (_is_weak(ans) or (hard and chosen == cheap)):
            nxt = strong if chosen != strong else mid
            if nxt != chosen:
                ans = _call(call_model, nxt, prompt)
                calls += 1
                chosen = nxt

        # 3) Optional verify: on code/hard tasks, ask a different model and keep the longer
        # grounded pool response when the first looks thin. Still returns pool text only.
        if (
            do_verify
            and calls < max_calls
            and (hard or (verify_on_code and codey))
            and chosen != strong
        ):
            alt = _call(call_model, strong, prompt)
            calls += 1
            if _is_weak(ans) or (not _is_weak(alt) and len(str(alt)) > len(str(ans)) * 1.15):
                ans = alt

        # Final grounding fallback: never return empty — one more strong call if needed.
        if _is_weak(ans) and calls < max_calls and chosen != strong:
            ans = _call(call_model, strong, prompt)

        return ans

    return agent