from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, List, Protocol

import torch

from .dtypes import parse_torch_dtype
from .transforms.unary import format_target_ref, resolve_target_names
from .transform import (
    StateDictProvider,
    TensorRef,
    TransformError,
    ensure_mapping_payload,
    must_model,
    parse_model_expr,
    parse_slice,
    select_tensor,
    validate_payload_keys,
)


class AssertTransformError(TransformError):
    pass


class AssertExpr(Protocol):
    def evaluate(self, provider: StateDictProvider) -> None: ...
    def collect_models(self) -> set[str]: ...


@dataclass(frozen=True)
class AssertExprHelp:
    name: str
    payload_kind: str
    allowed_keys: set[str] | None = None
    required_keys: set[str] | None = None
    description: str | None = None


AssertExprCompiler = Callable[[Any, str | None], AssertExpr]

_ASSERT_EXPR_COMPILERS: dict[str, AssertExprCompiler] = {}
_ASSERT_EXPR_HELP: dict[str, AssertExprHelp] = {}


def register_assert_expr(
    name: str,
    *,
    payload_kind: str,
    allowed_keys: set[str] | None = None,
    required_keys: set[str] | None = None,
    description: str | None = None,
) -> Callable[[AssertExprCompiler], AssertExprCompiler]:
    def decorator(fn: AssertExprCompiler) -> AssertExprCompiler:
        _ASSERT_EXPR_COMPILERS[name] = fn
        _ASSERT_EXPR_HELP[name] = AssertExprHelp(
            name=name,
            payload_kind=payload_kind,
            allowed_keys=None if allowed_keys is None else set(allowed_keys),
            required_keys=None if required_keys is None else set(required_keys),
            description=description,
        )
        return fn

    return decorator


def get_assert_expr_names() -> list[str]:
    return sorted(_ASSERT_EXPR_COMPILERS)


def get_assert_expr_help(name: str | None = None) -> dict[str, AssertExprHelp] | AssertExprHelp:
    if name is None:
        return dict(_ASSERT_EXPR_HELP)

    try:
        return _ASSERT_EXPR_HELP[name]
    except KeyError as exc:
        raise AssertTransformError(f"unknown assert op: {name!r}") from exc


@dataclass(frozen=True)
class ExistsExpr:
    ref: TensorRef

    def evaluate(self, provider: StateDictProvider) -> None:
        matches = resolve_matches(self.ref, provider, op_name="exists")
        if not matches:
            raise AssertTransformError(f"exists failed: {format_ref(self.ref)} matched zero tensors")

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


@dataclass(frozen=True)
class CountExpr:
    ref: TensorRef
    is_value: int

    def evaluate(self, provider: StateDictProvider) -> None:
        matches = resolve_matches(self.ref, provider, op_name="count.of")
        if len(matches) != self.is_value:
            raise AssertTransformError(
                f"count failed: {format_ref(self.ref)} matched {len(matches)} tensors, expected {self.is_value}"
            )

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


@dataclass(frozen=True)
class DtypeExpr:
    ref: TensorRef
    is_value: torch.dtype

    def evaluate(self, provider: StateDictProvider) -> None:
        tensor = resolve_single_tensor(self.ref, provider, op_name="dtype.of")
        if tensor.dtype != self.is_value:
            raise AssertTransformError(
                f"dtype failed: {format_ref(self.ref)} has dtype {tensor.dtype}, expected {self.is_value}"
            )

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


@dataclass(frozen=True)
class ShapeExpr:
    ref: TensorRef
    is_value: tuple[int, ...]

    def evaluate(self, provider: StateDictProvider) -> None:
        tensor = resolve_single_tensor(self.ref, provider, op_name="shape.of")
        if tuple(tensor.shape) != self.is_value:
            raise AssertTransformError(
                f"shape failed: {format_ref(self.ref)} has shape {tuple(tensor.shape)}, expected {self.is_value}"
            )

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


@dataclass(frozen=True)
class DimensionsExpr:
    ref: TensorRef
    is_value: int

    def evaluate(self, provider: StateDictProvider) -> None:
        tensor = resolve_single_tensor(self.ref, provider, op_name="dimensions.of")
        if len(tensor.shape) != self.is_value:
            raise AssertTransformError(
                f"dimensions failed: {format_ref(self.ref)} has {len(tensor.shape)} dims, expected {self.is_value}"
            )

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


@dataclass(frozen=True)
class IsZeroExpr:
    ref: TensorRef

    def evaluate(self, provider: StateDictProvider) -> None:
        tensor = resolve_single_tensor(self.ref, provider, op_name="iszero")
        if not torch.all(tensor == 0):
            raise AssertTransformError(f"iszero failed: {format_ref(self.ref)} is not all zeros")

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


@dataclass(frozen=True)
class EqualExpr:
    left: TensorRef
    right: TensorRef

    def evaluate(self, provider: StateDictProvider) -> None:
        left = resolve_single_tensor(self.left, provider, op_name="equal.left")
        right = resolve_single_tensor(self.right, provider, op_name="equal.right")

        if left.shape != right.shape:
            raise AssertTransformError(
                f"equal failed: shape mismatch {tuple(left.shape)} != {tuple(right.shape)}"
            )
        if left.dtype != right.dtype:
            raise AssertTransformError(f"equal failed: dtype mismatch {left.dtype} != {right.dtype}")
        if not torch.equal(left, right):
            raise AssertTransformError(f"equal failed: {format_ref(self.left)} != {format_ref(self.right)}")

    def collect_models(self) -> set[str]:
        return {must_model(self.left), must_model(self.right)}


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

    def evaluate(self, provider: StateDictProvider) -> None:
        try:
            self.expr.evaluate(provider)
        except AssertTransformError:
            return

        raise AssertTransformError(f"not failed: inner assertion succeeded: {format_expr(self.expr)}")

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


@dataclass(frozen=True)
class AllExpr:
    exprs: List[AssertExpr]

    def evaluate(self, provider: StateDictProvider) -> None:
        for expr in self.exprs:
            expr.evaluate(provider)

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


@dataclass(frozen=True)
class AnyExpr:
    exprs: List[AssertExpr]

    def evaluate(self, provider: StateDictProvider) -> None:
        errors: List[str] = []
        for expr in self.exprs:
            try:
                expr.evaluate(provider)
                return
            except AssertTransformError as exc:
                errors.append(str(exc))
        raise AssertTransformError("any failed: all alternatives failed:\n- " + "\n- ".join(errors))

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


def compile_assert_expr(raw: Any, default_model: str | None) -> AssertExpr:
    if not isinstance(raw, dict) or len(raw) != 1:
        raise AssertTransformError("assert expression must be a single-key mapping")

    op, payload = next(iter(raw.items()))

    try:
        compiler = _ASSERT_EXPR_COMPILERS[op]
    except KeyError as exc:
        raise AssertTransformError(f"unknown assert op: {op!r}") from exc

    return compiler(payload, default_model)


@register_assert_expr(
    "exists",
    payload_kind="tensor-ref",
    description="Succeeds if the reference matches at least one tensor.",
)
def compile_exists_expr(payload: Any, default_model: str | None) -> AssertExpr:
    return ExistsExpr(ref=compile_tensor_ref_expr(payload, default_model, "exists"))


@register_assert_expr(
    "count",
    payload_kind="mapping",
    allowed_keys={"of", "is"},
    required_keys={"of", "is"},
    description="Succeeds if the reference matches exactly the given number of tensors.",
)
def compile_count_expr(payload: Any, default_model: str | None) -> CountExpr:
    payload = ensure_mapping_payload(payload, "count")
    validate_payload_keys(
        payload,
        op_name="count",
        allowed_keys={"of", "is"},
        required_keys={"of", "is"},
    )
    ref = compile_tensor_ref_expr(payload["of"], default_model, "count.of")
    is_value = payload["is"]
    if not isinstance(is_value, int):
        raise AssertTransformError("count.is must be an integer")
    return CountExpr(ref=ref, is_value=is_value)


@register_assert_expr(
    "dtype",
    payload_kind="mapping",
    allowed_keys={"of", "is"},
    required_keys={"of", "is"},
    description="Succeeds if the tensor has the given dtype.",
)
def compile_dtype_expr(payload: Any, default_model: str | None) -> DtypeExpr:
    payload = ensure_mapping_payload(payload, "dtype")
    validate_payload_keys(
        payload,
        op_name="dtype",
        allowed_keys={"of", "is"},
        required_keys={"of", "is"},
    )
    ref = compile_tensor_ref_expr(payload["of"], default_model, "dtype.of")

    raw_dtype = payload["is"]
    if not isinstance(raw_dtype, str) or not raw_dtype:
        raise AssertTransformError("dtype.is must be a non-empty string")

    return DtypeExpr(
        ref=ref,
        is_value=parse_torch_dtype(
            raw_dtype,
            error_type=AssertTransformError,
            op_name="dtype",
            field_name="is",
        ),
    )


@register_assert_expr(
    "shape",
    payload_kind="mapping",
    allowed_keys={"of", "is"},
    required_keys={"of", "is"},
    description="Succeeds if the tensor has the given shape.",
)
def compile_shape_expr(payload: Any, default_model: str | None) -> ShapeExpr:
    payload = ensure_mapping_payload(payload, "shape")
    validate_payload_keys(
        payload,
        op_name="shape",
        allowed_keys={"of", "is"},
        required_keys={"of", "is"},
    )
    ref = compile_tensor_ref_expr(payload["of"], default_model, "shape.of")
    return ShapeExpr(ref=ref, is_value=compile_shape(payload["is"]))


@register_assert_expr(
    "dimensions",
    payload_kind="mapping",
    allowed_keys={"of", "is"},
    required_keys={"of", "is"},
    description="Succeeds if the tensor has the given number of dimensions.",
)
def compile_dimensions_expr(payload: Any, default_model: str | None) -> DimensionsExpr:
    payload = ensure_mapping_payload(payload, "dimensions")
    validate_payload_keys(
        payload,
        op_name="dimensions",
        allowed_keys={"of", "is"},
        required_keys={"of", "is"},
    )
    ref = compile_tensor_ref_expr(payload["of"], default_model, "dimensions.of")
    is_value = payload["is"]
    if not isinstance(is_value, int):
        raise AssertTransformError("dimensions.is must be an integer")
    return DimensionsExpr(ref=ref, is_value=is_value)


@register_assert_expr(
    "iszero",
    payload_kind="tensor-ref",
    description="Succeeds if the selected tensor is all zeros.",
)
def compile_iszero_expr(payload: Any, default_model: str | None) -> AssertExpr:
    return IsZeroExpr(ref=compile_tensor_ref_expr(payload, default_model, "iszero"))


@register_assert_expr(
    "equal",
    payload_kind="mapping",
    allowed_keys={"left", "right"},
    required_keys={"left", "right"},
    description="Succeeds if two tensors have the same shape, dtype, and values.",
)
def compile_equal_expr(payload: Any, default_model: str | None) -> EqualExpr:
    payload = ensure_mapping_payload(payload, "equal")
    validate_payload_keys(
        payload,
        op_name="equal",
        allowed_keys={"left", "right"},
        required_keys={"left", "right"},
    )
    left = compile_tensor_ref_expr(payload["left"], default_model, "equal.left")
    right = compile_tensor_ref_expr(payload["right"], default_model, "equal.right")
    return EqualExpr(left=left, right=right)


@register_assert_expr(
    "not",
    payload_kind="assert-expr",
    description="Succeeds if the inner assertion fails.",
)
def compile_not_expr(payload: Any, default_model: str | None) -> AssertExpr:
    return NotExpr(expr=compile_assert_expr(payload, default_model))


@register_assert_expr(
    "all",
    payload_kind="list[assert-expr]",
    description="Succeeds if all inner assertions succeed.",
)
def compile_all_expr(payload: Any, default_model: str | None) -> AllExpr:
    if not isinstance(payload, list) or not payload:
        raise AssertTransformError("all must be a non-empty list")
    return AllExpr(exprs=[compile_assert_expr(item, default_model) for item in payload])


@register_assert_expr(
    "any",
    payload_kind="list[assert-expr]",
    description="Succeeds if any inner assertion succeeds.",
)
def compile_any_expr(payload: Any, default_model: str | None) -> AnyExpr:
    if not isinstance(payload, list) or not payload:
        raise AssertTransformError("any must be a non-empty list")
    return AnyExpr(exprs=[compile_assert_expr(item, default_model) for item in payload])


def compile_tensor_ref_expr(raw: Any, default_model: str | None, op_name: str) -> TensorRef:
    if isinstance(raw, str):
        if not raw:
            raise AssertTransformError(
                f"{op_name} must be a non-empty string reference or non-empty list of strings"
            )
    elif isinstance(raw, list):
        if not raw or not all(isinstance(item, str) and item for item in raw):
            raise AssertTransformError(
                f"{op_name} must be a non-empty string reference or non-empty list of strings"
            )
    else:
        raise AssertTransformError(
            f"{op_name} must be a non-empty string reference or non-empty list of strings"
        )

    ref = parse_model_expr(raw, default_model=default_model)
    if ref.slice_spec is not None:
        parse_slice(ref.slice_spec)
    return ref


def compile_shape(raw: Any) -> tuple[int, ...]:
    if not isinstance(raw, list) or not all(isinstance(x, int) for x in raw):
        raise AssertTransformError("shape.is must be a list of integers")
    return tuple(raw)


def resolve_matches(
    ref: TensorRef,
    provider: StateDictProvider,
    *,
    op_name: str,
) -> List[str]:
    return resolve_target_names(
        target_ref=ref,
        provider=provider,
        op_name=op_name,
        error_type=AssertTransformError,
    )


def resolve_single_tensor(ref: TensorRef, provider: StateDictProvider, op_name: str) -> torch.Tensor:
    model = must_model(ref)
    sd = provider.get_state_dict(model)
    matches = resolve_matches(ref, provider, op_name=op_name)

    if len(matches) == 0:
        raise AssertTransformError(f"{op_name} failed: {format_ref(ref)} matched zero tensors")
    if len(matches) != 1:
        raise AssertTransformError(
            f"{op_name} failed: {format_ref(ref)} matched {len(matches)} tensors, expected 1"
        )

    tensor = sd[matches[0]]
    slice_spec = parse_slice(ref.slice_spec) if ref.slice_spec else None
    return select_tensor(tensor, slice_spec)


def collect_ref_models(ref: TensorRef) -> set[str]:
    return {must_model(ref)}


def collect_expr_models(exprs: list[AssertExpr]) -> set[str]:
    models: set[str] = set()
    for expr in exprs:
        models.update(expr.collect_models())
    return models


def format_ref(ref: TensorRef) -> str:
    return format_target_ref(ref)


def format_expr(expr: AssertExpr) -> str:
    if isinstance(expr, ExistsExpr):
        return f"exists({format_ref(expr.ref)})"
    if isinstance(expr, CountExpr):
        return f"count({format_ref(expr.ref)}) == {expr.is_value}"
    if isinstance(expr, DtypeExpr):
        return f"dtype({format_ref(expr.ref)}) == {expr.is_value}"
    if isinstance(expr, ShapeExpr):
        return f"shape({format_ref(expr.ref)}) == {expr.is_value}"
    if isinstance(expr, DimensionsExpr):
        return f"dimensions({format_ref(expr.ref)}) == {expr.is_value}"
    if isinstance(expr, IsZeroExpr):
        return f"iszero({format_ref(expr.ref)})"
    if isinstance(expr, EqualExpr):
        return f"equal({format_ref(expr.left)}, {format_ref(expr.right)})"
    if isinstance(expr, NotExpr):
        return f"not({format_expr(expr.expr)})"
    if isinstance(expr, AllExpr):
        return f"all({', '.join(format_expr(e) for e in expr.exprs)})"
    if isinstance(expr, AnyExpr):
        return f"any({', '.join(format_expr(e) for e in expr.exprs)})"
    return repr(expr)
