"""Tests for bash completion script generation and functional testing."""

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

import pytest

import tyro

pytestmark = pytest.mark.skipif(
    sys.platform == "win32", reason="Bash not available on Windows"
)


@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()


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 get_completions(self, words: List[str], cword: int) -> List[str]:
        """Simulate bash completion for a command line.

        Args:
            words: List of words on the command line (e.g., ["prog", "sub1", ""]).
            cword: Index of the word being completed (0-based).

        Returns:
            List of completions from COMPREPLY.
        """
        # Extract the completion function name from the script.
        # It's registered with: complete -F <func_name> <prog>
        match = re.search(r"complete -F (\S+)", self.completion_script)
        if not match:
            raise ValueError("Could not find completion function in script")
        func_name = match.group(1)

        # Build bash array syntax for COMP_WORDS.
        words_escaped = " ".join(f'"{w}"' for w in words)

        test_script = f"""
{self.completion_script}

# Set up completion environment.
COMP_WORDS=({words_escaped})
COMP_CWORD={cword}

# Call the completion function.
{func_name}

# Output completions, one per line.
printf '%s\\n' "${{COMPREPLY[@]}}"
"""

        proc = subprocess.Popen(
            ["bash", "-c", test_script],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = proc.communicate()

        if proc.returncode != 0:
            raise RuntimeError(
                f"Bash completion failed:\n{stderr.decode() if stderr else ''}"
            )

        # Parse output - one completion per line.
        output = stdout.decode().strip()
        if not output:
            return []
        return output.split("\n")

    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(backend: str) -> None:
    """Test basic bash completion script generation."""
    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_completion_uses_correct_prog_name(backend: str) -> None:
    """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.

        # 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)


def test_bash_functional_completion_simple(backend: str) -> None:
    """Test that bash completion actually works for a simple case."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion behavior")

    @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 completions at root level.
    completions = tester.get_completions(["prog", ""], 1)
    # Should have exactly these completions (flags for mode, verbose, help).
    expected_root = {"-h", "--help", "--mode", "--verbose", "--no-verbose"}
    assert set(completions) == expected_root, (
        f"Root completions mismatch.\nExpected: {expected_root}\nGot: {set(completions)}"
    )

    # Test completing after --mode shows only the choices.
    completions = tester.get_completions(["prog", "--mode", ""], 2)
    expected_choices = {"train", "eval", "test"}
    assert set(completions) == expected_choices, (
        f"Mode choices mismatch.\nExpected: {expected_choices}\nGot: {set(completions)}"
    )

    # Test partial completion of mode values filters correctly.
    completions = tester.get_completions(["prog", "--mode", "tr"], 2)
    assert set(completions) == {"train"}, (
        f"Partial completion should only show 'train', got: {completions}"
    )


def test_bash_functional_completion_with_subcommands(backend: str) -> None:
    """Test bash completion with subcommands."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion behavior")

    @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()
    tester = BashCompletionTester(completion_script)

    # Test root level shows only subcommands and help.
    completions = tester.get_completions(["prog", ""], 1)
    mnist_subcmd = next((c for c in completions if "mnist" in c.lower()), None)
    imagenet_subcmd = next((c for c in completions if "image" in c.lower()), None)
    assert mnist_subcmd is not None, f"mnist subcommand not found in {completions}"
    assert imagenet_subcmd is not None, (
        f"imagenet subcommand not found in {completions}"
    )

    # Root should only have help flags and subcommands - no other options.
    expected_root = {"-h", "--help", mnist_subcmd, imagenet_subcmd}
    assert set(completions) == expected_root, (
        f"Root completions mismatch.\nExpected: {expected_root}\nGot: {set(completions)}"
    )

    # Test that selecting mnist subcommand shows exactly its options.
    completions = tester.get_completions(["prog", mnist_subcmd, ""], 2)
    # Options have prefix like --dataset.binary.
    expected_mnist = {"-h", "--help", "--dataset.binary", "--dataset.no-binary"}
    assert set(completions) == expected_mnist, (
        f"Mnist completions mismatch.\nExpected: {expected_mnist}\nGot: {set(completions)}"
    )
    # Should not have imagenet's --subset option.
    assert not any("subset" in c for c in completions), (
        f"--subset should not appear in mnist context, got: {completions}"
    )

    # Test that selecting imagenet subcommand shows exactly its options.
    completions = tester.get_completions(["prog", imagenet_subcmd, ""], 2)
    expected_imagenet = {"-h", "--help", "--dataset.subset"}
    assert set(completions) == expected_imagenet, (
        f"ImageNet completions mismatch.\nExpected: {expected_imagenet}\nGot: {set(completions)}"
    )
    # Should not have mnist's --binary option.
    assert not any("binary" in c for c in completions), (
        f"--binary should not appear in imagenet context, got: {completions}"
    )


def test_bash_functional_completion_frontier_subcommands(backend: str) -> None:
    """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) -> None:
    """Test bash completion with CascadeSubcommandArgs marker.

    CascadeSubcommandArgs allows parent options to appear at the same level as subcommands,
    so users can specify them before selecting a subcommand.
    """
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion behavior")

    @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: SubcommandA | SubcommandB = dataclasses.field(
            default_factory=SubcommandA
        )

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

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

    # Test 1: At root level, should see subcommands, help, AND parent options together.
    # This is the key cascade behavior - options appear alongside subcommands.
    completions = tester.get_completions(["prog", ""], 1)
    subcmd_a = next((c for c in completions if "subcommand-a" in c.lower()), None)
    subcmd_b = next((c for c in completions if "subcommand-b" in c.lower()), None)
    assert subcmd_a is not None, f"subcommand-a not found in {completions}"
    assert subcmd_b is not None, f"subcommand-b not found in {completions}"

    # Root should have help, subcommands, AND the common option (cascade behavior).
    expected_root = {"-h", "--help", subcmd_a, subcmd_b, "--common", "-c"}
    assert set(completions) == expected_root, (
        f"Root completions mismatch.\nExpected: {expected_root}\nGot: {set(completions)}"
    )

    # Test 2: After selecting subcommand-a, should see exactly its options.
    completions = tester.get_completions(["prog", subcmd_a, ""], 2)
    expected_subcmd_a = {"-h", "--help", "--subcommand.arg-a"}
    assert set(completions) == expected_subcmd_a, (
        f"SubcommandA completions mismatch.\nExpected: {expected_subcmd_a}\nGot: {set(completions)}"
    )
    # Should not have subcommand-b's option.
    assert not any("arg-b" in c for c in completions), (
        f"--arg-b should not appear in subcommand-a context, got: {completions}"
    )

    # Test 3: After selecting subcommand-b, should see exactly its options.
    completions = tester.get_completions(["prog", subcmd_b, ""], 2)
    expected_subcmd_b = {"-h", "--help", "--subcommand.arg-b"}
    assert set(completions) == expected_subcmd_b, (
        f"SubcommandB completions mismatch.\nExpected: {expected_subcmd_b}\nGot: {set(completions)}"
    )
    # Should not have subcommand-a's option.
    assert not any("arg-a" in c for c in completions), (
        f"--arg-a should not appear in subcommand-b context, got: {completions}"
    )


def test_deeply_nested_subcommand_completion(backend: str) -> None:
    """Test completion for subcommands at depth >= 2.

    This tests that the completion context correctly navigates into nested
    subcommands. Previously, there was a bug where the code checked the root
    spec instead of the current spec when determining if a word is a subcommand.

    See: https://github.com/brentyi/tyro/pull/407
    """
    if backend != "tyro":
        pytest.skip("Testing tyro-specific completion behavior")

    # Create nested subcommand structure: root -> level1 -> level2.
    @dataclasses.dataclass
    class Level2A:
        level2_opt: int = 1

    @dataclasses.dataclass
    class Level2B:
        level2_opt: int = 2

    @dataclasses.dataclass
    class Level1A:
        nested: Level2A | Level2B

    @dataclasses.dataclass
    class Level1B:
        other: str = "x"

    @dataclasses.dataclass
    class Root:
        child: Level1A | Level1B

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

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

    # Test 1: At root level, should see exactly level1 subcommands and help.
    completions = tester.get_completions(["prog", ""], 1)
    level1a_name = next((c for c in completions if "level1-a" in c.lower()), None)
    level1b_name = next((c for c in completions if "level1-b" in c.lower()), None)
    assert level1a_name is not None, f"level1-a not found in {completions}"
    assert level1b_name is not None, f"level1-b not found in {completions}"

    expected_root = {"-h", "--help", level1a_name, level1b_name}
    assert set(completions) == expected_root, (
        f"Root completions mismatch.\nExpected: {expected_root}\nGot: {set(completions)}"
    )

    # Test 2: After selecting level1-a, should see exactly level2 subcommands and help.
    completions = tester.get_completions(["prog", level1a_name, ""], 2)
    level2a_name = next((c for c in completions if "level2-a" in c.lower()), None)
    level2b_name = next((c for c in completions if "level2-b" in c.lower()), None)
    assert level2a_name is not None, f"level2-a not found in {completions}"
    assert level2b_name is not None, f"level2-b not found in {completions}"

    expected_level1a = {"-h", "--help", level2a_name, level2b_name}
    assert set(completions) == expected_level1a, (
        f"Level1A completions mismatch.\nExpected: {expected_level1a}\nGot: {set(completions)}"
    )
    # Should not have level1b's --other option.
    assert "--other" not in completions, "--other should not appear in level1-a context"

    # Test 3: After selecting level1-a and level2-a, should see exactly level2's options.
    # This is the case that was broken before the fix - it would show level2 subcommands again.
    completions = tester.get_completions(["prog", level1a_name, level2a_name, ""], 3)
    # Options have full prefix like --child.nested.level2-opt.
    expected_level2a = {"-h", "--help", "--child.nested.level2-opt"}
    assert set(completions) == expected_level2a, (
        f"Level2A completions mismatch.\nExpected: {expected_level2a}\nGot: {set(completions)}"
    )
    # Should not have level2 subcommands anymore (the bug showed these).
    assert level2a_name not in completions and level2b_name not in completions, (
        f"Subcommands should not appear after selecting level2, got: {completions}"
    )

    # Test 4: Verify level1-b shows its own options, not level1-a's nested subcommands.
    completions = tester.get_completions(["prog", level1b_name, ""], 2)
    expected_level1b = {"-h", "--help", "--child.other"}
    assert set(completions) == expected_level1b, (
        f"Level1B completions mismatch.\nExpected: {expected_level1b}\nGot: {set(completions)}"
    )
    # Should not have level2 subcommands.
    assert level2a_name not in completions and level2b_name not in completions, (
        "Level2 subcommands should not appear in level1-b context"
    )


def test_path_completion_functional(backend: str) -> None:
    """Test that path arguments trigger file completion in bash."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific path completion")

    from pathlib import Path

    @dataclasses.dataclass
    class Config:
        input_file: Path
        name: str = "default"

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

    completion_script = target.getvalue()

    # Verify the script contains the marker handling logic.
    assert "__TYRO_COMPLETE_FILES__" in completion_script
    assert "compgen -f" in completion_script

    # Both Path and str arguments should be marked as path type in the spec.
    assert "'type': 'path'" in completion_script


def test_bash_path_completion_marker(backend: str) -> None:
    """Test that bash completion returns path marker for Path/str arguments."""
    if backend != "tyro":
        pytest.skip("Testing tyro-specific path completion")

    @dataclasses.dataclass
    class Config:
        input_file: pathlib.Path
        output: str = "out.txt"
        count: int = 5

    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 1: At root level, should see all options.
    completions = tester.get_completions(["prog", ""], 1)
    assert "--input-file" in completions
    assert "--output" in completions
    assert "--count" in completions

    # Test 2: After --input-file, bash should use file completion.
    # The Python script returns __TYRO_COMPLETE_FILES__, which bash
    # converts to actual file completions via compgen -f.
    completions = tester.get_completions(["prog", "--input-file", ""], 2)
    # Should have file completions from cwd (the tyro repo has many files).
    assert len(completions) > 0, "Expected file completions after --input-file"
    # Verify we get actual files, not flags or other completions.
    assert not any(c.startswith("-") for c in completions), (
        f"Expected file completions, not flags: {completions}"
    )

    # Test 3: After --count (int), should return no completions.
    completions = tester.get_completions(["prog", "--count", ""], 2)
    assert completions == []
