"""Tests for tuples containing unions with different lengths."""

from __future__ import annotations

import dataclasses
from typing import List, Literal, Optional, Tuple

import pytest

import tyro
from tyro.constructors._primitive_spec import (
    PrimitiveConstructorSpec,
    PrimitiveTypeInfo,
)


@dataclasses.dataclass
class _TestInner:
    values: Tuple[float, float] | Tuple[float, float, float]


@dataclasses.dataclass
class _TestOuter:
    inner: _TestInner
    extra: Tuple[str] | Tuple[str, str]


def test_basic_union_of_tuples():
    """Test basic union of tuples with different lengths."""

    def main(x: Tuple[int, int] | Tuple[int, int, int]) -> Tuple[int, ...]:
        return x

    # Test with 2 elements.
    assert tyro.cli(main, args=["--x", "1", "2"]) == (1, 2)

    # Test with 3 elements.
    assert tyro.cli(main, args=["--x", "1", "2", "3"]) == (1, 2, 3)

    # Test with wrong number of elements.
    with pytest.raises(SystemExit):
        tyro.cli(main, args=["--x", "1"])

    with pytest.raises(SystemExit):
        tyro.cli(main, args=["--x", "1", "2", "3", "4"])


def test_union_of_tuples_with_strings():
    """Test union of tuples with different types."""

    def main(x: Tuple[str, str] | Tuple[int, int, int]) -> tuple:
        return x

    # Test with 2 strings.
    assert tyro.cli(main, args=["--x", "hello", "world"]) == ("hello", "world")

    # Test with 3 integers.
    assert tyro.cli(main, args=["--x", "1", "2", "3"]) == (1, 2, 3)


def test_nested_tuple_with_union() -> None:
    """Test nested tuple containing a union with different lengths."""

    @dataclasses.dataclass
    class Config:
        x: Tuple[Tuple[int, int] | Tuple[int, int, int]]

    # Test with inner tuple of 2 elements.
    result = tyro.cli(Config, args=["--x", "1", "2"])
    assert result.x == ((1, 2),)

    # Test with inner tuple of 3 elements.
    result = tyro.cli(Config, args=["--x", "1", "2", "3"])
    assert result.x == ((1, 2, 3),)


def test_optional_nested_tuple_with_union() -> None:
    """Test optional nested tuple containing a union."""

    @dataclasses.dataclass
    class Config:
        x: Optional[Tuple[Tuple[int, int] | Tuple[int, int, int]]]

    # Test with None.
    result = tyro.cli(Config, args=["--x", "None"])
    assert result.x is None

    # Test with tuple of 2 elements.
    result = tyro.cli(Config, args=["--x", "1", "2"])
    assert result.x == ((1, 2),)

    # Test with tuple of 3 elements.
    result = tyro.cli(Config, args=["--x", "1", "2", "3"])
    assert result.x == ((1, 2, 3),)


def test_multiple_unions_in_tuple():
    """Test tuple with multiple union elements."""

    def main(x: Tuple[Tuple[int] | Tuple[int, int], str]) -> tuple:
        return x

    # Test with single int and string.
    assert tyro.cli(main, args=["--x", "1", "hello"]) == ((1,), "hello")

    # Test with pair of ints and string.
    assert tyro.cli(main, args=["--x", "1", "2", "world"]) == ((1, 2), "world")


def test_deeply_nested_unions():
    """Test deeply nested structures with variable-length unions."""

    # Test a simple list of variable-length tuples.
    def main(x: List[Tuple[int, int] | Tuple[int, int, int]]) -> list:
        return x

    # Test with a single tuple.
    result = tyro.cli(main, args=["--x", "1", "2"])
    assert result == [(1, 2)]

    # Test with a different length tuple.
    result = tyro.cli(main, args=["--x", "3", "4", "5"])
    assert result == [(3, 4, 5)]


def test_union_with_literals():
    """Test union of tuples containing literals."""

    def main(
        x: tuple[Literal["a"], int] | tuple[Literal["b"], int, int],
    ) -> tuple:
        return x

    # Test with literal "a".
    assert tyro.cli(main, args=["--x", "a", "1"]) == ("a", 1)

    # Test with literal "b".
    assert tyro.cli(main, args=["--x", "b", "1", "2"]) == ("b", 1, 2)

    # Test with wrong literal.
    with pytest.raises(SystemExit):
        tyro.cli(main, args=["--x", "c", "1"])


def test_complex_nested_structure():
    """Test complex nested structure with variable-length tuples."""

    # Test with all minimum lengths.
    result = tyro.cli(
        _TestOuter, args=["--inner.values", "1.0", "2.0", "--extra", "hello"]
    )
    assert result.inner.values == (1.0, 2.0)
    assert result.extra == ("hello",)

    # Test with all maximum lengths.
    result = tyro.cli(
        _TestOuter,
        args=["--inner.values", "1.0", "2.0", "3.0", "--extra", "hello", "world"],
    )
    assert result.inner.values == (1.0, 2.0, 3.0)
    assert result.extra == ("hello", "world")


def test_helptext_for_variable_length_tuples():
    """Test that help text is generated correctly."""

    def main(x: Tuple[int, int] | Tuple[int, int, int]) -> tuple:
        """Test function.

        Args:
            x: A tuple with either 2 or 3 integers.
        """
        return x

    # Just check that help doesn't crash.
    with pytest.raises(SystemExit) as exc_info:
        tyro.cli(main, args=["--help"])
    assert exc_info.value.code == 0


def test_primitive_spec_nargs_computation():
    """Test that nargs is computed correctly for variable-length tuples."""
    from tyro.constructors._registry import ConstructorRegistry

    registry = ConstructorRegistry()

    # Test simple union of tuples.
    type_info = PrimitiveTypeInfo.make(
        Tuple[int, int] | Tuple[int, int, int],  # type: ignore
        parent_markers=set(),
    )
    spec = registry.get_primitive_spec(type_info)
    assert isinstance(spec, PrimitiveConstructorSpec)
    assert spec.nargs == (2, 3)

    # Test nested tuple with union.
    type_info = PrimitiveTypeInfo.make(
        Tuple[Tuple[int, int] | Tuple[int, int, int]], parent_markers=set()
    )
    spec = registry.get_primitive_spec(type_info)
    assert isinstance(spec, PrimitiveConstructorSpec)
    assert spec.nargs == (2, 3)


def test_error_on_star_nargs():
    """Test that star nargs in tuples raises appropriate error."""
    from tyro.constructors._primitive_spec import UnsupportedTypeAnnotationError
    from tyro.constructors._registry import ConstructorRegistry

    registry = ConstructorRegistry()

    # This should raise an error because list[int] has nargs="*".
    type_info = PrimitiveTypeInfo.make(Tuple[List[int], int], parent_markers=set())
    spec = registry.get_primitive_spec(type_info)
    assert not isinstance(spec, UnsupportedTypeAnnotationError)
    assert spec.nargs == "*"


def test_backtracking_with_choices():
    """Test that backtracking works correctly with choices."""

    def main(
        x: Tuple[Literal["a", "b"], int] | Tuple[Literal["c"], int, int],
    ) -> tuple:
        return x

    # Test valid choices.
    assert tyro.cli(main, args=["--x", "a", "1"]) == ("a", 1)
    assert tyro.cli(main, args=["--x", "b", "1"]) == ("b", 1)
    assert tyro.cli(main, args=["--x", "c", "1", "2"]) == ("c", 1, 2)

    # Test invalid choice.
    with pytest.raises(SystemExit):
        tyro.cli(main, args=["--x", "d", "1"])


def test_ambiguous_parsing():
    """Test cases where parsing could be ambiguous."""

    def main(
        x: Tuple[Tuple[int] | Tuple[int, int], Tuple[int] | Tuple[int, int]],
    ) -> tuple:
        return x

    # This should parse as ((1,), (2, 3)).
    assert tyro.cli(main, args=["--x", "1", "2", "3"]) == ((1,), (2, 3))

    # This should parse as ((1, 2), (3,)).
    assert tyro.cli(main, args=["--x", "1", "2", "3"]) == ((1,), (2, 3))
    # Note: Due to greedy left-to-right parsing, this will actually be ((1,), (2, 3)).
    # This is expected behavior with the backtracking algorithm.


def test_multiple_variable_length_elements():
    """Test tuple with multiple variable-length elements."""

    def main(
        x: Tuple[
            Tuple[int, int] | Tuple[int, int, int],
            Tuple[str] | Tuple[str, str],
            Tuple[float, float, float] | Tuple[float, float, float, float],
        ],
    ) -> tuple:
        return x

    # Test with all minimum lengths.
    result = tyro.cli(main, args=["--x", "1", "2", "hello", "1.0", "2.0", "3.0"])
    assert result == ((1, 2), ("hello",), (1.0, 2.0, 3.0))

    # Test with mixed lengths.
    result = tyro.cli(
        main, args=["--x", "1", "2", "3", "hello", "world", "1.0", "2.0", "3.0", "4.0"]
    )
    assert result == ((1, 2, 3), ("hello", "world"), (1.0, 2.0, 3.0, 4.0))


if __name__ == "__main__":
    # Run a simple test when executed directly.
    test_basic_union_of_tuples()
    print("Basic test passed!")
