"""Compilers building list functions on top of hugr standard operations, that involve
multiple nodes.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import hugr.std.collections.list
from hugr import Wire, ops
from hugr import tys as ht
from hugr.std.collections.list import List, ListVal

from guppylang.definition.custom import CustomCallCompiler
from guppylang.definition.value import CallReturnWires
from guppylang.error import InternalGuppyError
from guppylang.std._internal.compiler.arithmetic import (
    convert_ifromusize,
    convert_itousize,
)
from guppylang.std._internal.compiler.prelude import (
    build_unwrap,
    build_unwrap_left,
    build_unwrap_right,
)
from guppylang.tys.arg import TypeArg

if TYPE_CHECKING:
    from hugr.build.dfg import DfBase


# ------------------------------------------------------
# --------- std.collections operations -----------------
# ------------------------------------------------------


def _instantiate_list_op(
    name: str, elem_type: ht.Type, inp: list[ht.Type], out: list[ht.Type]
) -> ops.ExtOp:
    op_def = hugr.std.collections.list.EXTENSION.get_op(name)
    return ops.ExtOp(
        op_def,
        ht.FunctionType(inp, out),
        [ht.TypeTypeArg(elem_type)],
    )


def list_pop(elem_type: ht.Type) -> ops.ExtOp:
    """Returns a list `pop` operation."""
    list_type = List(elem_type)
    return _instantiate_list_op(
        "pop", elem_type, [list_type], [list_type, ht.Option(elem_type)]
    )


def list_push(elem_type: ht.Type) -> ops.ExtOp:
    """Returns a list `push` operation."""
    list_type = List(elem_type)
    return _instantiate_list_op("push", elem_type, [list_type, elem_type], [list_type])


def list_get(elem_type: ht.Type) -> ops.ExtOp:
    """Returns a list `get` operation."""
    list_type = List(elem_type)
    return _instantiate_list_op(
        "get", elem_type, [list_type, ht.USize()], [ht.Option(elem_type)]
    )


def list_set(elem_type: ht.Type) -> ops.ExtOp:
    """Returns a list `set` operation."""
    list_type = List(elem_type)
    return _instantiate_list_op(
        "set",
        elem_type,
        [list_type, ht.USize(), elem_type],
        [list_type, ht.Either([elem_type], [elem_type])],
    )


def list_insert(elem_type: ht.Type) -> ops.ExtOp:
    """Returns a list `insert` operation."""
    list_type = List(elem_type)
    return _instantiate_list_op(
        "insert",
        elem_type,
        [list_type, ht.USize(), elem_type],
        [list_type, ht.Either([elem_type], [ht.Unit])],
    )


def list_length(elem_type: ht.Type) -> ops.ExtOp:
    """Returns a list `length` operation."""
    list_type = List(elem_type)
    return _instantiate_list_op(
        "length", elem_type, [list_type], [list_type, ht.USize()]
    )


# ------------------------------------------------------
# --------- Custom compilers for non-native ops --------
# ------------------------------------------------------


class ListGetitemCompiler(CustomCallCompiler):
    """Compiler for the `list.__getitem__` function."""

    def build_classical_getitem(
        self,
        list_wire: Wire,
        idx: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.__getitem__` for classical lists."""
        idx = self.builder.add_op(convert_itousize(), idx)
        result = self.builder.add_op(list_get(elem_ty), list_wire, idx)
        elem = build_unwrap(self.builder, result, "List index out of bounds")
        return CallReturnWires(regular_returns=[elem], inout_returns=[list_wire])

    def build_linear_getitem(
        self,
        list_wire: Wire,
        idx: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.__getitem__` for linear lists."""
        # Swap out the element at the given index with `None`. The `to_hugr`
        # implementation of the list type ensures that linear element types are turned
        # into optionals.
        elem_opt_ty = ht.Option(elem_ty)
        none = self.builder.add_op(ops.Tag(0, elem_opt_ty))
        idx = self.builder.add_op(convert_itousize(), idx)
        list_wire, result = self.builder.add_op(
            list_set(elem_opt_ty), list_wire, idx, none
        )
        elem_opt = build_unwrap_right(self.builder, result, "List index out of bounds")
        elem = build_unwrap(
            self.builder, elem_opt, "Linear list element has already been used"
        )
        return CallReturnWires(regular_returns=[elem], inout_returns=[list_wire])

    def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires:
        [list_wire, idx] = args
        [elem_ty_arg] = self.type_args
        assert isinstance(elem_ty_arg, TypeArg)
        if elem_ty_arg.ty.linear:
            return self.build_linear_getitem(list_wire, idx, elem_ty_arg.ty.to_hugr())
        else:
            return self.build_classical_getitem(
                list_wire, idx, elem_ty_arg.ty.to_hugr()
            )

    def compile(self, args: list[Wire]) -> list[Wire]:
        raise InternalGuppyError("Call compile_with_inouts instead")


class ListSetitemCompiler(CustomCallCompiler):
    """Compiler for the `list.__setitem__` function."""

    def build_classical_setitem(
        self,
        list_wire: Wire,
        idx: Wire,
        elem: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.__setitem__` for classical lists."""
        idx = self.builder.add_op(convert_itousize(), idx)
        list_wire, result = self.builder.add_op(list_set(elem_ty), list_wire, idx, elem)
        # Unwrap the result, but we don't have to hold onto the returned old value
        build_unwrap_right(self.builder, result, "List index out of bounds")
        return CallReturnWires(regular_returns=[], inout_returns=[list_wire])

    def build_linear_setitem(
        self,
        list_wire: Wire,
        idx: Wire,
        elem: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `array.__setitem__` for linear arrays."""
        # Embed the element into an optional
        elem_opt_ty = ht.Option(elem_ty)
        elem = self.builder.add_op(ops.Some(elem_ty), elem)
        idx = self.builder.add_op(convert_itousize(), idx)
        list_wire, result = self.builder.add_op(
            list_set(elem_opt_ty), list_wire, idx, elem
        )
        old_elem_opt = build_unwrap_right(
            self.builder, result, "List index out of bounds"
        )
        # Check that the old element was `None`
        build_unwrap_left(
            self.builder, old_elem_opt, "Linear list element has not been used"
        )
        return CallReturnWires(regular_returns=[], inout_returns=[list_wire])

    def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires:
        [list_wire, idx, elem] = args
        [elem_ty_arg] = self.type_args
        assert isinstance(elem_ty_arg, TypeArg)
        if elem_ty_arg.ty.linear:
            return self.build_linear_setitem(
                list_wire, idx, elem, elem_ty_arg.ty.to_hugr()
            )
        else:
            return self.build_classical_setitem(
                list_wire, idx, elem, elem_ty_arg.ty.to_hugr()
            )

    def compile(self, args: list[Wire]) -> list[Wire]:
        raise InternalGuppyError("Call compile_with_inouts instead")


class ListPopCompiler(CustomCallCompiler):
    """Compiler for the `list.pop` function."""

    def build_classical_pop(
        self,
        list_wire: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.pop` for classical lists."""
        list_wire, result = self.builder.add_op(list_pop(elem_ty), list_wire)
        elem = build_unwrap(self.builder, result, "List index out of bounds")
        return CallReturnWires(regular_returns=[elem], inout_returns=[list_wire])

    def build_linear_pop(
        self,
        list_wire: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.pop` for linear lists."""
        elem_opt_ty = ht.Option(elem_ty)
        list_wire, result = self.builder.add_op(list_pop(elem_opt_ty), list_wire)
        elem_opt = build_unwrap(self.builder, result, "List index out of bounds")
        elem = build_unwrap(
            self.builder, elem_opt, "Linear list element has already been used"
        )
        return CallReturnWires(regular_returns=[elem], inout_returns=[list_wire])

    def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires:
        [list_wire] = args
        [elem_ty_arg] = self.type_args
        assert isinstance(elem_ty_arg, TypeArg)
        if elem_ty_arg.ty.linear:
            return self.build_linear_pop(list_wire, elem_ty_arg.ty.to_hugr())
        else:
            return self.build_classical_pop(list_wire, elem_ty_arg.ty.to_hugr())

    def compile(self, args: list[Wire]) -> list[Wire]:
        raise InternalGuppyError("Call compile_with_inouts instead")


class ListPushCompiler(CustomCallCompiler):
    """Compiler for the `list.push` function."""

    def build_classical_push(
        self,
        list_wire: Wire,
        elem: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.push` for classical lists."""
        list_wire = self.builder.add_op(list_push(elem_ty), list_wire, elem)
        return CallReturnWires(regular_returns=[], inout_returns=[list_wire])

    def build_linear_push(
        self,
        list_wire: Wire,
        elem: Wire,
        elem_ty: ht.Type,
    ) -> CallReturnWires:
        """Lowers a call to `list.push` for linear lists."""
        # Wrap element into an optional
        elem_opt_ty = ht.Option(elem_ty)
        elem_opt = self.builder.add_op(ops.Some(elem_ty), elem)
        list_wire = self.builder.add_op(list_push(elem_opt_ty), list_wire, elem_opt)
        return CallReturnWires(regular_returns=[], inout_returns=[list_wire])

    def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires:
        [list_wire, elem] = args
        [elem_ty_arg] = self.type_args
        assert isinstance(elem_ty_arg, TypeArg)
        if elem_ty_arg.ty.linear:
            return self.build_linear_push(list_wire, elem, elem_ty_arg.ty.to_hugr())
        else:
            return self.build_classical_push(list_wire, elem, elem_ty_arg.ty.to_hugr())

    def compile(self, args: list[Wire]) -> list[Wire]:
        raise InternalGuppyError("Call compile_with_inouts instead")


class ListLengthCompiler(CustomCallCompiler):
    """Compiler for the `list.__len__` function."""

    def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires:
        [list_wire] = args
        [elem_ty_arg] = self.type_args
        assert isinstance(elem_ty_arg, TypeArg)
        elem_ty = elem_ty_arg.ty.to_hugr()
        if elem_ty_arg.ty.linear:
            elem_ty = ht.Option(elem_ty)
        list_wire, length = self.builder.add_op(list_length(elem_ty), list_wire)
        length = self.builder.add_op(convert_ifromusize(), length)
        return CallReturnWires(regular_returns=[length], inout_returns=[list_wire])

    def compile(self, args: list[Wire]) -> list[Wire]:
        raise InternalGuppyError("Call compile_with_inouts instead")


def list_new(
    builder: DfBase[ops.DfParentOp], elem_type: ht.Type, args: list[Wire]
) -> Wire:
    if elem_type.type_bound() == ht.TypeBound.Any:
        return _list_new_linear(builder, elem_type, args)
    else:
        return _list_new_classical(builder, elem_type, args)


def _list_new_classical(
    builder: DfBase[ops.DfParentOp], elem_type: ht.Type, args: list[Wire]
) -> Wire:
    # This may be simplified in the future with a `new` or `with_capacity` list op
    # See https://github.com/CQCL/hugr/issues/1508
    lst = builder.load(ListVal([], elem_ty=elem_type))
    push_op = list_push(elem_type)
    for elem in args:
        lst = builder.add_op(push_op, lst, elem)
    return lst


def _list_new_linear(
    builder: DfBase[ops.DfParentOp], elem_type: ht.Type, args: list[Wire]
) -> Wire:
    elem_opt_ty = ht.Option(elem_type)
    lst = builder.load(ListVal([], elem_ty=elem_opt_ty))
    push_op = list_push(elem_opt_ty)
    for elem in args:
        elem_opt = builder.add_op(ops.Some(elem_type), elem)
        lst = builder.add_op(push_op, lst, elem_opt)
    return lst
