"""Borrow array types and operations."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, cast

import hugr.model as model
from hugr import tys, val
from hugr.std._util import _load_extension
from hugr.utils import comma_sep_str

if TYPE_CHECKING:
    from hugr.ext import ExtensionRegistry, ExtensionResolutionResult

EXTENSION = _load_extension("collections.borrow_arr")


@dataclass(eq=False)
class BorrowArray(tys.ExtType):
    """Fixed `size` borrow array of `ty` elements."""

    def __init__(self, ty: tys.Type, size: int | tys.TypeArg) -> None:
        if isinstance(size, int):
            size = tys.BoundedNatArg(size)

        err_msg = (
            f"Borrow array size must be a bounded natural or a nat variable, not {size}"
        )
        match size:
            case tys.BoundedNatArg(_n):
                pass
            case tys.VariableArg(_idx, param):
                if not isinstance(param, tys.BoundedNatParam):
                    raise ValueError(err_msg)  # noqa: TRY004
            case _:
                raise ValueError(err_msg)

        ty_arg = tys.TypeTypeArg(ty)

        self.type_def = EXTENSION.types["borrow_array"]
        self.args = [size, ty_arg]

    @property
    def ty(self) -> tys.Type:
        assert isinstance(
            self.args[1], tys.TypeTypeArg
        ), "Borrow array elements must have a valid type"
        return self.args[1].ty

    @property
    def size(self) -> int | None:
        """If the borrow array has a concrete size, return it.

        Otherwise, return None.
        """
        if isinstance(self.args[0], tys.BoundedNatArg):
            return self.args[0].n
        return None

    def type_bound(self) -> tys.TypeBound:
        return tys.TypeBound.Linear

    def _resolve_used_extensions(
        self, registry: ExtensionRegistry | None = None
    ) -> tuple[BorrowArray, ExtensionResolutionResult]:
        ext_type, result = super()._resolve_used_extensions(registry)

        assert isinstance(
            ext_type, tys.ExtType
        ), "HUGR internal error, expected resolved type to be extension type."
        assert (
            ext_type.type_def == EXTENSION.types["borrow_array"]
        ), "HUGR internal error, expected resolved type to be borrow array."

        borrow_array = BorrowArray(tys.Unit, 0)
        borrow_array.args = ext_type.args
        return borrow_array, result


# Note that only borrow array values with no elements borrowed should be emitted.
@dataclass
class BorrowArrayVal(val.ExtensionValue):
    """Constant value for a statically sized borrow array of elements."""

    v: list[val.Value]
    ty: BorrowArray

    def __init__(self, v: list[val.Value], elem_ty: tys.Type) -> None:
        self.v = v
        self.ty = BorrowArray(elem_ty, len(v))

    def to_value(self) -> val.Extension:
        name = "BorrowArrayValue"
        # The value list must be serialized at this point, otherwise the
        # `Extension` value would not be serializable.
        vs = [v._to_serial_root() for v in self.v]
        element_ty = self.ty.ty._to_serial_root()
        serial_val = {"values": vs, "typ": element_ty}
        return val.Extension(name, typ=self.ty, val=serial_val)

    def __str__(self) -> str:
        return f"borrow_array({comma_sep_str(self.v)})"

    def to_model(self) -> model.Term:
        return model.Apply(
            "collections.borrow_array.const",
            [
                model.Literal(len(self.v)),
                cast(model.Term, self.ty.ty.to_model()),
                model.List([value.to_model() for value in self.v]),
            ],
        )

    def _resolve_used_extensions_inplace(
        self, registry: ExtensionRegistry | None = None
    ) -> ExtensionResolutionResult:
        resolved_ty, result = self.ty._resolve_used_extensions(registry)
        assert isinstance(
            resolved_ty, BorrowArray
        ), "HUGR internal error, expected resolved type to be borrow array."
        self.ty = resolved_ty
        for value in self.v:
            result.extend(value._resolve_used_extensions_inplace(registry))
        return result
