"""Tests for zsh completion script generation.

TODO: it would be helpful to have end-to-end tests that actually run the generated
completion scripts in a zsh shell.
"""

import contextlib
import dataclasses
import io
import shutil
from typing import Annotated, Literal, Optional

import pytest

import tyro


@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: TypeA | TypeB = TypeA()


def test_zsh(backend: str) -> None:
    """Test basic zsh completion script generation."""
    target = io.StringIO()
    with pytest.raises(SystemExit), contextlib.redirect_stdout(target):
        tyro.cli(Wrapper, args=["--tyro-print-completion", "zsh"])

    output = target.getvalue()
    # Check for appropriate signature based on backend.
    if backend == "argparse":
        assert "# AUTOMATICALLY GENERATED by `shtab`" in output
    else:  # tyro backend.
        assert "# AUTOMATICALLY GENERATED by tyro" in output
        # Check that zsh-specific syntax is present.
        assert "#compdef" in output
        # New tyro backend uses embedded Python, not _arguments.
        assert "PYTHON_EOF" in output or "_arguments" in output


def test_completion_zsh(backend: str) -> None:
    """Test zsh completion with choices.

    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)

    # Check for appropriate signature based on backend.
    if backend == "argparse":
        assert "# AUTOMATICALLY GENERATED by `shtab`" in completion_script
    else:  # tyro backend.
        assert "# AUTOMATICALLY GENERATED by tyro" in completion_script

    # Both backends should have choices.
    if backend == "argparse":
        assert "preset_choices=(" in completion_script
        assert "p_choices=(" in completion_script
        assert "frame_choices=(" in completion_script
        assert "f_choices=(" in completion_script
    else:  # tyro backend uses Python spec.
        # Check that choices are in the embedded Python spec.
        assert "'choices':" in completion_script
        assert "rgb" in completion_script
        assert "depth" in completion_script
        assert "world" in completion_script
        assert "base" in completion_script


def test_zsh_functional_completion_simple(backend: str) -> None:
    """Test that zsh completion script is generated correctly."""

    # Skip if zsh is not available.
    if shutil.which("zsh") is None:
        pytest.skip("zsh not available")

    @dataclasses.dataclass
    class Config:
        mode: Literal["train", "eval", "test"] = "train"
        verbose: bool = False

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

    completion_script = target.getvalue()

    # Verify the script structure is correct.
    assert "PYTHON_EOF" in completion_script or "_arguments" in completion_script
    assert "train" in completion_script or "eval" in completion_script
