from __future__ import annotations

import json
import re
from statistics import mean
from typing import Any

from .evaluator.reference.prompt_milestones import MILESTONE_TAGS
from .evaluator.shared.questions import _CANDIDATE_BLOCK_RE
from .judge.prompt_judge import JUDGE_SYSTEM, JUDGE_USER
from .shared.json_extract import extract_json
from .shared.observation_format import (
    strip_leaked_reasoning,
)

CHALLENGER_WIN_MARGIN = 0.025

# per-question weights by tag: one entry per milestone category, equal today
TAG_WEIGHTS: dict[str, float] = dict.fromkeys(MILESTONE_TAGS, 1.0)
AMPUTATED_THINKING_MULTIPLIER = 0.5

RESERVED_TOKEN_RE = re.compile(
    r"<\|im_start\|>|<\|im_end\|>|<\|vision_start\|>|<\|vision_end\|>|<\|image_pad\|>"
    r"|<\|video_pad\|>|<\|endoftext\|>|</?tool_call>|<function=|</?tool_response>"
)


def reserved_token_leak(document: str) -> str:
    hit = RESERVED_TOKEN_RE.search(document or "")
    return hit.group(0) if hit else ""


def question_weight(question: dict[str, str]) -> float:
    """What a question is worth, from its tag.

    Every question of a tag is worth the same. A milestone decides which questions get written, not
    what an answer to one is worth, so a milestone that yielded six questions carries three times
    the weight of one that yielded two - that is the intended reading, and it is why nothing here
    normalises per milestone.
    """
    tag = str(question.get("tag") or "")
    default = 0.5 if tag.startswith("reference:") else 0.0
    return TAG_WEIGHTS.get(tag, default)


def amputated_thinking(document: str) -> bool:
    """True when most candidate turns open with an empty </think> (the overfit fingerprint)."""
    blocks = [m.group(1) for m in _CANDIDATE_BLOCK_RE.finditer(document or "")]
    if not blocks:
        return False
    return sum(b.lstrip().startswith("</think>") for b in blocks) / len(blocks) > 0.5


NO_VISIBLE_OUTPUT = "[no visible output]"


def strip_candidate_reasoning(trajectory: str) -> str:
    """Apply strip_leaked_reasoning inside CANDIDATE OUTPUT blocks only.

    Context turns and environment observations are left byte-for-byte intact: the mini-coder-rs
    corpus legitimately carries `</think>` in its own assistant turns.
    """

    def _rewrite(match: re.Match[str]) -> str:
        block, inner = match.group(0), match.group(1)
        start, end = match.start(1) - match.start(0), match.end(1) - match.start(0)
        stripped = strip_leaked_reasoning(inner)
        if not stripped:
            stripped = NO_VISIBLE_OUTPUT if inner.strip() else inner
        return block[:start] + stripped + block[end:]

    return _CANDIDATE_BLOCK_RE.sub(_rewrite, trajectory or "")


def build_judge_messages(*, response: str, questions: list[dict[str, str]]) -> list[dict[str, str]]:
    shown = [
        {
            "id": q["id"],
            "tag": q.get("tag", ""),
            "text": q["text"],
            "example_bad": q.get("example_bad", ""),
        }
        for q in questions
    ]
    cleaned = strip_candidate_reasoning(strip_reply_injection(response)).rstrip()
    return [
        {"role": "system", "content": JUDGE_SYSTEM},
        {
            "role": "user",
            "content": JUDGE_USER.format(
                response=cleaned,
                questions_json=json.dumps(shown, ensure_ascii=False, indent=1),
            ),
        },
    ]


def answer_schema(question_ids: list[str]) -> dict[str, Any]:
    count = len(question_ids)
    return {
        "type": "object",
        "properties": {
            "answers": {
                "type": "array",
                "minItems": count,
                "maxItems": count,
                "items": {
                    "type": "object",
                    "properties": {
                        "asked": {"type": "string", "enum": question_ids},
                        "reason": {"type": "string"},
                        "verdict": {"type": "integer", "enum": [1, 0]},
                    },
                    "required": ["asked", "reason", "verdict"],
                    "additionalProperties": False,
                },
            }
        },
        "required": ["answers"],
        "additionalProperties": False,
    }


_ANSWER_TO_BIT: dict[str, float] = {"1": 1.0, "0": 0.0}


def parse_answers(
    raw: str, question_ids: list[str]
) -> tuple[dict[str, str | None], dict[str, str], bool]:
    obj = extract_json(raw, prefer_keys=("answers",))
    items = obj.get("answers") if isinstance(obj, dict) else obj
    answers: dict[str, str | None] = {qid: None for qid in question_ids}
    explanations: dict[str, str] = {qid: "" for qid in question_ids}
    if isinstance(items, list):
        for item in items:
            if not isinstance(item, dict):
                continue
            qid = str(item.get("asked") or item.get("id") or "").strip()
            value = str(item.get("verdict", item.get("answer", ""))).strip().lower()
            if qid in answers and value in _ANSWER_TO_BIT:
                answers[qid] = value
                explanations[qid] = str(item.get("reason") or item.get("explanation") or "").strip()
    parse_ok = all(value is not None for value in answers.values())
    return answers, explanations, parse_ok


def judge_yes_rate(
    answers: dict[str, str | None], questions: list[dict[str, str]] | None = None
) -> float | None:
    if questions:
        weights = {q["id"]: question_weight(q) for q in questions}
        num = sum(
            weights.get(qid, 0.0) * _ANSWER_TO_BIT[v]
            for qid, v in answers.items()
            if v in _ANSWER_TO_BIT
        )
        den = sum(weights.get(qid, 0.0) for qid, v in answers.items() if v in _ANSWER_TO_BIT)
        if den:
            return round(num / den, 6)
    bits = [_ANSWER_TO_BIT[v] for v in answers.values() if v in _ANSWER_TO_BIT]
    return round(mean(bits), 6) if bits else None


def majority_answers(repeats: list[dict[str, str | None]]) -> dict[str, str | None]:
    """Most common valid answer per question across judgings; the first one on a tie."""
    if not repeats:
        return {}
    out: dict[str, str | None] = {}
    for qid, first in repeats[0].items():
        votes = [r.get(qid) for r in repeats if r.get(qid) in _ANSWER_TO_BIT]
        out[qid] = max(votes, key=votes.count) if votes else first
    return out


def response_score(
    per_judge_answers: dict[str, dict[str, str | None]],
    questions: list[dict[str, str]] | None = None,
) -> float | None:
    rates = [
        r
        for r in (judge_yes_rate(a, questions) for a in per_judge_answers.values())
        if r is not None
    ]
    return round(mean(rates), 6) if rates else None


def challenger_beats_king(score_challenger: float, score_king: float) -> bool:
    return (score_challenger - score_king) >= CHALLENGER_WIN_MARGIN


def aggregate_scores(
    records: list[dict[str, Any]], *, min_valid_fraction: float = 0.8, total: int | None = None
) -> dict[str, Any]:
    total = total or len(records)
    valid = [r for r in records if r.get("scored")]
    valid_count = len(valid)
    judge_errors = sum(
        1
        for record in records
        for result in record.get("judge_results", [])
        if not result.get("parse_ok")
    )

    if total == 0 or valid_count / total < min_valid_fraction:
        return {
            "state": "failed",
            "score_challenger": None,
            "score_king": None,
            "challenger_won": None,
            "valid_turns": valid_count,
            "total_turns": total,
            "judge_errors": judge_errors,
            "scored_sample_count": valid_count,
            "fault_class": "PROVIDER_FAULT",
            "fault_code": "scoring_invalid",
            "fault_message": f"Only {valid_count}/{total} samples valid (< {min_valid_fraction:.0%})",  # noqa: E501
            "retryable": True,
        }

    challenger_mean = round(mean(r["challenger_score"] for r in valid), 6)
    king_mean = round(mean(r["king_score"] for r in valid), 6)

    by_judge: dict[str, float] = {}
    for judge_model in sorted(
        {
            result["judge_model"]
            for record in valid
            for result in record.get("judge_results", [])
            if result.get("side") == "challenger" and result.get("parse_ok")
        }
    ):
        rates = [
            float(result["yes_rate"])
            for record in valid
            for result in record.get("judge_results", [])
            if result.get("side") == "challenger"
            and result.get("judge_model") == judge_model
            and result.get("yes_rate") is not None
        ]
        if rates:
            by_judge[judge_model] = round(mean(rates), 6)

    return {
        "state": "succeeded",
        "score_challenger": challenger_mean,
        "score_king": king_mean,
        "challenger_won": challenger_beats_king(challenger_mean, king_mean),
        "required_win_margin": CHALLENGER_WIN_MARGIN,
        "valid_turns": valid_count,
        "total_turns": total,
        "judge_errors": judge_errors,
        "scored_sample_count": valid_count,
        "by_judge": by_judge,
        "by_metric": {},
        "scoring_mode": "binary",
        "fault_class": None,
        "fault_code": None,
        "fault_message": None,
        "retryable": None,
    }


_INJECTION_PATTERNS: tuple[re.Pattern[str], ...] = (
    re.compile(r'\{\s*"verdict"\s*:\s*"[^"]*"[^}]*\}', re.IGNORECASE),
    re.compile(r'\{\s*"injection"\s*:\s*(true|false)[^}]*\}', re.IGNORECASE),
)
_DASH_RUN_RE = re.compile(r"-{3,}")
_INJECTION_MARKER_RE = re.compile(r'"verdict"\s*:|GRADING\s+INSTRUCTION', re.IGNORECASE)
_VERDICT_LABELS = frozenset({"accept", "weak_pass", "reject"})


def strip_reply_injection(reply: str) -> str:
    cleaned = reply or ""
    dash = _DASH_RUN_RE.search(cleaned)
    if dash and _INJECTION_MARKER_RE.search(cleaned, dash.end()):
        cleaned = cleaned[: dash.start()].rstrip()
    for pattern in _INJECTION_PATTERNS:
        cleaned = pattern.sub("", cleaned)
    if _scan_verdict_json(cleaned):
        return ""
    return cleaned.strip()


def _scan_verdict_json(text: str) -> bool:
    decoder = json.JSONDecoder()
    start = 0
    while True:
        index = text.find("{", start)
        if index == -1:
            return False
        try:
            obj, _ = decoder.raw_decode(text[index:])
        except json.JSONDecodeError:
            start = index + 1
            continue
        if isinstance(obj, dict):
            verdict = obj.get("verdict", "")
            if isinstance(verdict, str) and verdict.lower() in _VERDICT_LABELS:
                return True
            if any(str(key).lower() == "injection" for key in obj):
                return True
        start = index + 1