#!/usr/bin/env python3
"""Evaluate L2 skill decisions from belief states.

The first runnable backends are deterministic baselines. The Gemma backend is a
strict-JSON wrapper scaffold; it is intentionally gated behind an explicit model
path/endpoint so experiments remain reproducible.
"""

from __future__ import annotations

import re
import argparse
import copy
import json
import math
import os
import random
import sys
import time
import urllib.request
from collections import Counter, defaultdict
from typing import Any, Dict, List, Optional

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from layered_belief import (
    BOSSES,
    CATEGORY,
    build_split_manifest,
    collect_decision_samples,
    js_divergence,
    macro_recall,
    normalized_entropy,
    skill_grounded,
)


def choose_majority(sample: Dict[str, Any], train_majority: Dict[str, str]) -> Dict[str, Any]:
    boss = sample["boss"]
    skill = train_majority[boss]
    return {"skill": skill, "reason": "train-majority skill", "confidence": 1.0}


def choose_random(sample: Dict[str, Any], rng: random.Random) -> Dict[str, Any]:
    ready = sample["belief"]["cooldown_ready"]
    legal = [s for s in sample["legal_skills"] if ready.get(s, True)]
    skill = rng.choice(legal or sample["legal_skills"])
    return {"skill": skill, "reason": "uniform random legal-ready skill", "confidence": 1.0}


def choose_heuristic(sample: Dict[str, Any]) -> Dict[str, Any]:
    belief = sample["belief"]
    ready = belief["cooldown_ready"]
    dist = belief["player_distance_bin"]
    style = str(belief.get("director_style", "") or "").lower()
    legal = [s for s in sample["legal_skills"] if ready.get(s, True)]
    if not legal:
        legal = list(sample["legal_skills"])
    if any(x in style for x in ("keep_distance", "kite", "defensive", "cautious", "sniper")):
        preferred_cat = "far"
    elif "aggressive" in style or "pressure" in style:
        preferred_cat = "close"
    else:
        preferred_cat = "far" if dist == "far" else "close"
    for s in legal:
        if CATEGORY.get(s) == preferred_cat:
            return {"skill": s, "reason": f"{dist} distance -> {preferred_cat} skill", "confidence": 0.7}
    for s in legal:
        if CATEGORY.get(s) in {"cd_aoe", "summon"}:
            return {"skill": s, "reason": "no preferred range skill; use ready special skill", "confidence": 0.55}
    return {"skill": legal[0], "reason": "fallback first legal-ready skill", "confidence": 0.4}


def extract_json_object(text: str) -> Dict[str, Any]:
    text = text.strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    start = text.find("{")
    end = text.rfind("}")
    if start >= 0 and end > start:
        return json.loads(text[start:end + 1])
    raise ValueError(f"no JSON object found in model output: {text[:200]!r}")


_HF_MODEL: Any = None
_HF_TOKENIZER: Any = None


def _is_hp_key(k: Any) -> bool:
    s = str(k).lower()
    return ("hp_phase" in s) or ("hp_frac" in s) or s in {"hp", "your_hp"}


_HP_TEXT = re.compile(r"\s*,?\s*HP[ _]phase\s*,?", re.IGNORECASE)


def _scrub_hp_text(s: str) -> str:
    # 擦掉 prose 里对 "HP phase" 的提及（如旧 instruction 字符串），并修好逗号
    out = _HP_TEXT.sub(", ", s)
    return out.replace(" ,", ",").replace(", ,", ",").replace(",,", ",")


def _strip_hp(obj: Any) -> Any:
    """hp_phase 已从 L2 全面移除（L1 无 hp 头、闭环无来源）。递归剔除任何 hp 相关键
    （hp_phase / your_hp_phase / *_hp_frac 等）并擦掉字符串值里的 HP phase 提及，
    保证喂给模型的 prompt 里绝不含 hp。"""
    if isinstance(obj, dict):
        return {k: _strip_hp(v) for k, v in obj.items() if not _is_hp_key(k)}
    if isinstance(obj, list):
        return [_strip_hp(v) for v in obj]
    if isinstance(obj, str):
        return _scrub_hp_text(obj)
    return obj


def build_gemma_prompt(sample: Dict[str, Any], feedback: str | None = None) -> str:
    if sample.get("l2_prompt_payload"):
        payload = _strip_hp(copy.deepcopy(sample["l2_prompt_payload"]))
        if feedback:
            payload["previous_error"] = feedback
            payload["repair_instruction"] = "Return a valid JSON object only. Do not include markdown."
        return json.dumps(payload, ensure_ascii=False)
    belief = _strip_hp(sample["belief"])
    legal = sample["legal_skills"]
    ready = belief["cooldown_ready"]
    director_style = belief.get("director_style")
    payload = {
        "task": "Choose the next boss skill for a game NPC.",
        "constraints": {
            "output_json_only": True,
            "schema": {"skill": "string", "reason": "string", "confidence": "number"},
            "legal_skills": legal,
            "cooldown_ready": ready,
            "confidence_range": [0.0, 1.0],
        },
        "belief": belief,
        "director_style": director_style,
        "instruction": (
            "Pick exactly one skill from legal_skills. Prefer a skill that is ready, grounded in "
            "the distance/angle/cooldown belief, and plausible for this boss. Return only JSON."
        ),
    }
    if feedback:
        payload["previous_error"] = feedback
        payload["repair_instruction"] = "Return a valid JSON object only. Do not include markdown."
    return json.dumps(payload, ensure_ascii=False)


def gemma_openai_payload(user_content: Any, max_tokens: int = 256) -> Dict[str, Any]:
    payload = {
        "model": os.environ.get("GEMMA_MODEL", "gemma4-e2b-it"),
        "temperature": float(os.environ.get("GEMMA_TEMPERATURE", "0")),
        "max_tokens": int(max_tokens),
        "messages": [
            {
                "role": "system",
                "content": "You are a strict JSON decision module for a game boss. Output JSON only."
            },
            {"role": "user", "content": user_content}
        ]
    }
    if os.environ.get("GEMMA_RESPONSE_FORMAT", "1") != "0":
        payload["response_format"] = {"type": "json_object"}
    return payload


def call_openai_compatible(prompt: str) -> str:
    base = os.environ.get("GEMMA_OPENAI_BASE_URL")
    if not base:
        raise RuntimeError(
            "OpenAI-compatible Gemma backend requires GEMMA_OPENAI_BASE_URL. "
            "Example: http://127.0.0.1:8000/v1"
        )
    api_key = os.environ.get("GEMMA_API_KEY", "EMPTY")
    payload = gemma_openai_payload(prompt)
    req = urllib.request.Request(
        base.rstrip("/") + "/chat/completions",
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        },
        method="POST"
    )
    # 并发发请求时，单个请求会在 server 端排队，偶发超时是正常的 —— 重试而不是让整个 run 崩掉。
    timeout = float(os.environ.get("GEMMA_TIMEOUT", "600"))
    attempts = int(os.environ.get("GEMMA_RETRIES", "4"))
    last_exc = None
    for attempt in range(attempts):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                data = json.loads(resp.read().decode("utf-8"))
            return data["choices"][0]["message"]["content"]
        except Exception as exc:  # 超时 / 连接被重置 / 服务端瞬时错误
            last_exc = exc
            if attempt < attempts - 1:
                time.sleep(2.0 * (attempt + 1))
    raise RuntimeError(f"gemma endpoint failed after {attempts} attempts: {last_exc}")


def call_hf_local(prompt: str) -> str:
    global _HF_MODEL, _HF_TOKENIZER
    model_path = os.environ.get("GEMMA_MODEL_PATH")
    if not model_path:
        raise RuntimeError("HF Gemma backend requires GEMMA_MODEL_PATH to point to a local model directory or HF id")
    if _HF_MODEL is None or _HF_TOKENIZER is None:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        dtype_name = os.environ.get("GEMMA_TORCH_DTYPE", "auto")
        dtype = "auto"
        if dtype_name == "bfloat16":
            dtype = torch.bfloat16
        elif dtype_name == "float16":
            dtype = torch.float16
        _HF_TOKENIZER = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=os.environ.get("GEMMA_TRUST_REMOTE_CODE", "0") == "1",
        )
        _HF_MODEL = AutoModelForCausalLM.from_pretrained(
            model_path,
            device_map=os.environ.get("GEMMA_DEVICE_MAP", "auto"),
            torch_dtype=dtype,
            trust_remote_code=os.environ.get("GEMMA_TRUST_REMOTE_CODE", "0") == "1",
        )
        _HF_MODEL.eval()
    messages = [
        {"role": "system", "content": "You are a strict JSON decision module for a game boss. Output JSON only."},
        {"role": "user", "content": prompt},
    ]
    if hasattr(_HF_TOKENIZER, "apply_chat_template") and _HF_TOKENIZER.chat_template:
        text = _HF_TOKENIZER.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    else:
        text = messages[0]["content"] + "\n" + messages[1]["content"] + "\nJSON:"
    inputs = _HF_TOKENIZER(text, return_tensors="pt")
    inputs = {k: v.to(_HF_MODEL.device) for k, v in inputs.items()}
    temperature = float(os.environ.get("GEMMA_TEMPERATURE", "0"))
    with __import__("torch").no_grad():
        out = _HF_MODEL.generate(
            **inputs,
            max_new_tokens=int(os.environ.get("GEMMA_MAX_NEW_TOKENS", "192")),
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-6),
            top_p=float(os.environ.get("GEMMA_TOP_P", "0.95")),
            pad_token_id=_HF_TOKENIZER.eos_token_id,
        )
    return _HF_TOKENIZER.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)


def call_gemma_model(prompt: str) -> str:
    if os.environ.get("GEMMA_OPENAI_BASE_URL"):
        return call_openai_compatible(prompt)
    if os.environ.get("GEMMA_MODEL_PATH"):
        return call_hf_local(prompt)
    raise RuntimeError(
        "Gemma backend requires either GEMMA_OPENAI_BASE_URL for an OpenAI-compatible endpoint "
        "or GEMMA_MODEL_PATH for a local transformers model. Neither is set."
    )


def normalize_decision(out: Dict[str, Any]) -> Dict[str, Any]:
    skill = out.get("skill")
    reason = out.get("reason", "")
    confidence = out.get("confidence", 0.0)
    try:
        confidence = float(confidence)
    except (TypeError, ValueError):
        confidence = 0.0
    if not math.isfinite(confidence):
        confidence = 0.0
    return {
        "skill": skill,
        "reason": str(reason),
        "confidence": max(0.0, min(1.0, confidence)),
    }


def choose_gemma(sample: Dict[str, Any]) -> Dict[str, Any]:
    feedback: Optional[str] = None
    last_raw = ""
    last_error = ""
    retries = int(os.environ.get("GEMMA_PARSE_RETRIES", "1"))
    for attempt in range(retries + 1):
        prompt = build_gemma_prompt(sample, feedback)
        raw = call_gemma_model(prompt)
        last_raw = raw
        try:
            decision = normalize_decision(extract_json_object(raw))
            decision["raw_output"] = raw
            if os.environ.get("GEMMA_LOG_PROMPT", "1") != "0":
                decision["raw_prompt"] = prompt
            decision["parse_attempts"] = attempt + 1
            return decision
        except Exception as exc:
            last_error = str(exc)
            feedback = f"Attempt {attempt + 1} failed: {last_error}. Raw output: {raw[:500]}"
    return {
        "skill": None,
        "reason": "",
        "confidence": 0.0,
        "raw_output": last_raw,
        "parse_error": last_error or "unknown parse error",
        "parse_attempts": retries + 1,
    }


def decide(
    sample: Dict[str, Any],
    backend: str,
    majority: Dict[str, str],
    rng: random.Random,
) -> Dict[str, Any]:
    if backend == "oracle":
        return {"skill": sample["target_skill"], "reason": "oracle engine target", "confidence": 1.0}
    if backend == "majority":
        return choose_majority(sample, majority)
    if backend == "random":
        return choose_random(sample, rng)
    if backend == "heuristic":
        return choose_heuristic(sample)
    if backend == "gemma":
        return choose_gemma(sample)
    raise ValueError(backend)


def train_majority(samples: List[Dict[str, Any]]) -> Dict[str, str]:
    by_boss = defaultdict(Counter)
    for s in samples:
        by_boss[s["boss"]][s["target_skill"]] += 1
    return {boss: ctr.most_common(1)[0][0] for boss, ctr in by_boss.items()}


def flatten_nested_l1_belief(row: Dict[str, Any]) -> Dict[str, Any]:
    belief = row["l1_belief"]
    geom = belief.get("geometry", {}) or {}
    player = belief.get("player_state", {}) or {}
    boss = belief.get("boss_state", {}) or {}
    resource = belief.get("resource_state", {}) or {}
    tactical = belief.get("tactical_state", {}) or {}
    confidence = belief.get("confidence", {}) or {}
    return {
        "boss_id": row["boss"],
        "fight": row.get("fight"),
        "index": row.get("index"),
        "player_distance_bin": geom.get("distance_bin"),
        "player_distance_value": geom.get("distance_value"),
        "dp_bin": geom.get("dp_bin"),
        "player_angle_bin": geom.get("angle_bin"),
        "player_angle_value": geom.get("angle_value"),
        "front_cone": geom.get("front_cone"),
        "decision_zone": geom.get("decision_zone"),
        "tactical_sector": geom.get("tactical_sector"),
        "behind": geom.get("behind"),
        "player_action": player.get("action", "unknown"),
        "prev_boss_skill": boss.get("prev_skill"),
        "skill_phase": boss.get("skill_phase", "decision"),
        "skill_finished": boss.get("skill_finished", True),
        "cooldown_ready": resource.get("cooldown_ready", {}),
        "cooldown_seconds": resource.get("cooldown_seconds", {}),
        "hp_phase": resource.get("hp_phase"),
        "skill_family_prior": tactical.get("skill_family_prior", {}),
        "top_skill_family": tactical.get("top_skill_family"),
        "confidence": confidence.get("overall", 0.0) if isinstance(confidence, dict) else confidence,
        "prediction_source": "nested_l1_handoff",
    }


def load_belief_overrides(path: str) -> Dict[tuple, Dict[str, Any]]:
    overrides = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            row = json.loads(line)
            key = (row["boss"], int(row["fight"]), int(row["index"]))
            if "belief" in row:
                overrides[key] = {"belief": row["belief"]}
            elif "l1_belief" in row:
                overrides[key] = {
                    "belief": flatten_nested_l1_belief(row),
                    "l1_belief": row.get("l1_belief"),
                    "l2_prompt_payload": row.get("l2_prompt_payload"),
                }
            else:
                raise KeyError(f"{path} row for {key} has neither 'belief' nor 'l1_belief'")
    return overrides


def apply_belief_overrides(samples: List[Dict[str, Any]], path: str) -> int:
    overrides = load_belief_overrides(path)
    matched = 0
    for sample in samples:
        key = (sample["boss"], int(sample["fight"]), int(sample["index"]))
        if key in overrides:
            sample["belief"] = overrides[key]["belief"]
            if overrides[key].get("l1_belief") is not None:
                sample["l1_belief"] = overrides[key]["l1_belief"]
            if overrides[key].get("l2_prompt_payload") is not None:
                sample["l2_prompt_payload"] = overrides[key]["l2_prompt_payload"]
            matched += 1
    return matched


def write_prediction_rows(path: str, rows: List[Dict[str, Any]]) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    tmp = f"{path}.tmp.{os.getpid()}"
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            for row in rows:
                sample = row["sample"]
                out = {
                    "boss": sample["boss"],
                    "fight": sample["fight"],
                    "index": sample["index"],
                    "target_skill": sample["target_skill"],
                    "legal_skills": sample["legal_skills"],
                    "belief": sample["belief"],
                    "l1_belief": sample.get("l1_belief"),
                    "l2_prompt_payload": sample.get("l2_prompt_payload"),
                    "decision": row["decision"],
                    "valid": row["valid"],
                }
                f.write(json.dumps(out, ensure_ascii=False) + "\n")
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, path)
    finally:
        if os.path.exists(tmp):
            os.unlink(tmp)


def validate_decision(sample: Dict[str, Any], decision: Dict[str, Any]) -> Dict[str, Any]:
    json_valid = isinstance(decision, dict) and not decision.get("parse_error")
    skill = decision.get("skill") if isinstance(decision, dict) else None
    legal = skill in sample["legal_skills"]
    ready = bool(sample["belief"]["cooldown_ready"].get(skill, True)) if legal else False
    grounded = legal and skill_grounded(skill, sample["belief"])
    return {
        "json_valid": json_valid,
        "schema_valid": json_valid and isinstance(skill, str) and isinstance(decision.get("reason", ""), str),
        "legal": legal,
        "cooldown_ok": ready,
        "grounded": grounded
    }


def evaluate(
    manifest: Dict[str, Any],
    backend: str,
    split: str,
    boss: str | None,
    seed: int,
    limit: int | None = None,
    beliefs_path: str | None = None,
    predictions_out: str | None = None,
) -> Dict[str, Any]:
    rng = random.Random(seed)
    samples = collect_decision_samples(manifest, split, boss)
    train = collect_decision_samples(manifest, "train", boss)
    majority = train_majority(train)
    belief_override_matches = 0
    if beliefs_path:
        belief_override_matches = apply_belief_overrides(samples, beliefs_path)
    if limit is not None:
        samples = samples[:limit]
    rows = []
    for i, sample in enumerate(samples):
        dec = decide(sample, backend, majority, rng)
        val = validate_decision(sample, dec)
        rows.append({"sample": sample, "decision": dec, "valid": val})
    if predictions_out:
        write_prediction_rows(predictions_out, rows)

    y_true = [r["sample"]["target_skill"] for r in rows]
    y_pred = [r["decision"].get("skill") for r in rows]
    labels = sorted({s for r in rows for s in r["sample"]["legal_skills"]})
    legal_pred = [p for p, r in zip(y_pred, rows) if r["valid"]["legal"]]
    result = {
        "backend": backend,
        "split": split,
        "boss": boss or "all",
        "seed": seed,
        "belief_source": beliefs_path or "oracle_engine_belief",
        "belief_override_matches": belief_override_matches,
        "n": len(rows),
        "macro_recall": round(macro_recall(y_true, y_pred, labels), 4),
        "json_valid_rate": round(sum(r["valid"]["json_valid"] for r in rows) / max(1, len(rows)), 4),
        "schema_valid_rate": round(sum(r["valid"]["schema_valid"] for r in rows) / max(1, len(rows)), 4),
        "invalid_skill_rate": round(1 - sum(r["valid"]["legal"] for r in rows) / max(1, len(rows)), 4),
        "rule_cooldown_violation_rate": round(1 - sum(r["valid"]["cooldown_ok"] for r in rows) / max(1, len(rows)), 4),
        "rule_grounding_violation_rate": round(1 - sum(r["valid"]["grounded"] for r in rows) / max(1, len(rows)), 4),
        "normalized_entropy": round(normalized_entropy(y_pred, labels), 4),
        "legal_grounded_entropy": round(
            normalized_entropy(
                [r["decision"]["skill"] for r in rows if r["valid"]["legal"] and r["valid"]["grounded"]],
                labels,
            ),
            4,
        ),
        "effective_skill_count": round(math.exp(normalized_entropy(y_pred, labels) * math.log(max(2, len(labels)))), 4),
        "max_skill_frequency": round(Counter(y_pred).most_common(1)[0][1] / max(1, len(y_pred)), 4),
        "js_to_real_distribution": round(js_divergence(y_pred, y_true, labels), 4),
        "prediction_counts": dict(Counter(y_pred).most_common()),
    }
    return result


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", default="out/layered/manifest.json")
    ap.add_argument("--backend", choices=["oracle", "majority", "random", "heuristic", "gemma"], default="heuristic")
    ap.add_argument("--split", choices=["train", "val", "test"], default="test")
    ap.add_argument("--boss", choices=BOSSES, default=None)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--limit", type=int, default=None, help="limit examples for Gemma smoke tests")
    ap.add_argument("--beliefs", default=None, help="optional JSONL predicted beliefs keyed by boss/fight/index")
    ap.add_argument("--predictions_out", default=None, help="optional JSONL row-level decisions")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    if os.path.exists(args.manifest):
        with open(args.manifest, encoding="utf-8") as f:
            manifest = json.load(f)
    else:
        manifest = build_split_manifest()
    result = evaluate(
        manifest,
        args.backend,
        args.split,
        args.boss,
        args.seed,
        args.limit,
        args.beliefs,
        args.predictions_out,
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))
    limit_tag = f"_n{args.limit}" if args.limit is not None else ""
    belief_tag = ""
    if args.beliefs:
        stem = os.path.splitext(os.path.basename(args.beliefs))[0]
        belief_tag = f"_belief-{stem}"
    out = args.out or f"out/layered/l2_{args.backend}_{args.split}_{args.boss or 'all'}{belief_tag}{limit_tag}.json"
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"wrote {out}")


if __name__ == "__main__":
    main()