import contextlib
import dataclasses
import io
from typing import Union

import pytest
from typing_extensions import Annotated, Literal, Optional

import tyro


# https://github.com/brentyi/tyro/issues/9
@dataclasses.dataclass(frozen=True)
class Subtype:
    data: int = 1


@dataclasses.dataclass(frozen=True)
class TypeA:
    subtype: Subtype = Subtype(1)


@dataclasses.dataclass(frozen=True)
class TypeB:
    subtype: Subtype = Subtype(2)


@dataclasses.dataclass(frozen=True)
class Wrapper:
    supertype: Union[TypeA, TypeB] = TypeA()


def test_bash():
    target = io.StringIO()
    with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
        tyro.cli(Wrapper, args=["--tyro-print-completion", "bash"])
    assert "# AUTOMATICALLY GENERATED by `shtab`" in target.getvalue()


def test_zsh():
    target = io.StringIO()
    with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
        tyro.cli(Wrapper, args=["--tyro-print-completion", "zsh"])
    assert "# AUTOMATICALLY GENERATED by `shtab`" in target.getvalue()


def test_completion_zsh():
    """https://github.com/brentyi/tyro/issues/158"""

    def start_device(
        preset: Annotated[
            Optional[Literal["rgb", "depth", "ir"]], tyro.conf.arg(aliases=["-p"])
        ] = None,
        frame: Annotated[
            Literal["world", "base"], tyro.conf.arg(aliases=["-f"])
        ] = "world",
    ) -> None:
        """
        Start device with the given preset.

        :param preset: device preset to use.
        :param frame: coordinate frame to use.
        """

    target = io.StringIO()
    with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
        tyro.cli(start_device, args=["--tyro-print-completion", "bash"])

    completion_script = target.getvalue()
    print(completion_script)
    assert "# AUTOMATICALLY GENERATED by `shtab`" in completion_script
    assert "preset_choices=(" in completion_script
    assert "p_choices=(" in completion_script
    assert "frame_choices=(" in completion_script
    assert "f_choices=(" in completion_script
