"""Scorer primitives: compare an extracted prediction against the target.

Contract:  fn(pred: str, target, sample: Sample, ctx: ScoreContext) -> (scores, details)
           scores: {'acc': 1.0}  details: {'acc': {...audit info...}}
Register:  @register_scorer('exact')
Look up:   get_scorer('exact') / make_scorers({'acc': 'exact', ...})

Four scoring paradigms (mirrors the benchmark survey):
  text compare      -- exact / math_equal / em_f1 / alias_match   [implemented]
  LLM-as-judge      -- llm_judge                                  [needs ModelAdapter; wired via ctx.judge]
  code execution    -- execution                                  [needs Sandbox layer; raises NotReady]
  environment reward-- env_reward                                 [needs Agent loop; raises NotReady]
"""

import re
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from pydantic import BaseModel

from ..data.sample import Sample
from .registry import EvalRegistry

ScorerFn = Callable[[str, Any, Sample, 'ScoreContext'], Tuple[Dict[str, float], Dict[str, Any]]]

SCORER_REGISTRY = EvalRegistry('scorer')


def register_scorer(name: str):
    def decorator(fn: ScorerFn) -> ScorerFn:
        SCORER_REGISTRY.register(name, fn)
        return fn

    return decorator


def get_scorer(name: str) -> ScorerFn:
    return SCORER_REGISTRY.get(name)


class ScoreContext(BaseModel):
    """Everything a scorer may need beyond (pred, target, sample).

    judge: optional callable(prompt_messages) -> str. Provided by the runner
    when a ModelAdapter is configured; llm_judge scorers raise NotReady if
    it is None (explicit, never silently wrong).
    """

    class Config:
        arbitrary_types_allowed = True

    judge: Optional[Callable] = None
    judge_model: str = ''
    params: Dict[str, Any] = {}


class LayerNotReady(RuntimeError):
    """A scoring paradigm needs a layer that is not built yet (sandbox/agent)."""


# ------------------------- normalization helpers -------------------------


def _strip_string(s: str) -> str:
    """Light math normalization (subset of Hendrycks/Qwen strip_string)."""
    s = s.strip()
    s = re.sub(r'\\text\{(.+?)\}', r'\1', s)
    s = re.sub(r'\\!|\\,|\\;|\\ ', '', s)
    s = s.replace('\\%', '%').replace('\\$', '$').replace('$', '').replace('%', '')
    s = s.replace('^{\\circ}', '').replace('^\\circ', '')
    s = re.sub(r'(\d),(\d{3})', r'\1\2', s)
    s = re.sub(r'\.0+(?=$|[^0-9])', '', s)
    if len(s) > 1 and s[0] == '{' and s[-1] == '}':
        s = s[1:-1]
    s = s.replace(' ', '').lower()
    return s


def normalize_text(value: str, mode: str = 'math') -> str:
    if mode == 'raw':
        return (value or '').strip()
    if mode == 'numeric':
        v = _strip_string(value)
        try:
            f = float(v.replace(',', ''))
            return str(int(f)) if f == int(f) else str(f)
        except ValueError:
            return v
    return _strip_string(value)


def _targets_list(target) -> List[str]:
    if target is None:
        return []
    if isinstance(target, list):
        return [str(t) for t in target]
    return [str(target)]


# ------------------------- text-compare scorers -------------------------


@register_scorer('exact')
def exact(pred: str, target, sample: Sample, ctx: ScoreContext):
    """Exact match after shared normalization. mode: raw|math|numeric."""
    mode = ctx.params.get('mode', 'math')
    norm = normalize_text(pred or '', mode)
    hit = int(any(norm == normalize_text(t, mode) for t in _targets_list(target)))
    return {'acc': float(hit)}, {'acc': {'mode': mode, 'normalized_pred': norm}}


@register_scorer('math_equal')
def math_equal(pred: str, target, sample: Sample, ctx: ScoreContext):
    """Normalized equality, then optional sympy symbolic equivalence.

    sympy path is tried only when available and only for non-identical
    strings; identical normalized strings short-circuit (fast + exact).
    """
    targets = _targets_list(target)
    norm = _strip_string(pred or '')
    if norm and any(norm == _strip_string(t) for t in targets):
        return {'acc': 1.0}, {'acc': {'path': 'normalized_string'}}

    details: Dict[str, Any] = {'acc': {'path': 'none', 'normalized_pred': norm}}
    if ctx.params.get('sympy', True):
        try:
            from .math_grader import grade_answer  # lazily imported, optional dep

            if any(grade_answer(pred or '', t) for t in targets):
                details['acc'] = {'path': 'sympy'}
                return {'acc': 1.0}, details
        except ImportError:
            details['acc']['sympy'] = 'not installed (pip install sympy pylatexenc)'
    return {'acc': 0.0}, details


def _token_bag(text: str) -> List[str]:
    from string import punctuation

    text = (text or '').lower()
    text = re.sub(r'\b(a|an|the)\b', ' ', text)
    text = re.sub(f'[{re.escape(punctuation)}]', ' ', text)
    return [t for t in text.split() if t]


def _drop_normalize(text: str) -> List[str]:
    """Official DROP normalization: tokenize on space/hyphen, per-token number
    normalization (float str), punctuation strip, article strip, lowercase."""
    from string import punctuation

    out = []
    for token in re.split(r'[ |-]', text or ''):
        token = token.lower()
        if _is_number_official(token):
            token = str(float(token))
        else:
            token = ''.join(c for c in token if c not in punctuation)
        token = re.sub(r'\b(a|an|the)\b', ' ', token)
        token = ' '.join(token.split())
        if token:
            out.append(token)
    return out


def _is_number_official(text: str) -> bool:
    try:
        float(text)
        return True
    except ValueError:
        return False


def _drop_f1(pred_set: set, gold_set: set) -> float:
    intersection = len(gold_set & pred_set)
    precision = intersection / len(pred_set) if pred_set else 1.0
    recall = intersection / len(gold_set) if gold_set else 1.0
    if precision == 0.0 and recall == 0.0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


def _drop_metrics(predicted: List[str], gold: List[str]):
    """Official get_drop_metrics: EM on normalized span sets; F1 via optimal
    1-1 bag alignment (Hungarian), numbers must intersect."""
    pred_norm = [_drop_normalize(p) for p in predicted if (p or '').strip()]
    gold_norm = [_drop_normalize(g) for g in gold if (g or '').strip()]
    if not gold_norm:
        return 0.0, 0.0
    # EM compares the SET of normalized spans (order-insensitive)
    em = 1.0 if pred_norm and set(tuple(p) for p in pred_norm) == set(tuple(g) for g in gold_norm) else 0.0

    pred_bags = [set(' '.join(p).split()) for p in pred_norm]
    gold_bags = [set(' '.join(g).split()) for g in gold_norm]
    try:
        import numpy as np
        from scipy.optimize import linear_sum_assignment
    except ImportError as e:
        raise LayerNotReady("official DROP scoring needs numpy+scipy: pip install 'evalharness[exec]'") from e

    n, m = len(gold_bags), len(pred_bags)
    if m == 0:
        return 0.0, 0.0
    score = np.zeros((n, m))
    for gi, gb in enumerate(gold_bags):
        for pi, pb in enumerate(pred_bags):
            gold_nums = {w for w in gb if _is_number_official(w)}
            pred_nums = {w for w in pb if _is_number_official(w)}
            if not gold_nums or (gold_nums & pred_nums):
                score[gi, pi] = _drop_f1(pb, gb)
    rows, cols = linear_sum_assignment(-score)
    per_bag = [0.0] * max(n, m)
    for r, c in zip(rows, cols):
        per_bag[r] = max(per_bag[r], score[r, c])
    f1 = round(float(np.mean(per_bag)) * 100, 2)
    return em, f1


@register_scorer('em_f1')
def em_f1(pred: str, target, sample: Sample, ctx: ScoreContext):
    """DROP official EM/F1 over gold spans.

    Gold spans are alternative acceptable answers (OR semantics): the
    prediction is split on newlines into predicted spans, then scored
    against EACH gold alternative; the best gold wins (official
    get_drop_metrics called once per alternative).
    """
    import re as _re

    golds: List[List[str]] = []
    if isinstance(target, list):
        # list of alternative answers -> each is one full gold-answer option
        golds = [[str(t)] for t in target]
    else:
        golds = [[str(target)]]
    predicted = [p for p in _re.split(r'\n', pred or '') if p.strip()]
    best_em, best_f1 = 0.0, 0.0
    for gold in golds:
        em, f1 = _drop_metrics(predicted, gold)
        best_em, best_f1 = max(best_em, em), max(best_f1, f1 / 100.0)
    return {'em': best_em, 'f1': best_f1}, {'em': {'alternatives': len(golds)}, 'f1': {'official': True}}


@register_scorer('alias_match')
def alias_match(pred: str, target, sample: Sample, ctx: ScoreContext):
    """TriviaQA style: normalized pred equals any alias, or alias contained."""
    targets = _targets_list(target)
    norm = normalize_text(pred or '', 'math')
    hit = 0
    for t in targets:
        tn = normalize_text(t, 'math')
        if norm == tn or (len(norm) >= 3 and norm in tn):
            hit = 1
            break
    return {'em': float(hit)}, {'em': {'aliases': len(targets)}}


# ------------------------- judge / execution / env (paradigm slots) -------------------------


@register_scorer('llm_judge')
def llm_judge(pred: str, target, sample: Sample, ctx: ScoreContext):
    """LLM-as-judge with an explicit label->scores contract.

    params: prompt_template (with {prediction} {target} {question}),
            labels: {'C': {'acc': 1.0}, 'I': {'acc': 0.0}},
            primary: 'acc'
    """
    if ctx.judge is None:
        raise LayerNotReady(
            "llm_judge needs a judge model; configure a ModelAdapter (runner judge=...) "
            'before running recipes that use it'
        )
    template = ctx.params.get('prompt_template', '{prediction}')
    prompt = template.format(prediction=pred or '', target=target or '', question=sample.input_text)
    raw = ctx.judge([{'role': 'user', 'content': prompt}])
    import json as _json

    try:  # judge callable may return a parsed object
        raw_text = raw if isinstance(raw, str) else _json.dumps(raw)
    except Exception:
        raw_text = str(raw)
    labels: Dict[str, Dict[str, float]] = ctx.params.get('labels') or {}
    upper = (raw_text or '').upper()
    chosen = None
    for label in labels:
        if label.upper() and label.upper() in upper:
            chosen = label
            break
    primary = ctx.params.get('primary', 'acc')
    default_label = ctx.params.get('default_label')  # e.g. SimpleQA's C on parse failure
    if chosen is None and default_label and default_label in labels:
        chosen = default_label
    if chosen is None:
        scores = {k: 0.0 for k in (labels.get(next(iter(labels))) or {})}
        return scores, {primary: {'judge_raw': raw_text[:500], 'parse': 'failed'}}
    return dict(labels[chosen]), {primary: {'judge_label': chosen, 'judge_raw': raw_text[:500]}}


@register_scorer('execution')
def execution(pred: str, target, sample: Sample, ctx: ScoreContext):
    """Run code in a sandbox: sandbox.exec(files built by params['harness']).

    Layering: the sandbox is generic (run files, report exit/stdout/stderr);
    the bench-specific program assembly (completion + tests + checker) is a
    ``harness(sample, pred) -> {filename: content}`` closure provided by the
    recipe. params: harness (required), sandbox ('docker'|'local'),
    entry, timeout_s.
    """
    harness = ctx.params.get('harness')
    if harness is None:
        raise LayerNotReady(
            "execution scorer needs params['harness']: a recipe-provided "
            '(sample, pred) -> {filename: content} builder'
        )
    from ..sandbox import get_sandbox

    sbx = get_sandbox(ctx.params.get('sandbox', 'local'))
    files = harness(sample, pred or '')
    result = sbx.exec(files, entry=ctx.params.get('entry', 'main.py'),
                      timeout_s=ctx.params.get('timeout_s', 30),
                      image=ctx.params.get('image', ''))
    ok = result.ok
    return ({'pass': 1.0} if ok else {'pass': 0.0}), {'pass': {
        'exit_code': result.exit_code,
        'timed_out': result.timed_out,
        'duration_s': result.duration_s,
        'stderr_tail': result.stderr[-400:],
        'sandbox': sbx.name,
    }}


@register_scorer('env_reward')
def env_reward(pred: str, target, sample: Sample, ctx: ScoreContext):
    """Score an agent trajectory by environment final state.

    Consumes ctx.params['env_state'] (set by the runner from the trajectory);
    today: bfcl-style call-sequence comparison against official ground truth.
    tau2/swe get dedicated envs later; this scorer stays the entry point.
    """
    env_state = ctx.params.get('env_state') or {}
    if not env_state:
        raise LayerNotReady(
            'env_reward needs env_state from an agent trajectory '
            '(run with run_eval(loop=True) or a bench env)'
        )
    calls = env_state.get('calls', [])
    backend = ctx.params.get('backend', 'auto')  # auto | bfcl_official | bfcl_native | ...
    if backend != 'native':
        from ..agent.envs.bfcl_mock import BACKEND_REGISTRY, get_backend

        name = backend if backend != 'auto' else 'bfcl_official'
        try:
            verdict = get_backend(name)(calls, sample)
        except KeyError:
            verdict = None
        if verdict is not None:
            return {'acc': float(verdict)}, {'acc': {'mode': name, 'backend': name}}
        if backend != 'auto':
            raise LayerNotReady(f'backend {backend!r} unavailable '
                                '(missing dependency?)')
    gt_calls = (env_state.get('ground_truth') or {}).get('tool_calls')
    if gt_calls is None:
        # irrelevance categories: correct behavior is calling NOTHING
        hit = int(len(calls) == 0)
        return {'acc': float(hit)}, {'acc': {'mode': 'no_calls', 'n_calls': len(calls)}}

    def norm(call: Dict[str, Any]) -> str:
        return json.dumps({'name': call.get('name'),
                           'arguments': call.get('arguments') or call.get('parameters', {})},
                          sort_keys=True, ensure_ascii=False)

    want = [norm(c) for c in gt_calls]
    got = [norm(c) for c in calls]
    hit = int(want == got)
    return {'acc': float(hit)}, {'acc': {
        'mode': 'call_sequence', 'expected': want[:5], 'got': got[:5],
        'n_expected': len(want), 'n_got': len(got),
    }}


# ------------------------- resolution -------------------------

ScorerSpec = Union[str, ScorerFn, Dict[str, Any]]


def make_scorer(metric: str, spec: ScorerSpec) -> ScorerFn:
    """Resolve one metric's scorer spec (name / fn / {'name', **params})."""
    params: Dict[str, Any] = {}
    if isinstance(spec, dict):
        spec = dict(spec)
        params = {k: v for k, v in spec.items() if k != 'name'}
        spec = spec.get('name')
    if callable(spec):
        return spec
    if isinstance(spec, str):
        base = get_scorer(spec)
        if not params:
            return base

        def with_params(pred, target, sample, ctx):
            merged = ScoreContext(judge=ctx.judge, judge_model=ctx.judge_model,
                                  params={**ctx.params, **params})
            return base(pred, target, sample, merged)

        return with_params
    raise TypeError(f'bad scorer spec for {metric!r}: {spec!r}')