import contextlib
import dataclasses
import io
import pathlib
import re
import shutil
import subprocess
import sys
import tempfile
from typing import Annotated, List, Literal, Optional

import pytest

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


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

    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


def test_zsh(backend: str):
    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):
    """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 _strip_comments_and_normalize(script: str) -> str:
    """Strip comments and normalize whitespace for comparison.

    This removes:
    - Lines starting with # (comments)
    - Empty lines
    - Leading/trailing whitespace

    Then normalizes remaining whitespace for comparison.
    """
    lines = []
    for line in script.split("\n"):
        stripped = line.strip()
        # Skip empty lines and comment lines.
        if not stripped or stripped.startswith("#"):
            continue
        lines.append(stripped)
    return "\n".join(lines)


def test_completion_output_parity():
    """Test that argparse and tyro backends produce equivalent completion scripts.

    The scripts should be functionally identical (ignoring comments and formatting).
    """

    def simple_function(
        x: int = 5,
        y: Annotated[str, tyro.conf.arg(aliases=["-y"])] = "hello",
        flag: bool = False,
    ) -> None:
        """Simple function for testing completion parity."""

    # Generate completion with argparse backend.
    original_backend = tyro._experimental_options["backend"]
    tyro._experimental_options["backend"] = "argparse"

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

    # Generate completion with tyro backend.
    tyro._experimental_options["backend"] = "tyro"

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

    # Restore original backend.
    tyro._experimental_options["backend"] = original_backend

    # Strip comments and compare.
    argparse_normalized = _strip_comments_and_normalize(target_argparse.getvalue())
    tyro_normalized = _strip_comments_and_normalize(target_tyro.getvalue())

    # Different architectures but same functionality.
    # Argparse backend uses bash variables, tyro backend uses embedded Python.
    assert "_option_strings=" in argparse_normalized
    assert "COMPLETION_SPEC" in tyro_normalized
    assert "PYTHON_EOF" in tyro_normalized

    # Check that both have the same options (though in possibly different order/format).
    assert "--x" in argparse_normalized and "--x" in tyro_normalized
    assert "-y" in argparse_normalized and "-y" in tyro_normalized
    assert "--flag" in argparse_normalized and "--flag" in tyro_normalized
    assert "--no-flag" in argparse_normalized and "--no-flag" in tyro_normalized


def test_completion_parity_with_subcommands():
    """Test completion parity for simple subcommands (no CascadeSubcommandArgs)."""
    import dataclasses

    @dataclasses.dataclass
    class SubA:
        value_a: int = 1

    @dataclasses.dataclass
    class SubB:
        value_b: str = "hello"

    @dataclasses.dataclass
    class Config:
        sub: SubA | SubB

    # Generate completion with argparse backend.
    original_backend = tyro._experimental_options["backend"]
    tyro._experimental_options["backend"] = "argparse"

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

    # Generate completion with tyro backend.
    tyro._experimental_options["backend"] = "tyro"

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

    # Restore original backend.
    tyro._experimental_options["backend"] = original_backend

    argparse_output = target_argparse.getvalue()
    tyro_output = target_tyro.getvalue()

    # Both should have subcommands.
    assert "sub:sub-a" in argparse_output or "sub_a" in argparse_output
    assert "sub:sub-b" in argparse_output or "sub_b" in argparse_output
    assert "sub:sub-a" in tyro_output or "sub_a" in tyro_output
    assert "sub:sub-b" in tyro_output or "sub_b" in tyro_output

    # Both should have the subcommand arguments (possibly with prefix).
    assert (
        "--value-a" in argparse_output
        or "--value_a" in argparse_output
        or "--sub.value-a" in argparse_output
        or "--sub.value_a" in argparse_output
    )
    assert (
        "--value-b" in argparse_output
        or "--value_b" in argparse_output
        or "--sub.value-b" in argparse_output
        or "--sub.value_b" in argparse_output
    )
    assert (
        "--value-a" in tyro_output
        or "--value_a" in tyro_output
        or "--sub.value-a" in tyro_output
        or "--sub.value_a" in tyro_output
    )
    assert (
        "--value-b" in tyro_output
        or "--value_b" in tyro_output
        or "--sub.value-b" in tyro_output
        or "--sub.value_b" in tyro_output
    )


def test_path_completion(backend: str):
    """Test that path arguments get proper file/directory completion."""
    from pathlib import Path

    def path_function(
        input_file: Path,
        output_dir: Path,
        config_path: str,
    ) -> None:
        """Function with path arguments."""

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

    output = target.getvalue()

    if backend == "tyro":
        # New Python-based backend: check for path types in the spec.
        assert "'type': 'path'" in output
        assert "'is_directory': True" in output or "'is_directory': False" in output
    else:
        # Argparse backend uses shtab helper functions.
        assert "_shtab_compgen_files" in output or "_shtab_compgen_dirs" in output


def test_cascaded_subcommand_completion(backend: str):
    """Test completion for cascaded subcommands.

    This tests the tyro-specific CascadeSubcommandArgs feature, which allows
    more flexible argument ordering. The completion should include all available
    options at each level.
    """
    import dataclasses

    @dataclasses.dataclass
    class SubA:
        sub_value: int = 1

    @dataclasses.dataclass
    class SubB:
        sub_value: int = 2

    @dataclasses.dataclass
    class Config:
        parent_arg: str = "default"
        sub: SubA | SubB = dataclasses.field(default_factory=SubA)

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

    output = target.getvalue()

    # Both backends should generate completions with subcommands.
    assert "sub:sub-a" in output or "sub_a" in output
    assert "sub:sub-b" in output or "sub_b" in output
    # Parent argument should be available.
    assert "--parent-arg" in output or "--parent_arg" in output


def test_completion_uses_correct_prog_name(backend: str):
    """Test that completion scripts use the correct program name from sys.argv[0].

    The completion script should be registered for the actual script being run,
    not for a hardcoded 'tyro' command.
    """

    # Create a temporary script.
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False
    ) as script_file:
        script_file.write(
            """
import tyro

def main(x: int = 5):
    pass

if __name__ == "__main__":
    tyro.cli(main)
"""
        )
        script_path = script_file.name

    try:
        # Generate bash completion.
        result_bash = subprocess.run(
            [sys.executable, script_path, "--tyro-print-completion", "bash"],
            capture_output=True,
            text=True,
            check=False,
        )
        assert result_bash.returncode == 0
        bash_output = result_bash.stdout

        # Generate zsh completion.
        result_zsh = subprocess.run(
            [sys.executable, script_path, "--tyro-print-completion", "zsh"],
            capture_output=True,
            text=True,
            check=False,
        )
        assert result_zsh.returncode == 0
        zsh_output = result_zsh.stdout

        # Verify bash completion uses the script name.
        assert "complete -F" in bash_output
        assert script_path in bash_output
        if backend == "argparse":
            assert "complete -o filenames -F _tyro_tyro tyro" not in bash_output
        # For tyro backend, we just check that the script path is in the completion command.

        # Verify zsh completion uses the script name.
        assert f"#compdef {script_path}" in zsh_output
        assert "#compdef tyro" not in zsh_output

    finally:
        # Clean up.

        pathlib.Path(script_path).unlink(missing_ok=True)


class BashCompletionTester:
    """Helper class for testing bash completions functionally.

    Based on shtab's testing approach.
    """

    def __init__(self, completion_script: str):
        """Initialize with a completion script.

        Args:
            completion_script: The bash completion script to test.
        """
        self.completion_script = completion_script

    def test_compgen(
        self, compgen_cmd: str, word: str, expected_completions: str
    ) -> None:
        """Test that compgen produces expected completions.

        Args:
            compgen_cmd: The compgen command (e.g., '-W "foo bar"').
            word: The word to complete.
            expected_completions: Space-separated expected completions.
        """

        test_script = f'{self.completion_script}\n[[ "$(echo $(compgen {compgen_cmd} -- "{word}"))" = "{expected_completions}" ]]'

        proc = subprocess.Popen(
            ["bash", "-o", "pipefail", "-euc", test_script],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = proc.communicate()
        assert proc.returncode == 0, f"""compgen test failed:
compgen {compgen_cmd} -- "{word}"
Expected: {expected_completions}
=== stdout ===
{stdout.decode() if stdout else ""}
=== stderr ===
{stderr.decode() if stderr else ""}"""


def test_bash_functional_completion_simple(backend: str):
    """Test that bash completion actually works for a simple case."""

    @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", "bash"])

    completion_script = target.getvalue()
    tester = BashCompletionTester(completion_script)

    # Test that mode choices are available.
    # Look for the variable name in the completion script.
    match = re.search(r"(\w+___mode_choices)=\(", completion_script)
    if match:
        var_name = match.group(1)
        # Test completion of mode values.
        tester.test_compgen(f'-W "${{{var_name}[*]}}"', "tr", "train")
        tester.test_compgen(f'-W "${{{var_name}[*]}}"', "e", "eval")
        tester.test_compgen(f'-W "${{{var_name}[*]}}"', "te", "test")


def test_bash_functional_completion_with_subcommands(backend: str):
    """Test bash completion with subcommands."""

    @dataclasses.dataclass
    class DatasetMnist:
        binary: bool = False

    @dataclasses.dataclass
    class DatasetImageNet:
        subset: Literal[50, 100, 1000] = 50

    @dataclasses.dataclass
    class OptimizerConfig:
        dataset: DatasetMnist | DatasetImageNet

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

    completion_script = target.getvalue()
    BashCompletionTester(completion_script)

    # Test subcommand completion.
    # The completion script should have subparser/positional choices.
    match = re.search(
        r"(\w+_(?:subparsers|pos_0_choices))=\(([^)]+)\)", completion_script
    )
    if match:
        match.group(1)
        choices_str = match.group(2)
        # Extract the actual subcommand names from the choices.
        # They could be like 'dataset:dataset-mnist' 'dataset:dataset-image-net'.
        choices = [c.strip("' ") for c in re.findall(r"'([^']+)'", choices_str)]

        # Verify that the expected subcommands are present.
        assert any("mnist" in c.lower() for c in choices), (
            f"mnist not found in {choices}"
        )
        assert any("image" in c.lower() for c in choices), (
            f"image-net not found in {choices}"
        )


def test_bash_functional_completion_frontier_subcommands(backend: str):
    """Test bash completion with frontier (multiple Union) subcommands.

    Tests the scenario where subcommands can be selected from multiple Union fields.
    This is a key difference between backends:
    - tyro backend: Shows ALL frontier subcommands at once (flat)
    - argparse backend: Shows nested subcommands (sequential selection)
    """

    @dataclasses.dataclass
    class DatasetMnist:
        pass

    @dataclasses.dataclass
    class DatasetImageNet:
        pass

    @dataclasses.dataclass
    class OptimizerAdam:
        learning_rate: float = 1e-4

    @dataclasses.dataclass
    class OptimizerSgd:
        learning_rate: float = 1e-3

    @dataclasses.dataclass
    class Config:
        dataset: DatasetMnist | DatasetImageNet
        optimizer: OptimizerAdam | OptimizerSgd

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

    completion_script = target.getvalue()

    # Both backends should have all subcommands somewhere in the completion script.
    assert "mnist" in completion_script.lower()
    assert "image" in completion_script.lower()
    assert "adam" in completion_script.lower()
    assert "sgd" in completion_script.lower()

    # Check the key difference: frontier behavior.
    if backend == "tyro":
        # Tyro backend uses Python spec with frontier_groups.
        assert "'frontier_groups':" in completion_script
        # Should have 2 groups (datasets and optimizers).
        assert "[[" in completion_script  # Nested lists indicate frontier groups.

        # Verify all 4 subcommands are in the spec.
        assert (
            "'dataset" in completion_script.lower()
            or "mnist" in completion_script.lower()
        )
        assert (
            "'optimizer" in completion_script.lower()
            or "adam" in completion_script.lower()
        )
    else:  # argparse backend
        # Find the main subparsers array definition.
        main_subparsers_match = re.search(
            r"(_\w+)_subparsers=\(([^)]+)\)", completion_script
        )
        assert main_subparsers_match, "Could not find main subparsers array"

        main_subparsers_content = main_subparsers_match.group(2)

        # Argparse backend: Nested structure - only shows dataset subcommands initially.
        # Should have 2 subcommands in main array (just datasets).
        subcommand_count = len(
            re.findall(r"'[^']*(?:mnist|image|adam|sgd)[^']*'", main_subparsers_content)
        )
        assert subcommand_count == 2, (
            f"Argparse backend should show 2 dataset subcommands in main array, "
            f"found {subcommand_count}: {main_subparsers_content}"
        )
        # Optimizer subcommands should be in nested arrays.
        assert re.search(
            r"_dataset_[^_]+_subparsers=\([^)]*adam[^)]*\)", completion_script
        ), "Argparse backend should have nested optimizer subparsers"


def test_bash_functional_completion_cascade_subcommand_args(backend: str):
    """Test bash completion with CascadeSubcommandArgs marker.

    CascadeSubcommandArgs allows arguments to appear before or after subcommands.
    """

    @dataclasses.dataclass
    class SubcommandA:
        arg_a: int = 1

    @dataclasses.dataclass
    class SubcommandB:
        arg_b: str = "hello"

    @dataclasses.dataclass
    class Config:
        """Config with CascadeSubcommandArgs."""

        common: Annotated[int, tyro.conf.arg(aliases=["-c"])] = 5
        subcommand: Annotated[
            SubcommandA | SubcommandB, tyro.conf.CascadeSubcommandArgs
        ] = dataclasses.field(default_factory=SubcommandA)

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

    completion_script = target.getvalue()
    BashCompletionTester(completion_script)

    # Verify that both the common argument and subcommands are in the completion.
    assert "--common" in completion_script or "-c" in completion_script
    assert "subcommand-a" in completion_script.lower() or "a" in completion_script

    # Verify subcommand choices exist.
    match = re.search(
        r"(\w+_(?:subparsers|pos_0_choices))=\(([^)]+)\)", completion_script
    )
    if match:
        match.group(1)
        choices_str = match.group(2)
        choices = [c.strip("' ") for c in re.findall(r"'([^']+)'", choices_str)]

        # Both subcommands should be available.
        assert any("a" in c.lower() for c in choices), (
            f"subcommand-a not found in {choices}"
        )
        assert any("b" in c.lower() for c in choices), (
            f"subcommand-b not found in {choices}"
        )


class ZshCompletionTester:
    """Helper class for testing zsh completions functionally using zpty.

    Based on research from Stack Overflow and Unix Stack Exchange about zsh completion testing.
    """

    def __init__(self, completion_script: str, prog: str):
        """Initialize with a completion script.

        Args:
            completion_script: The zsh completion script to test.
            prog: The program name (for compdef).
        """
        self.completion_script = completion_script
        self.prog = prog

    def test_completions(
        self, command_prefix: str, expected_completions: List[str]
    ) -> None:
        """Test that zsh produces expected completions for a command prefix.

        Args:
            command_prefix: The command to complete (e.g., "prog --mo").
            expected_completions: List of expected completion strings.
        """

        # Create a zsh script that uses zpty to test completions.
        # This is based on the Stack Exchange pattern for testing zsh completions.
        test_script = f"""
zmodload zsh/zpty

# Load the completion system.
autoload -U compinit
compinit

# Source the completion script.
{self.completion_script}

# Create a test function that captures completions.
comptest() {{
    # Override compadd to capture completions.
    typeset -ga completions
    completions=()

    compadd() {{
        local -a reply
        builtin compadd -O reply "$@"
        completions+=("${{reply[@]}}")
        builtin compadd "$@"
    }}

    # Bind tab to completion widget.
    bindkey "^I" complete-word
    zle -C {{,,}}complete-word
    complete-word() {{
        unset 'compstate[vared]'
        _main_complete "$@"
        # Output completions with markers.
        print -n $'\\002'  # Start marker (^B).
        print -nlr -- "${{completions[@]}}"
        print -n $'\\003'  # End marker (^C).
        exit
    }}

    # Use vared to activate the completion context.
    stty rows 100000 cols 2000
    vared -c tmp
}}

# Create pty and run the test.
zpty pty comptest
zpty -w pty "{command_prefix}"$'\\t'
zpty -r pty REPLY $'*\\002'  # Read until start marker.
zpty -r pty REPLY $'*\\003'  # Read until end marker.
print -r -- "${{REPLY%$'\\003'}}"  # Output completions without end marker.
zpty -d pty
"""

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".zsh", delete=False
        ) as script_file:
            script_file.write(test_script)
            script_path = script_file.name

        proc: subprocess.Popen | None = None
        try:
            proc = subprocess.Popen(
                ["zsh", script_path],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            stdout, stderr = proc.communicate(timeout=5)

            # Parse completions from output.
            actual_completions = [
                line.strip() for line in stdout.strip().split("\n") if line.strip()
            ]

            # Check if expected completions are in actual completions.
            for expected in expected_completions:
                assert any(
                    expected in actual for actual in actual_completions
                ), f"""Zsh completion test failed:
Command: {command_prefix}
Expected (at least one match for): {expected}
Actual completions: {actual_completions}
Stderr: {stderr}"""

        except subprocess.TimeoutExpired:
            if proc is not None:
                proc.kill()
            raise AssertionError(
                f"Zsh completion test timed out for command: {command_prefix}"
            )
        finally:
            pathlib.Path(script_path).unlink(missing_ok=True)


def test_zsh_functional_completion_simple(backend: str):
    """Test that zsh completion actually works for a simple case."""

    # 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()
    ZshCompletionTester(completion_script, "test_prog")

    # Test that mode completions work.
    # Note: Actual zsh completion testing is complex and may require special setup.
    # For now, we'll just 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


def test_cascade_marker_detection(backend: str):
    """Test that CascadeSubcommandArgs marker is properly detected in completion spec."""
    if backend != "tyro":
        pytest.skip("Cascade marker detection is tyro-specific")

    @dataclasses.dataclass
    class Config:
        regular_field: int = 5
        cascade_field: tyro.conf.CascadeSubcommandArgs[str] = "default"

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

    completion_script = target.getvalue()

    # The completion spec is embedded in the script.
    # Verify it contains the fields and cascade info.
    assert "regular-field" in completion_script or "regular_field" in completion_script
    assert "cascade-field" in completion_script or "cascade_field" in completion_script

    # Verify the spec has cascade markers.
    # The cascade field should be tracked in the spec (using Python dict syntax with single quotes).
    assert "'cascade'" in completion_script


def test_nargs_with_choices_completion(backend: str):
    """Test that nargs is properly tracked for choice options in completion spec."""
    if backend != "tyro":
        pytest.skip("Choice nargs tracking is tyro-specific")

    @dataclasses.dataclass
    class Config:
        # Single choice.
        mode: Literal["train", "eval"] = "train"
        # Multiple choices.
        modes: List[Literal["train", "eval", "test"]] = dataclasses.field(
            default_factory=lambda: ["train"]
        )

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

    completion_script = target.getvalue()

    # Verify completion spec has choices and nargs.
    assert "'choices'" in completion_script
    assert "train" in completion_script
    assert "eval" in completion_script
    assert "test" in completion_script

    # Verify nargs is tracked for the list field.
    assert "'nargs'" in completion_script


def test_metavar_in_description(backend: str):
    """Test that metavar is included in option descriptions."""
    if backend != "tyro":
        pytest.skip("Metavar formatting is tyro-specific")

    @dataclasses.dataclass
    class Config:
        count: int = 5
        name: str = "default"

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

    completion_script = target.getvalue()

    # Verify metavar (type hints) appear in descriptions.
    # The descriptions should include INT and STR metavars.
    assert "INT" in completion_script or "int" in completion_script
    assert "STR" in completion_script or "str" in completion_script


@dataclasses.dataclass
class _ConfigForBulletTest:
    # No custom helptext, just default.
    simple: int = 5
    # Custom helptext.
    documented: int = 10
    """This is a custom help message."""


def test_smart_bullet_separator(backend: str):
    """Test that bullet separator is only used when there's custom helptext."""
    if backend != "tyro":
        pytest.skip("Bullet separator logic is tyro-specific")

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

    completion_script = target.getvalue()

    # Verify bullet separator (•) is used for fields with custom helptext.
    # The custom help message should have the bullet.
    assert "custom help" in completion_script.lower()
    assert "•" in completion_script


@dataclasses.dataclass
class _PreprocessingA:
    scale: float = 1.0


@dataclasses.dataclass
class _PreprocessingB:
    normalize: bool = True


@dataclasses.dataclass
class _AugmentationA:
    rotate: bool = True


@dataclasses.dataclass
class _AugmentationB:
    flip: bool = False


@dataclasses.dataclass
class _DatasetMNIST:
    """Dataset with its own frontier groups."""

    preprocessing: _PreprocessingA | _PreprocessingB
    augmentation: _AugmentationA | _AugmentationB


@dataclasses.dataclass
class _DatasetImageNet:
    resolution: int = 224


@dataclasses.dataclass
class _SubcommandWithFrontier:
    """A subcommand that has subcommands, one of which has frontier groups."""

    dataset: _DatasetMNIST | _DatasetImageNet


@dataclasses.dataclass
class _NestedMainCommand:
    # This creates a subcommand that has frontier groups.
    config: _SubcommandWithFrontier


def test_nested_subcommand_coverage(backend: str):
    """Test nested subcommands for coverage of _build_subcommand_spec recursion."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion generation")

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

    completion_script = target.getvalue()
    # Verify nested subcommands are present.
    assert "config" in completion_script
    # Check for deeply nested frontier groups within subcommands.
    assert "dataset" in completion_script or "preprocessing" in completion_script


def test_positional_argument_coverage(backend: str):
    """Test that positional arguments are skipped in completion spec."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion generation")

    def main(input_file: tyro.conf.Positional[str], output: str = "out.txt") -> str:
        """Test with positional argument.

        Args:
            input_file: Input file path (positional).
            output: Output file path.
        """
        return f"{input_file} -> {output}"

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

    completion_script = target.getvalue()
    # Positional args shouldn't appear in options list, only regular flags.
    assert "--output" in completion_script


def test_count_action_coverage(backend: str):
    """Test count action type for coverage."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion generation")

    def main(verbose: tyro.conf.UseCounterAction[int]) -> int:
        """Test function with counter action.

        Args:
            verbose: Verbosity level.
        """
        return verbose

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

    completion_script = target.getvalue()
    # Count action creates a flag-type option.
    assert "flag" in completion_script.lower() or "verbose" in completion_script


def test_metavar_and_helptext_edge_cases(backend: str):
    """Test that completion works even with minimal help text."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion generation")

    def main(value: int = 5, flag: bool = False) -> int:
        """Test completion with basic arguments."""
        return value

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

    completion_script = target.getvalue()
    # Verify basic completion generation works.
    assert "value" in completion_script or "flag" in completion_script


def test_unsupported_shell_error(backend: str):
    """Test that unsupported shell types raise an error."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific error handling")

    def main(value: int = 5) -> int:
        """Test function."""
        return value

    # Test that tcsh (supported by argparse but not tyro) raises ValueError.
    with pytest.raises(ValueError, match="Unsupported shell.*tcsh"):
        tyro.cli(main, args=["--tyro-print-completion", "tcsh"])


def test_reconstruct_colon_words_basic():
    """Test basic word reconstruction for colon-separated subcommands."""
    from tyro._backends._completion._completion_script import reconstruct_colon_words

    # Spec with colon-separated subcommands.
    spec = {
        "subcommands": {
            "dataset:mnist": {},
            "dataset:image-net": {},
            "optimizer:adam": {},
        }
    }

    # Test 1: Basic reconstruction of "dataset:mnist".
    # Bash splits this as ["dataset", ":", "mnist"].
    words = ["dataset", ":", "mnist"]
    reconstructed, new_cword = reconstruct_colon_words(words, 2, spec)
    assert reconstructed == ["dataset:mnist"]
    assert new_cword == 0

    # Test 2: Partial completion "dataset:".
    # Bash splits this as ["dataset", ":"].
    words = ["dataset", ":"]
    reconstructed, new_cword = reconstruct_colon_words(words, 1, spec)
    assert reconstructed == ["dataset:"]
    assert new_cword == 0

    # Test 3: Partial prefix "dataset:m".
    # Bash splits this as ["dataset", ":", "m"].
    words = ["dataset", ":", "m"]
    reconstructed, new_cword = reconstruct_colon_words(words, 2, spec)
    assert reconstructed == ["dataset:m"]
    assert new_cword == 0


def test_reconstruct_colon_words_no_match():
    """Test that words are NOT reconstructed when they don't match subcommands."""
    from tyro._backends._completion._completion_script import reconstruct_colon_words

    spec = {"subcommands": {"dataset:mnist": {}}}

    # Test 1: Non-matching colon pattern (e.g., option value "key:value").
    # Should NOT be merged since "key:value" is not a known subcommand.
    words = ["--config", "key", ":", "value"]
    reconstructed, new_cword = reconstruct_colon_words(words, 3, spec)
    # Since "key:value" doesn't match any subcommand, keep them separate.
    # Standalone colons are skipped.
    assert reconstructed == ["--config", "key", "value"]
    assert new_cword == 2


def test_reconstruct_colon_words_multiple():
    """Test reconstruction with multiple colon-separated subcommands."""
    from tyro._backends._completion._completion_script import reconstruct_colon_words

    spec = {
        "subcommands": {
            "dataset:mnist": {},
            "optimizer:adam": {},
        }
    }

    # Test: Multiple subcommands "dataset:mnist optimizer:adam".
    # Bash splits as ["dataset", ":", "mnist", "optimizer", ":", "adam"].
    words = ["dataset", ":", "mnist", "optimizer", ":", "adam"]
    reconstructed, new_cword = reconstruct_colon_words(words, 5, spec)
    assert reconstructed == ["dataset:mnist", "optimizer:adam"]
    assert (
        new_cword == 1
    )  # Cursor on "adam" -> after reconstruction it's on second word.


def test_reconstruct_colon_words_with_options():
    """Test reconstruction with flags mixed in."""
    from tyro._backends._completion._completion_script import reconstruct_colon_words

    spec = {
        "subcommands": {
            "dataset:mnist": {},
        }
    }

    # Test: Subcommand with options "dataset:mnist --lr 0.001".
    # Bash splits as ["dataset", ":", "mnist", "--lr", "0.001"].
    words = ["dataset", ":", "mnist", "--lr", "0.001"]
    reconstructed, new_cword = reconstruct_colon_words(words, 4, spec)
    assert reconstructed == ["dataset:mnist", "--lr", "0.001"]
    assert new_cword == 2  # Cursor on "0.001" -> after merging, it's at index 2.


def test_reconstruct_colon_words_cursor_on_colon():
    """Test cursor position when it's on the colon itself."""
    from tyro._backends._completion._completion_script import reconstruct_colon_words

    spec = {"subcommands": {"dataset:mnist": {}}}

    # Test: Cursor is on the colon character.
    # Words: ["dataset", ":"], cursor at index 1 (the colon).
    words = ["dataset", ":"]
    reconstructed, new_cword = reconstruct_colon_words(words, 1, spec)
    # The colon should be merged with "dataset" to form "dataset:".
    assert reconstructed == ["dataset:"]
    # Cursor should be on the merged word.
    assert new_cword == 0


def test_reconstruct_colon_words_empty():
    """Test reconstruction with empty word list."""
    from tyro._backends._completion._completion_script import reconstruct_colon_words

    spec = {"subcommands": {}}

    # Empty word list.
    words = []
    reconstructed, new_cword = reconstruct_colon_words(words, 0, spec)
    assert reconstructed == []
    assert new_cword == 0
