from collections.abc import Callable

import pytest

from hugr import tys as ht

from hugr import Wire

from guppylang.decorator import guppy
from guppylang.definition.custom import CustomCallCompiler
from guppylang.std.builtins import array
from guppylang.std.quantum import qubit


def test_id(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(x: int) -> int:
        return foo(x)

    validate(guppy.compile(main))


def test_id_nested(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(x: int) -> int:
        return foo(foo(foo(x)))

    validate(guppy.compile(main))


def test_use_twice(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(x: int, y: bool) -> None:
        foo(x)
        foo(y)

    validate(guppy.compile(main))


def test_define_twice(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy.declare
    def bar(x: T) -> T:  # Reuse same type var!
        ...

    @guppy
    def main(x: bool, y: float) -> None:
        foo(x)
        foo(y)

    validate(guppy.compile(main))


def test_return_tuple_implicit(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(x: int) -> tuple[int, int]:
        return foo((x, 0))

    validate(guppy.compile(main))


def test_same_args(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T, y: T) -> None: ...

    @guppy
    def main(x: int) -> None:
        foo(x, 42)

    validate(guppy.compile(main))


def test_different_args(validate):
    S = guppy.type_var("S")
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: S, y: T, z: tuple[S, T]) -> T: ...

    @guppy
    def main(x: int, y: float) -> float:
        return foo(x, y, (x, y)) + foo(y, 42.0, (0.0, y))

    validate(guppy.compile(main))


def test_nat_args(validate):
    n = guppy.nat_var("n")

    @guppy.declare
    def foo(x: array[int, n]) -> array[int, n]: ...

    @guppy
    def main(x: array[int, 42]) -> array[int, 42]:
        return foo(x)

    validate(guppy.compile(main))


def test_infer_basic(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo() -> T: ...

    @guppy
    def main() -> None:
        x: int = foo()

    validate(guppy.compile(main))


def test_infer_list(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo() -> T: ...

    @guppy
    def main() -> None:
        xs: list[int] = [foo()]
        ys = [1.0, foo()]

    validate(guppy.compile(main))


def test_infer_nested(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo() -> T: ...

    @guppy.declare
    def bar(x: T) -> T: ...

    @guppy
    def main() -> None:
        x: int = bar(foo())

    validate(guppy.compile(main))


def test_infer_left_to_right(validate):
    S = guppy.type_var("S")
    T = guppy.type_var("T")

    @guppy.declare
    def foo() -> T: ...

    @guppy.declare
    def bar(x: T, y: T, z: S, a: tuple[T, S]) -> None: ...

    @guppy
    def main() -> None:
        bar(42, foo(), False, foo())

    validate(guppy.compile(main))


def test_type_apply_basic(validate):
    S = guppy.type_var("S")
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy.declare
    def bar(x: S, y: T) -> S: ...

    @guppy
    def main() -> tuple[int, float, float]:
        x = foo[int](0)
        y = foo[float](1.0)
        z = bar[float, int](y, x)
        return x, y, z

    validate(guppy.compile(main))


def test_type_apply_higher_order(validate):
    S = guppy.type_var("S")
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy.declare
    def bar(x: S, y: T) -> S: ...

    @guppy
    def main() -> tuple[int, float, float]:
        f = foo[int]
        g = foo[float]
        h = bar[float, int]
        x = f(0)
        y = g(1.0)
        z = h(y, x)
        return x, y, z

    validate(guppy.compile(main))


def test_type_apply_nat(validate):
    n = guppy.nat_var("n")

    @guppy.declare
    def foo(x: array[int, n]) -> int: ...

    @guppy
    def main() -> int:
        return foo[0](array()) + foo[2](array(1, 2))

    validate(guppy.compile(main))


def test_type_apply_empty_tuple(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> None:
        ...

    @guppy
    def main() -> None:
        # `()` is the type of an empty tuple (`tuple[]` is not syntactically valid)
        foo[()]

    validate(guppy.compile(main))


def test_pass_poly_basic(validate):
    T = guppy.type_var("T")

    @guppy.declare
    def foo(f: Callable[[T], T]) -> None: ...

    @guppy.declare
    def bar(x: int) -> int: ...

    @guppy
    def main() -> None:
        foo(bar)

    validate(guppy.compile(main))


def test_pass_poly_cross(validate):
    S = guppy.type_var("S")
    T = guppy.type_var("T")

    @guppy.declare
    def foo(f: Callable[[S], int]) -> None: ...

    @guppy.declare
    def bar(x: bool) -> T: ...

    @guppy
    def main() -> None:
        foo(bar)

    validate(guppy.compile(main))


def test_linear(validate):
    T = guppy.type_var("T", copyable=False, droppable=False)

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(q: qubit) -> qubit:
        return foo(q)

    validate(guppy.compile(main))


def test_affine(validate):
    T = guppy.type_var("T", copyable=False, droppable=True)

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(a: array[int, 7]) -> None:
        foo(a)

    validate(guppy.compile(main))


def test_relevant(validate):
    T = guppy.type_var("T", copyable=True, droppable=False)

    @guppy.type(ht.Bool, copyable=True, droppable=False)
    class R: ...

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(r: R) -> R:
        r_copy = r
        return foo(r_copy)

    validate(guppy.compile(main))


def test_pass_nonlinear(validate):
    T = guppy.type_var("T", copyable=False, droppable=False)

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy
    def main(x: int) -> None:
        foo(x)

    validate(guppy.compile(main))


def test_pass_linear(validate):
    T = guppy.type_var("T", copyable=False, droppable=False)

    @guppy.declare
    def foo(f: Callable[[T], T]) -> None: ...

    @guppy.declare
    def bar(q: qubit) -> qubit: ...

    @guppy
    def main() -> None:
        foo(bar)

    validate(guppy.compile(main))


def test_custom_higher_order():
    class CustomCompiler(CustomCallCompiler):
        def compile(self, args: list[Wire]) -> list[Wire]:
            return args


    T = guppy.type_var("T")

    @guppy.custom(CustomCompiler())
    def foo(x: T) -> T: ...

    @guppy
    def main(x: int) -> int:
        f: Callable[[int], int] = foo
        return f(x)


@pytest.mark.skip("Higher-order polymorphic functions are not yet supported")
def test_higher_order_value(validate):

    T = guppy.type_var("T")

    @guppy.declare
    def foo(x: T) -> T: ...

    @guppy.declare
    def bar(x: T) -> T: ...

    @guppy
    def main(b: bool) -> int:
        f = foo if b else bar
        return f(42)

    validate(guppy.compile(main))
