import copy
import json
from dataclasses import replace

import jsonschema
import pytest
import torch

from pcd.cache import fork_cache, state_tensors
from pcd.config import Limits
from pcd.prompting import prompt_tokens
from pcd.schema import evaluate, validate_schema

SCHEMA = {
    "type": "object",
    "properties": {
        "service": {
            "type": "string",
            "enum": ["express", "express plus", "日本語", 'quote"slash\\'],
        },
        "insured": {
            "type": "boolean",
            "description": "Whether insurance was requested",
        },
    },
    "required": ["service", "insured"],
    "additionalProperties": False,
}


@pytest.mark.parametrize("mode", ["sequence", "token"])
def test_scores_match_uncached_and_microbatch_invariant(engine, mode):
    context = "Use express plus with insurance."
    reference = engine.reference(context, SCHEMA, mode=mode)
    limits = engine.limits
    try:
        for batch_size in (1, 3, 32):
            engine.limits = replace(limits, branch_batch_size=batch_size)
            actual = engine.constrained(context, SCHEMA, mode=mode)
            assert type(actual["object"]["insured"]) is bool
            assert actual["calibrated"] is False
            for field, scores in reference.items():
                torch.testing.assert_close(
                    torch.tensor([c["score"] for c in actual["fields"][field]["candidates"]]),
                    torch.tensor(scores),
                    atol=2e-5,
                    rtol=2e-5,
                )
                assert (
                    abs(sum(c["probability"] for c in actual["fields"][field]["candidates"]) - 1)
                    < 1e-10
                )
            assert (
                actual["forward_calls"] == 1 + (actual["branches"] + batch_size - 1) // batch_size
            )
    finally:
        engine.limits = limits


@torch.inference_mode()
def test_cache_isolation_and_both_convolution_paths(engine):
    assert engine.validate_cache()["prefix_unchanged"]
    base = engine._prefill(engine.tokenizer.encode("prefix is here"))
    snapshot = [t.clone() for t in state_tensors(base)]
    branch = fork_cache(base, 2)
    for tensor in state_tensors(branch):
        old_second = tensor[1].clone()
        tensor[0].zero_()
        assert torch.equal(tensor[1], old_second)
    for old, current in zip(snapshot, state_tensors(base)):
        assert torch.equal(old, current)
    with pytest.raises(ValueError):
        fork_cache(base, 0)


def test_strict_schema_and_output():
    assert len(validate_schema(SCHEMA)) == 2
    variants = []
    for key, value in [
        ("allOf", [{}]),
        ("additionalProperties", True),
        ("required", ["insured"]),
    ]:
        variants.append({**SCHEMA, key: value})
    variants.extend(
        [
            {
                **SCHEMA,
                "properties": {
                    "service": {"type": "string", "enum": ["x", "x"]},
                    "insured": {"type": "boolean"},
                },
            },
            {
                **SCHEMA,
                "properties": {
                    "service": {"type": "string"},
                    "insured": {"type": "boolean"},
                },
            },
            {
                **SCHEMA,
                "properties": {
                    "service": {"type": "number"},
                    "insured": {"type": "boolean"},
                },
            },
        ]
    )
    for schema in variants:
        with pytest.raises((ValueError, jsonschema.SchemaError)):
            validate_schema(schema)
    good = {"service": "express", "insured": False}
    assert evaluate(json.dumps(good), SCHEMA, good)["exact_match"]
    assert not evaluate(json.dumps({**good, "insured": 0}), SCHEMA, good)["schema_compliant"]
    assert not evaluate("```json\n" + json.dumps(good) + "\n```", SCHEMA, good)["syntax_valid"]
    assert not evaluate(json.dumps({**good, "extra": True}), SCHEMA, good)["schema_compliant"]


def test_limits_reject_before_forward(engine, monkeypatch):
    def forbidden(*args, **kwargs):
        raise AssertionError("model must not run")

    monkeypatch.setattr(engine.model.model, "forward", forbidden)
    with pytest.raises(ValueError):
        engine.constrained("x" * 70000, SCHEMA)
    with pytest.raises(ValueError):
        engine.constrained("x", SCHEMA, mode="unknown")
    huge = copy.deepcopy(SCHEMA)
    huge["properties"]["service"]["enum"] = [str(i) for i in range(257)]
    with pytest.raises(ValueError):
        engine.constrained("x", huge)
    with pytest.raises(ValueError):
        Limits(branch_batch_size=0)


def test_atomic_codes_and_reasoning_prefix(engine):
    compiled = engine.compile(SCHEMA, "token")
    for field in compiled.fields:
        assert len(set(field.token_ids)) == len(field.values)
        assert all(len(engine.tokenizer.encode(label)) == 1 for label in field.labels)
    prefix = engine.tokenizer.decode(
        prompt_tokens(engine.tokenizer, compiled, "test", engine.limits)
    )
    assert not prefix.endswith("<think>")
    for field in compiled.fields:
        assert "<think></think>\n" in engine.tokenizer.decode(field.suffix)
    sequence = engine.compile(SCHEMA, "sequence")
    assert engine.tokenizer.decode(
        prompt_tokens(engine.tokenizer, sequence, "test", engine.limits)
    ).endswith("<think></think>\n{\n")
    native = engine.tokenizer.decode(
        prompt_tokens(engine.tokenizer, compiled, "test", engine.limits, native=True)
    )
    assert native.endswith("<think>")


def test_reference_target_shift_uses_preceding_logits(engine):
    context = "express"
    field = engine.compile(SCHEMA, "sequence").fields[0]
    prefix = prompt_tokens(
        engine.tokenizer, engine.compile(SCHEMA, "sequence"), context, engine.limits
    )
    value = field.value_tokens[0]
    tokens = prefix + list(field.suffix + value)
    with torch.inference_mode():
        logits = engine.model(engine.tensor([tokens]), use_cache=False).logits[0].float()
    start = len(prefix) + len(field.suffix)
    expected = sum(float(logits[start + i - 1].log_softmax(-1)[t]) for i, t in enumerate(value))
    actual = engine.reference(context, SCHEMA)[field.name][0]
    assert abs(actual - expected) < 2e-5


def test_token_ar_codec_is_not_json_repair(engine, monkeypatch):
    coded = {f.name: f.labels[0] for f in engine.compile(SCHEMA, "token").fields}
    continuation = json.dumps(coded)[1:]

    def generate(ids, **kwargs):
        return torch.cat([ids, engine.tensor([engine.tokenizer.encode(continuation)])], dim=1)

    monkeypatch.setattr(engine.model, "generate", generate)
    actual = engine.autoregressive("test", SCHEMA, mode="token")
    assert json.loads(actual["text"]) == {"service": "express", "insured": True}