from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from ..expressions import AssertExpr, AssertTransformError, compile_assert_expr, get_assert_expr_names
from ..transform import (
    StateDictProvider,
    TypedTransform,
    TransformResult,
    ensure_mapping_payload,
    register_transform,
)


@dataclass(frozen=True)
class AssertSpec:
    expr: AssertExpr

    def collect_models(self) -> set[str]:
        return self.expr.collect_models()


class AssertTransform(TypedTransform[AssertSpec]):
    name = "assert"
    error_type = AssertTransformError
    spec_type = AssertSpec
    help_text = (
        "Checks conditions on tensors using a single assert expression. "
        "The operation fails if the assertion does not hold and does not modify tensors.\n"
        "\n"
        "Expressions can match tensors by name or pattern and may include slicing "
        "(written after '::') and nesting.\n"
        "\n"
        "Examples:\n"
        "  assert: { equal: { left: a.weight, right: b.weight } }\n"
        "  assert: { not: { equal: { left: a.weight, right: b.weight } } }\n"
        "  assert: { dtype: { of: ln_f.weight, is: float32 } }\n"
        "  assert: { shape: { of: 'ln_f.weight::[:8]', is: [8] } }\n"
        "  assert: { all: [ { exists: '.*weight' }, { dimensions: { of: ln_f.weight, is: 1 } } ] }"
    )

    def completion_payload_start_candidates(self, prefix_text: str) -> list[str] | None:
        candidates = [name for name in get_assert_expr_names() if name.startswith(prefix_text)]
        if not prefix_text:
            return candidates + ["{ "]
        return candidates

    def completion_key_candidates(self, before_cursor: str, prefix_text: str) -> list[str] | None:
        del before_cursor
        return [f"{name}: " for name in get_assert_expr_names() if name.startswith(prefix_text)]

    def compile(self, payload: Any, default_model: str | None) -> AssertSpec:
        payload = ensure_mapping_payload(payload, self.name)
        expr = compile_assert_expr(payload, default_model)
        return AssertSpec(expr=expr)

    def apply(self, spec: object, provider: StateDictProvider) -> TransformResult:
        typed = self.require_spec(spec)
        typed.expr.evaluate(provider)
        return TransformResult(name=self.name, count=1)

    def infer_output_model(self, spec: object) -> str:
        models = self.require_spec(spec).collect_models()
        if len(models) != 1:
            raise self.error_type(
                f"{self.name} must refer to exactly one model to infer output, got {sorted(models)}"
            )

        return next(iter(models))


def _unit_test_assert_compile_rejects_unknown_op() -> None:
    try:
        AssertTransform().compile({"unknown": {}}, default_model="model")
    except AssertTransformError as exc:
        assert "unknown assert op" in str(exc)
    else:  # pragma: no cover
        raise AssertionError("expected unknown assert op error")


def _unit_test_assert_compile_builds_spec() -> None:
    spec = AssertTransform().compile({"exists": "x"}, default_model="model")
    assert isinstance(spec, AssertSpec)
    assert spec.collect_models() == {"model"}


def _unit_test_assert_infer_output_model_requires_single_model() -> None:
    spec = AssertTransform().compile(
        {"equal": {"left": "a::x", "right": "b::x"}},
        default_model=None,
    )
    try:
        AssertTransform().infer_output_model(spec)
    except AssertTransformError as exc:
        assert "exactly one model" in str(exc)
    else:  # pragma: no cover
        raise AssertionError("expected single-model inference error")


__unit_tests__ = [
    _unit_test_assert_compile_rejects_unknown_op,
    _unit_test_assert_compile_builds_spec,
    _unit_test_assert_infer_output_model_requires_single_model,
]


register_transform(AssertTransform())
