"""Guided decoding: lambda = 0 is greedy, lookahead is exact, regex outputs always match."""

from __future__ import annotations

import random

import pytest
import regex
import torch

from conftest import ARCHS
from decoding.common import KV
from decoding.guided import Constraint, rollout, run_classifier, run_regex

BASE = {"max_new_tokens": 16, "mode": "greedy", "temperature": 1.0, "seed": 0, "rep": 1.0}


class FakeScorer:
    """Scores texts by how many times a marker substring occurs (deterministic, no model)."""

    def __init__(self, marker: str):
        self.marker = marker

    def __call__(self, texts, _target):
        return [float(t.count(self.marker)) for t in texts]


def scorers(marker="e"):
    return {"sentiment": FakeScorer(marker), "topic": FakeScorer(marker)}


@pytest.mark.parametrize("arch", ARCHS)
@pytest.mark.parametrize("lookahead", [0, 2])
def test_lambda_zero_is_greedy(tiny_family, arch, lookahead):
    fam = tiny_family(arch)
    P = {**BASE, "prompt": "Write a sentence.", "attribute": "sentiment", "target": "positive", "lam": 0.0,
         "k": 8, "lookahead": lookahead}
    res = run_classifier(fam, P, scorers())
    assert res["runs"]["method"]["ids"] == res["runs"]["baseline"]["ids"]
    assert res["summary"]["changed"] == 0


def test_large_lambda_steers_toward_marker(tiny_family):
    fam = tiny_family("llama")
    P = {**BASE, "prompt": "Write a sentence.", "attribute": "topic", "topic": "x", "lam": 50.0, "k": 20,
         "lookahead": 0}
    res = run_classifier(fam, P, scorers("e"))
    traj = res["trajectory"]
    assert res["summary"]["changed"] > 0
    assert traj["method"][-1] >= traj["baseline"][-1]


@pytest.mark.parametrize("arch", ARCHS)
def test_batched_rollout_equals_sequential(tiny_family, arch):
    fam = tiny_family(arch)
    prefix = [random.Random(1).randrange(3, 1000) for _ in range(9)]
    kv = KV(fam.main)
    kv.run(prefix)
    cand = torch.tensor([5, 17, 99, 250])
    batched = rollout(kv, cand, 4)
    assert kv.toks == prefix  # the original cache is untouched
    for c, cont in zip(cand.tolist(), batched):
        seq, ref = prefix + [c], KV(fam.main)
        expect = []
        for _ in range(4):
            t = int(ref.run(seq).logits[0, -1].argmax())
            expect.append(t)
            seq = seq + [t]
        assert cont == expect


PATTERNS = [
    r"\d{4}-\d{2}-\d{2}",
    r"(A|B|C|D)",
    r'\{"name": "[a-z]{1,8}", "age": [0-9]{1,3}\}',
    r"[a-z]{1,8}@[a-z]{1,8}\.(com|org)",
]


@pytest.mark.parametrize("arch", ARCHS)
@pytest.mark.parametrize("pattern", PATTERNS)
def test_regex_output_fullmatches(tiny_family, arch, pattern):
    fam = tiny_family(arch)
    P = {**BASE, "max_new_tokens": 40, "prompt": "Answer:", "pattern": pattern, "top_k": 32}
    res = run_regex(fam, P)
    meth = res["runs"]["method"]
    assert res["summary"]["valid"]["method"], (pattern, meth["text"], meth["finish"])
    assert regex.fullmatch(pattern, meth["text"])
    for st in meth["steps"]:
        if st["id"] in fam.stop_ids:
            assert st["x"]["complete"]  # end-of-sequence only after a complete match


def test_unbounded_pattern_output_stays_a_viable_prefix(tiny_family):
    """With `+` a random model may never close the pattern; the output must still be a valid prefix."""
    fam = tiny_family("llama")
    pattern = r"[a-z]+@[a-z]+\.(com|org)"
    res = run_regex(fam, {**BASE, "max_new_tokens": 12, "prompt": "Answer:", "pattern": pattern, "top_k": 32})
    text = res["runs"]["method"]["text"]
    assert regex.fullmatch(pattern, text, partial=True)
    assert res["summary"]["valid"]["method"] or res["runs"]["method"]["finish"] == "length"


def test_constraint_rejects_bad_patterns():
    with pytest.raises(ValueError):
        Constraint("(unclosed")
    with pytest.raises(ValueError):
        Constraint("a" * 400)
    con = Constraint(r"\d{3}")
    assert con.viable("12") and not con.viable("1a") and con.complete("123") and not con.complete("12")