"""Smoke tests for the MLX benchmark adapters and bench functions."""
from __future__ import annotations

import pytest

pytest.importorskip("mlx.core", reason="MLX not available")

import mlx.core as mx

from sgjm.bench.mlx_bench import (
    BenchResult,
    MLXBackboneAdapter,
    MLXDrafterAdapter,
    MLXJudgeAdapter,
    run_ar_bench,
    run_sgjm_bench,
)
from sgjm.training.config import TrainingConfig
from sgjm.training.mlx_backend.model import SGJM


def _smoke_model() -> tuple[SGJM, TrainingConfig]:
    cfg = TrainingConfig.smoke()
    model = SGJM(cfg.model)
    mx.eval(model.parameters())
    return model, cfg


def test_backbone_adapter_encode_returns_valid_state():
    model, cfg = _smoke_model()
    adapter = MLXBackboneAdapter(model)
    state = adapter.encode([10, 20, 30])
    assert state.tokens == (10, 20, 30)
    assert len(state.latent) == cfg.model.d_model
    assert all(isinstance(v, float) for v in state.latent)


def test_backbone_adapter_step_appends_token():
    model, cfg = _smoke_model()
    adapter = MLXBackboneAdapter(model)
    state = adapter.encode([1, 2])
    stepped = adapter.step(state, 99)
    assert stepped.tokens == (1, 2, 99)
    assert len(stepped.latent) == cfg.model.d_model


def test_drafter_adapter_returns_k_samples():
    model, cfg = _smoke_model()
    backbone = MLXBackboneAdapter(model)
    drafter = MLXDrafterAdapter(model, seed=7)
    state = backbone.encode([5, 6, 7, 8])
    samples = drafter.draft(state, k=3, block=cfg.model.block_size)
    assert len(samples) == 3
    for s in samples:
        assert len(s.tokens) == cfg.model.block_size
        assert len(s.latent) == cfg.model.d_model
        assert isinstance(s.log_prob, float)


def test_judge_adapter_returns_scalar():
    model, cfg = _smoke_model()
    judge = MLXJudgeAdapter(model)
    D = cfg.model.d_model
    parent = [0.1] * D
    child = [0.2] * D
    score = judge.score(parent, child)
    assert isinstance(score, float)


def test_run_sgjm_bench_smoke():
    model, cfg = _smoke_model()
    prompt = list(range(16))
    result = run_sgjm_bench(model, cfg.model, prompt, n_steps=2)
    assert isinstance(result, BenchResult)
    assert result.steps_completed >= 1
    assert 0.0 <= result.acceptance_rate <= 1.0
    assert result.elapsed_sec > 0.0


def test_run_ar_bench_smoke():
    model, _ = _smoke_model()
    prompt = list(range(8))
    result = run_ar_bench(model, prompt, n_steps=2)
    assert result.tokens_generated == 2
    assert result.steps_completed == 2
    assert result.elapsed_sec > 0.0
    assert result.tokens_per_sec > 0.0