# ruff: noqa: I001

import hugr
from hugr.package import Package
import hugr.ops as ops
import hugr.tys as tys
from hugr import ext, val

from hugr.build import Dfg
from hugr.std.collections.list import List
from hugr.std.float import FLOAT_TYPES_EXTENSION
from hugr.std.int import CONVERSIONS_EXTENSION, INT_T, INT_TYPES_EXTENSION
import pytest

from .conftest import H, QUANTUM_EXT, TEST_EXT, TEST_TYPE_OPAQUE, TEST_OP_OPAQUE


def test_extension_ops() -> None:
    h = Dfg(tys.Qubit)
    (q,) = h.inputs()

    h_node = h.add(H(q))
    h.set_outputs(h_node.out(0))

    op_exts = h.hugr[h_node].op.used_extensions().ids()
    assert set(op_exts) == {QUANTUM_EXT.name, FLOAT_TYPES_EXTENSION.name, "prelude"}

    exts = h.hugr.used_extensions().ids()
    assert set(exts) == {QUANTUM_EXT.name, FLOAT_TYPES_EXTENSION.name, "prelude"}


def test_core_ops() -> None:
    h = Dfg(INT_T, tys.Bool)
    a, b = h.inputs()

    tuple_n = h.add(ops.MakeTuple()(a, b))
    h.set_outputs(tuple_n.out(0))

    op_exts = h.hugr[tuple_n].op.used_extensions().ids()
    assert set(op_exts) == {INT_TYPES_EXTENSION.name, "prelude"}

    exts = h.hugr.used_extensions().ids()
    assert set(exts) == {INT_TYPES_EXTENSION.name, "prelude"}


def test_transitive_extension_dependencies() -> None:
    op = ops.ExtOp(CONVERSIONS_EXTENSION.get_op("itousize"), args=[])
    exts = op.used_extensions().ids()
    assert CONVERSIONS_EXTENSION.name in exts
    assert FLOAT_TYPES_EXTENSION.name in exts


def test_func_calls() -> None:
    h = Dfg(tys.Bool)
    [a] = h.inputs()

    f_decl = h.module_root_builder().declare_function(
        "f",
        tys.PolyFuncType(
            params=[tys.TypeTypeParam(bound=tys.TypeBound.Copyable)],
            body=tys.FunctionType([tys.Qubit], []),
        ),
    )

    call = h.call(
        f_decl,
        a,
        instantiation=tys.FunctionType([tys.Qubit], []),
        type_args=[tys.TypeTypeArg(INT_T)],
    )
    h.set_outputs()

    op_exts = h.hugr[call].op.used_extensions().ids()
    assert set(op_exts) == {INT_TYPES_EXTENSION.name, "prelude"}

    exts = h.hugr.used_extensions().ids()
    assert set(exts) == {INT_TYPES_EXTENSION.name, "prelude"}


def test_type_arguments() -> None:
    list_ty = List(tys.Qubit)
    h = Dfg(list_ty)
    (items,) = h.inputs()
    h.set_outputs(items)

    exts = h.hugr.used_extensions().ids()
    assert set(exts) == {hugr.std.collections.list.EXTENSION.name, "prelude"}


def test_op_signature_contains_same_extension() -> None:
    """Test that used_extensions() works when an op's signature contains
    types from the same extension as the operation itself.

    Regression test for https://github.com/Quantinuum/hugr/issues/2829
    """

    # Create an extension with a custom type
    my_ext = ext.Extension("test.self_referential", ext.Version(0, 1, 0))
    my_type_def = my_ext.add_type_def(
        ext.TypeDef(
            name="MyType",
            description="A custom type",
            params=[],
            bound=ext.ExplicitBound(tys.TypeBound.Copyable),
        )
    )

    # Create an ExtType from this extension's type definition
    my_type = tys.ExtType(my_type_def, args=[])

    # Add an operation to the same extension that uses MyType in its signature
    my_ext.add_op_def(
        ext.OpDef(
            name="ProcessMyType",
            description="An operation that uses MyType",
            signature=ext.OpDefSig(tys.FunctionType([my_type], [my_type])),
        )
    )

    # Create an ExtOp from this definition
    op = ops.ExtOp(my_ext.get_op("ProcessMyType"), args=[])

    # This would have failed with ExtensionRegistry.ExtensionExists before the fix
    # because the signature contains my_type (from my_ext) and we also add my_ext
    exts = op.used_extensions()
    assert my_ext.name in exts.ids()


@pytest.mark.parametrize(
    ("typ", "extensions"),
    [
        # Opaque type
        (TEST_TYPE_OPAQUE, [TEST_EXT]),
        (tys.Sum([[TEST_TYPE_OPAQUE], [tys.Bool]]), [TEST_EXT]),
        (tys.Tuple(TEST_TYPE_OPAQUE, tys.Bool), [TEST_EXT]),
        (tys.Option(TEST_TYPE_OPAQUE), [TEST_EXT]),
        (tys.Either([TEST_TYPE_OPAQUE], [tys.Bool]), [TEST_EXT]),
        (tys.UnitSum(3), []),
        (tys.FunctionType([TEST_TYPE_OPAQUE], [tys.Bool]), [TEST_EXT]),
        (
            tys.PolyFuncType(
                [tys.FloatParam()], tys.FunctionType([TEST_TYPE_OPAQUE], [tys.Bool])
            ),
            [TEST_EXT],
        ),
        (tys.Variable(0, tys.TypeBound.Copyable), []),
        (tys.RowVariable(0, tys.TypeBound.Copyable), []),
        (tys.Alias("MyAlias", tys.TypeBound.Copyable), []),
    ],
)
def test_type_resolution(typ: tys.Type, extensions: list[ext.Extension]) -> None:
    """Test that Opaque types are tracked and resolved correctly."""
    _, result = typ._resolve_used_extensions()
    for extension in extensions:
        assert extension.name in result.unresolved_extensions
    if extensions:
        assert (TEST_EXT.name, "TestType") in result.unresolved_types
    assert not result.unresolved_ops

    test_registry = ext.ExtensionRegistry()
    for extension in extensions:
        test_registry.add_extension(extension)
    _, result = typ._resolve_used_extensions(test_registry)
    for extension in extensions:
        assert extension.name in result.used_extensions.ids()
    assert not result.unresolved_types
    assert not result.unresolved_ops


@pytest.mark.parametrize(
    ("arg", "extensions"),
    [
        (tys.TypeTypeArg(TEST_TYPE_OPAQUE), [TEST_EXT]),
        (tys.ListArg([tys.TypeTypeArg(TEST_TYPE_OPAQUE)]), [TEST_EXT]),
        (
            tys.TupleArg([tys.TypeTypeArg(TEST_TYPE_OPAQUE), tys.BoundedNatArg(5)]),
            [TEST_EXT],
        ),
        (
            tys.ListConcatArg([tys.ListArg([tys.TypeTypeArg(TEST_TYPE_OPAQUE)])]),
            [TEST_EXT],
        ),
        (
            tys.TupleConcatArg([tys.TupleArg([tys.TypeTypeArg(TEST_TYPE_OPAQUE)])]),
            [TEST_EXT],
        ),
        (tys.VariableArg(0, tys.ConstParam(TEST_TYPE_OPAQUE)), [TEST_EXT]),
        (tys.BoundedNatArg(42), []),
        (tys.StringArg("hello"), []),
        (tys.FloatArg(3.14), []),
        (tys.BytesArg(b"bytes"), []),
    ],
)
def test_type_arg_resolution(arg: tys.TypeArg, extensions: list[ext.Extension]) -> None:
    """Test that type arguments are tracked and resolved correctly."""
    _, result = arg._resolve_used_extensions()
    for extension in extensions:
        assert extension.name in result.unresolved_extensions
    if extensions:
        assert (TEST_EXT.name, "TestType") in result.unresolved_types
    assert not result.unresolved_ops

    test_registry = ext.ExtensionRegistry()
    for extension in extensions:
        test_registry.add_extension(extension)
    _, result = arg._resolve_used_extensions(test_registry)
    for extension in extensions:
        assert extension.name in result.used_extensions.ids()
    assert not result.unresolved_types
    assert not result.unresolved_ops


@pytest.mark.parametrize(
    ("param", "extensions"),
    [
        (tys.ListParam(tys.ConstParam(TEST_TYPE_OPAQUE)), [TEST_EXT]),
        (
            tys.TupleParam([tys.ConstParam(TEST_TYPE_OPAQUE), tys.BoundedNatParam()]),
            [TEST_EXT],
        ),
        (tys.ConstParam(TEST_TYPE_OPAQUE), [TEST_EXT]),
        (tys.TypeTypeParam(tys.TypeBound.Copyable), []),
        (tys.BoundedNatParam(100), []),
        (tys.StringParam(), []),
        (tys.FloatParam(), []),
        (tys.BytesParam(), []),
    ],
)
def test_type_param_resolution(
    param: tys.TypeParam, extensions: list[ext.Extension]
) -> None:
    """Test that type parameters are tracked and resolved correctly."""
    _, result = param._resolve_used_extensions()
    for extension in extensions:
        assert extension.name in result.unresolved_extensions
    if extensions:
        assert (TEST_EXT.name, "TestType") in result.unresolved_types
    assert not result.unresolved_ops

    test_registry = ext.ExtensionRegistry()
    for extension in extensions:
        test_registry.add_extension(extension)
    _, result = param._resolve_used_extensions(test_registry)
    for extension in extensions:
        assert extension.name in result.used_extensions.ids()
    assert not result.unresolved_types
    assert not result.unresolved_ops


@pytest.mark.parametrize(
    ("op", "extensions"),
    [
        (TEST_OP_OPAQUE, [TEST_EXT]),
        (ops.MakeTuple([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.UnpackTuple([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.Tag(0, tys.Sum([[TEST_TYPE_OPAQUE], []])), [TEST_EXT]),
        (ops.LoadConst(TEST_TYPE_OPAQUE), [TEST_EXT]),
        (ops.Conditional(tys.Sum([[TEST_TYPE_OPAQUE], []]), [tys.Bool]), [TEST_EXT]),
        (ops.Case([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.TailLoop([TEST_TYPE_OPAQUE], [tys.Bool]), [TEST_EXT]),
        (ops.DFG([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.CFG([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.DataflowBlock([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.ExitBlock([TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (ops.FuncDefn("test_fn", [TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (
            ops.FuncDecl(
                "test_fn",
                tys.PolyFuncType([], tys.FunctionType([TEST_TYPE_OPAQUE], [])),
            ),
            [TEST_EXT],
        ),
        (ops.CallIndirect(tys.FunctionType([TEST_TYPE_OPAQUE], [])), [TEST_EXT]),
        (ops.AliasDefn("MyAlias", TEST_TYPE_OPAQUE), [TEST_EXT]),
    ],
)
def test_op_resolution(op: ops.Op, extensions: list[ext.Extension]) -> None:
    """Test that ops are tracked and resolved correctly."""
    _, result = op._resolve_used_extensions()
    for extension in extensions:
        assert extension.name in result.unresolved_extensions

    if isinstance(op, ops.Custom):
        assert (op.extension, op.op_name) in result.unresolved_ops
        assert not result.unresolved_types
    elif extensions:
        assert not result.unresolved_ops
        assert (TEST_EXT.name, "TestType") in result.unresolved_types
    else:
        assert not result.unresolved_ops
        assert not result.unresolved_types

    test_ext_registry = ext.ExtensionRegistry()
    for extension in extensions:
        test_ext_registry.add_extension(extension)
    _, result = op._resolve_used_extensions(test_ext_registry)
    for extension in extensions:
        assert extension.name in result.used_extensions.ids()
    assert not result.unresolved_ops
    assert not result.unresolved_types


@pytest.mark.parametrize(
    ("value", "extensions"),
    [
        (val.Sum(0, tys.Sum([[TEST_TYPE_OPAQUE], []]), []), [TEST_EXT]),
        (val.Tuple(val.TRUE), []),
        (val.Some(val.TRUE), []),
        (val.None_(TEST_TYPE_OPAQUE), [TEST_EXT]),
        (val.Left([], [TEST_TYPE_OPAQUE]), [TEST_EXT]),
        (val.Right([TEST_TYPE_OPAQUE], []), [TEST_EXT]),
        (val.UnitSum(0, 3), []),
        (val.Extension("test", TEST_TYPE_OPAQUE, {}), [TEST_EXT]),
    ],
)
def test_value_resolution(value: val.Value, extensions: list[ext.Extension]) -> None:
    """Test that values are tracked and resolved correctly."""
    result = value._resolve_used_extensions_inplace()
    for extension in extensions:
        assert extension.name in result.unresolved_extensions

    # Reset the value for the second test by creating a fresh copy
    test_ext_registry = ext.ExtensionRegistry()
    for extension in extensions:
        test_ext_registry.add_extension(extension)
    result = value._resolve_used_extensions_inplace(test_ext_registry)
    for extension in extensions:
        assert extension.name in result.used_extensions.ids()


def test_hugr_with_opaque_type_resolve() -> None:
    """Test that Hugr.used_extensions can resolve opaque types."""
    # Create a DFG with the opaque type
    h = Dfg(TEST_TYPE_OPAQUE)
    (inp,) = h.inputs()
    h.set_outputs(inp)

    # Without resolve_from, unresolved extensions are tracked
    result_no_resolve = h.hugr.used_extensions()
    assert TEST_EXT.name in result_no_resolve.unresolved_extensions

    # With resolve_from, should work
    test_ext_registry = ext.ExtensionRegistry()
    test_ext_registry.add_extension(TEST_EXT)
    exts = h.hugr.used_extensions(resolve_from=test_ext_registry)
    assert TEST_EXT.name in exts.used_extensions.ids()


def test_lower_funcs_resolve() -> None:
    """Test that extensions in `OpDef.lower_func`s are resolved correctly."""
    # A custom extension used by inside lower function.
    inner_op_def = ext.OpDef(
        name="InnerOp",
        description="inner custom op",
        signature=ext.OpDefSig(tys.FunctionType.endo([tys.Bool])),
    )
    inner_ext = ext.Extension(
        version=ext.Version(0, 1, 0),
        name="inner",
        types={},
    )
    inner_ext.add_op_def(inner_op_def)

    # Build a simple identity HUGR to use as a lowering implementation.
    dfg = Dfg(tys.Bool)
    [b] = dfg.inputs()
    b = dfg.add_op(inner_ext.get_op("InnerOp").instantiate(args=[]), b).out(0)
    dfg.set_outputs(b)
    lower_hugr = dfg.hugr

    # Wrap it in a FixedHugr lower function.
    fixed_hugr = ext.FixedHugr(extensions=[inner_ext.name], hugr=lower_hugr)

    # Build an extension whose op carries the lower_func.
    outer_op_def = ext.OpDef(
        name="OuterOp",
        description="outer op with lowering",
        signature=ext.OpDefSig(tys.FunctionType.endo([tys.Bool])),
        lower_funcs=[fixed_hugr],
    )
    outer_ext = ext.Extension(
        version=ext.Version(0, 1, 0),
        name="outer",
        types={},
    )
    outer_ext.add_op_def(outer_op_def)

    # Include the extension in a package and check that the lower func is resolved.
    package = Package(
        modules=[],
        extensions=[outer_ext],
    )
    used_exts = package.used_extensions()
    assert outer_ext.name in used_exts.used_extensions
    assert inner_ext.name in used_exts.used_extensions

    bytes = package.to_bytes()
    package = Package.from_bytes(bytes, extensions=used_exts.used_extensions)
    assert len(package.extensions[0].operations["OuterOp"].lower_funcs) == 1

    used_exts = package.used_extensions()
    assert outer_ext.name in used_exts.used_extensions.ids()
    assert inner_ext.name in used_exts.used_extensions.ids()
