"""Focused contract checks for ``generate_mtp(..., capture_trace=True)``."""

import os
import sys
from types import SimpleNamespace

import mlx.core as mx
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sample import generate_mtp  # noqa: E402
from scripts.rollout_metrics import (  # noqa: E402
    token_ids_sha256,
    validate_generation_trace,
)


TRACE_KEYS = {
    "schema_version",
    "prompt_token_count",
    "max_tokens",
    "cycles",
}
CYCLE_KEYS = {
    "cycle_index",
    "output_start",
    "prefix_token_count",
    "prefix_sha256",
    "state_source",
    "bonus_token",
    "draft_attempts",
    "verification",
    "emitted_token_ids",
    "next_state",
}
ATTEMPT_KEYS = {
    "depth",
    "normalized_entropy",
    "issued",
    "draft_token",
}
VERIFICATION_KEYS = {
    "candidate_sha256",
    "base_row_index",
    "projected_row_count",
    "outcomes",
    "rejection_depth",
    "fully_accepted",
    "unused_drafts",
}
OUTCOME_KEYS = {
    "depth",
    "draft_token",
    "target_argmax_token",
    "target_row_index",
    "accepted",
    "emitted_token",
}


class ToyModel:
    """Tiny causal lookup model with a scripted MTP draft distribution."""

    def __init__(self, target_next, diffuse_second_draft=False):
        self.target_next = target_next
        self.diffuse_second_draft = diffuse_second_draft
        self.vocab_size = 12
        self.norm = SimpleNamespace(
            weight=mx.ones((1,), dtype=mx.float32)
        )

    def tok_emb(self, token_ids):
        return token_ids.astype(mx.float32)[..., None]

    def trunk(self, token_ids, mask):
        return token_ids.astype(mx.float32)[..., None], []

    def mtp(self, hidden, token_emb, mask):
        del token_emb, mask
        last = float(np.asarray(hidden[0, -1, 0]))
        marker = -1.0 if last >= 0 else last - 1.0
        return mx.full(hidden.shape, marker, dtype=mx.float32), []

    def head(self, hidden):
        values = np.asarray(hidden, dtype=np.float32)[0, :, 0]
        logits = np.full(
            (1, len(values), self.vocab_size), -10.0, dtype=np.float32
        )
        for row, value in enumerate(values):
            marker = int(round(float(value)))
            if marker < 0:
                if marker <= -2 and self.diffuse_second_draft:
                    logits[0, row, :] = 0.0
                else:
                    logits[0, row, 7] = 10.0
            else:
                logits[0, row, self.target_next[marker]] = 10.0
        return mx.array(logits)


def run(model, max_tokens, depth, policy="fixed", capture_trace=True):
    return generate_mtp(
        model,
        None,
        [1, 2],
        max_tokens,
        depth,
        0.0,
        0,
        1.0,
        0.0,
        np.random.default_rng(123),
        policy=policy,
        entropy_threshold=0.5,
        capture_trace=capture_trace,
    )


def assert_strict_schema(trace):
    assert set(trace) == TRACE_KEYS
    assert trace["schema_version"] == 1
    for index, cycle in enumerate(trace["cycles"]):
        assert set(cycle) == CYCLE_KEYS
        assert cycle["cycle_index"] == index
        assert cycle["state_source"] in {
            "full_recompute",
            "verification_reuse",
        }
        assert cycle["next_state"] in {
            "verification_reuse",
            "full_recompute",
            "terminal",
        }
        for attempt in cycle["draft_attempts"]:
            assert set(attempt) == ATTEMPT_KEYS
        verification = cycle["verification"]
        if verification is None:
            assert cycle["next_state"] == "terminal"
            assert cycle["draft_attempts"] == []
            assert cycle["emitted_token_ids"] == [cycle["bonus_token"]]
            continue
        assert set(verification) == VERIFICATION_KEYS
        for outcome in verification["outcomes"]:
            assert set(outcome) == OUTCOME_KEYS


def assert_offline_validator_accepts(tokens, stats):
    output = tokens[2:]
    validate_generation_trace(
        stats["generation_trace"],
        [1, 2],
        output,
        stats,
        top_level={
            "tokens": stats["tokens"],
            "accepted_drafts": sum(
                stats["rollout_accepted_per_depth"]
            ),
            "corrections": stats["corrections"],
            "drafts_issued": stats["drafts_issued"],
            "draft_recursions": stats["draft_recursions"],
            "verification_forwards": stats["verification_forwards"],
            "target_forwards": stats["target_forwards"],
            "elapsed_seconds": stats["elapsed_seconds"],
        },
        vocab_size=12,
    )


def test_acceptance_reuse_and_bonus_only_terminal():
    model = ToyModel({1: 2, 2: 3, 3: 7, 7: 3})
    tokens, stats = run(model, max_tokens=3, depth=1)
    assert tokens == [1, 2, 3, 7, 3]
    assert stats["tok_per_sec"] == (
        stats["tokens"] / max(stats["elapsed_seconds"], 1e-9)
    )

    trace = stats["generation_trace"]
    assert_strict_schema(trace)
    assert_offline_validator_accepts(tokens, stats)
    assert trace["prompt_token_count"] == 2
    assert trace["max_tokens"] == 3
    assert len(trace["cycles"]) == 2

    first, second = trace["cycles"]
    assert first == {
        "cycle_index": 0,
        "output_start": 0,
        "prefix_token_count": 2,
        "prefix_sha256": token_ids_sha256([1, 2]),
        "state_source": "full_recompute",
        "bonus_token": 3,
        "draft_attempts": [{
            "depth": 1,
            "normalized_entropy": first["draft_attempts"][0][
                "normalized_entropy"
            ],
            "issued": True,
            "draft_token": 7,
        }],
        "verification": {
            "candidate_sha256": token_ids_sha256([1, 2, 3, 7]),
            "base_row_index": 2,
            "projected_row_count": 2,
            "outcomes": [{
                "depth": 1,
                "draft_token": 7,
                "target_argmax_token": 7,
                "target_row_index": 2,
                "accepted": True,
                "emitted_token": 7,
            }],
            "rejection_depth": None,
            "fully_accepted": True,
            "unused_drafts": 0,
        },
        "emitted_token_ids": [3, 7],
        "next_state": "verification_reuse",
    }
    assert second == {
        "cycle_index": 1,
        "output_start": 2,
        "prefix_token_count": 4,
        "prefix_sha256": token_ids_sha256([1, 2, 3, 7]),
        "state_source": "verification_reuse",
        "bonus_token": 3,
        "draft_attempts": [],
        "verification": None,
        "emitted_token_ids": [3],
        "next_state": "terminal",
    }

    _, untraced = run(
        ToyModel({1: 2, 2: 3, 3: 7, 7: 3}),
        max_tokens=1,
        depth=1,
        capture_trace=False,
    )
    assert "generation_trace" not in untraced


def test_rejection_unused_draft_and_recompute():
    model = ToyModel({1: 2, 2: 3, 3: 8, 7: 7, 8: 4})
    tokens, stats = run(model, max_tokens=3, depth=2)
    assert tokens == [1, 2, 3, 8, 4]

    trace = stats["generation_trace"]
    assert_strict_schema(trace)
    assert_offline_validator_accepts(tokens, stats)
    first, second = trace["cycles"]
    assert first["draft_attempts"][0]["draft_token"] == 7
    assert first["draft_attempts"][1]["draft_token"] == 7
    assert first["verification"] == {
        "candidate_sha256": token_ids_sha256([1, 2, 3, 7, 7]),
        "base_row_index": 2,
        "projected_row_count": 3,
        "outcomes": [{
            "depth": 1,
            "draft_token": 7,
            "target_argmax_token": 8,
            "target_row_index": 2,
            "accepted": False,
            "emitted_token": 8,
        }],
        "rejection_depth": 1,
        "fully_accepted": False,
        "unused_drafts": 1,
    }
    assert first["emitted_token_ids"] == [3, 8]
    assert first["next_state"] == "full_recompute"
    assert second["prefix_sha256"] == token_ids_sha256([1, 2, 3, 8])
    assert second["state_source"] == "full_recompute"
    assert second["verification"] is None


def test_adaptive_stop_attempt_is_recorded():
    model = ToyModel(
        {1: 2, 2: 3, 3: 7, 7: 3},
        diffuse_second_draft=True,
    )
    tokens, stats = run(
        model,
        max_tokens=2,
        depth=2,
        policy="adaptive",
    )
    assert tokens == [1, 2, 3, 7]

    trace = stats["generation_trace"]
    assert_strict_schema(trace)
    assert_offline_validator_accepts(tokens, stats)
    cycle = trace["cycles"][0]
    assert cycle["draft_attempts"][0]["issued"] is True
    assert cycle["draft_attempts"][0]["draft_token"] == 7
    stopped = cycle["draft_attempts"][1]
    assert stopped["depth"] == 2
    assert stopped["normalized_entropy"] == 1.0
    assert stopped["issued"] is False
    assert stopped["draft_token"] is None
    assert cycle["verification"]["projected_row_count"] == 2
    assert cycle["verification"]["unused_drafts"] == 0
    assert cycle["next_state"] == "terminal"


def test_terminal_truncation_records_unconsumed_issued_draft():
    model = ToyModel({1: 2, 2: 3, 3: 7, 7: 7})
    tokens, stats = run(model, max_tokens=2, depth=2)
    assert tokens == [1, 2, 3, 7]

    cycle = stats["generation_trace"]["cycles"][0]
    assert_offline_validator_accepts(tokens, stats)
    assert cycle["verification"]["candidate_sha256"] == token_ids_sha256(
        [1, 2, 3, 7, 7]
    )
    assert len(cycle["draft_attempts"]) == 2
    assert len(cycle["verification"]["outcomes"]) == 1
    assert cycle["verification"]["rejection_depth"] is None
    assert cycle["verification"]["fully_accepted"] is False
    assert cycle["verification"]["unused_drafts"] == 1
    assert cycle["emitted_token_ids"] == [3, 7]
    assert cycle["next_state"] == "terminal"


def main():
    test_acceptance_reuse_and_bonus_only_terminal()
    test_rejection_unused_draft_and_recompute()
    test_adaptive_stop_attempt_is_recorded()
    test_terminal_truncation_records_unconsumed_issued_draft()
    print("mtp generation trace: PASS")


if __name__ == "__main__":
    main()