"""Tests for tacoshop.cli."""

from __future__ import annotations

import pytest

from tacoshop.cli import ParseResult, parse_args
from tacoshop.config import Config, FlagVar


def _make_config(**flag_defs: dict) -> Config:
    """Helper: create a Config with flag vars.

    Usage: _make_config(TRITON={"flag": "--triton", "value": "1"},
                        LOGS={"flag": "--logs"})
    """
    config = Config()
    for name, defn in flag_defs.items():
        fv = FlagVar(
            name=name,
            flag=defn["flag"],
            value=defn.get("value"),
            append=defn.get("append", False),
            separator=defn.get("separator", ""),
        )
        config.flag_vars[name] = fv
        config.flag_lookup[fv.flag] = fv
    return config


EMPTY_CONFIG = Config()

SAMPLE_CONFIG = _make_config(
    TRITON_ALWAYS_COMPILE={"flag": "--triton", "value": "1"},
    TORCH_LOGS={"flag": "--logs"},
)


class TestBasicParsing:
    def test_empty_argv(self):
        result = parse_args([], EMPTY_CONFIG)
        assert result.command == []
        assert result.subcommand is None

    def test_simple_command(self):
        result = parse_args(["python", "test.py"], EMPTY_CONFIG)
        assert result.command == ["python", "test.py"]

    def test_command_with_its_own_flags(self):
        result = parse_args(
            ["python", "test.py", "--verbose", "--help"], EMPTY_CONFIG
        )
        assert result.command == ["python", "test.py", "--verbose", "--help"]
        assert not result.verbose
        assert not result.help


class TestBuiltinFlags:
    def test_help(self):
        result = parse_args(["--help"], EMPTY_CONFIG)
        assert result.help is True

    def test_help_short(self):
        result = parse_args(["-h"], EMPTY_CONFIG)
        assert result.help is True

    def test_verbose(self):
        result = parse_args(["--verbose", "python"], EMPTY_CONFIG)
        assert result.verbose is True
        assert result.command == ["python"]

    def test_verbose_short(self):
        result = parse_args(["-v", "python"], EMPTY_CONFIG)
        assert result.verbose is True

    def test_dry_run(self):
        result = parse_args(["--dry-run", "python"], EMPTY_CONFIG)
        assert result.dry_run is True
        assert result.command == ["python"]

    def test_version(self):
        result = parse_args(["--version"], EMPTY_CONFIG)
        assert result.version is True

    def test_multiple_builtins(self):
        result = parse_args(["--verbose", "--dry-run", "python"], EMPTY_CONFIG)
        assert result.verbose is True
        assert result.dry_run is True
        assert result.command == ["python"]


class TestConfigFlags:
    def test_preset_flag(self):
        result = parse_args(["--triton", "python", "test.py"], SAMPLE_CONFIG)
        assert result.activated_flags == {"TRITON_ALWAYS_COMPILE": "1"}
        assert result.command == ["python", "test.py"]

    def test_value_required_flag(self):
        result = parse_args(
            ["--logs", "+dynamo", "python", "test.py"], SAMPLE_CONFIG
        )
        assert result.activated_flags == {"TORCH_LOGS": "+dynamo"}
        assert result.command == ["python", "test.py"]

    def test_multiple_flags(self):
        result = parse_args(
            ["--triton", "--logs", "+dynamo", "python"], SAMPLE_CONFIG
        )
        assert result.activated_flags == {
            "TRITON_ALWAYS_COMPILE": "1",
            "TORCH_LOGS": "+dynamo",
        }
        assert result.command == ["python"]

    def test_value_flag_consumes_next_token(self):
        # Even if next token looks like a flag, it's consumed as value
        result = parse_args(["--logs", "--verbose", "python"], SAMPLE_CONFIG)
        assert result.activated_flags == {"TORCH_LOGS": "--verbose"}
        assert result.command == ["python"]

    def test_value_flag_missing_value(self):
        with pytest.raises(SystemExit) as exc_info:
            parse_args(["--logs"], SAMPLE_CONFIG)
        assert exc_info.value.code == 2

    def test_unknown_flag_starts_command(self):
        result = parse_args(["--unknown", "python"], SAMPLE_CONFIG)
        assert result.command == ["--unknown", "python"]
        assert result.activated_flags == {}


class TestExplicitSeparator:
    def test_double_dash_separator(self):
        result = parse_args(
            ["--triton", "--", "python", "test.py"], SAMPLE_CONFIG
        )
        assert result.activated_flags == {"TRITON_ALWAYS_COMPILE": "1"}
        assert result.command == ["python", "test.py"]

    def test_separator_bypasses_subcommand(self):
        result = parse_args(["--", "show"], EMPTY_CONFIG)
        assert result.subcommand is None
        assert result.command == ["show"]

    def test_separator_with_nothing_after(self):
        result = parse_args(["--"], EMPTY_CONFIG)
        assert result.command == []


class TestSubcommands:
    def test_trace_subcommand(self):
        result = parse_args(["trace"], EMPTY_CONFIG)
        assert result.subcommand == "trace"
        assert result.subcommand_args == []

    def test_init_subcommand(self):
        result = parse_args(["init"], EMPTY_CONFIG)
        assert result.subcommand == "init"
        assert result.subcommand_args == []

    def test_help_subcommand(self):
        result = parse_args(["help"], EMPTY_CONFIG)
        assert result.subcommand == "help"
        assert result.subcommand_args == []

    def test_print_subcommand(self):
        result = parse_args(["print", "python", "train.py"], EMPTY_CONFIG)
        assert result.subcommand == "print"
        assert result.subcommand_args == ["python", "train.py"]

    def test_print_with_flags(self):
        result = parse_args(
            ["--triton", "print", "python", "train.py"], SAMPLE_CONFIG
        )
        assert result.activated_flags == {"TRITON_ALWAYS_COMPILE": "1"}
        assert result.subcommand == "print"
        assert result.subcommand_args == ["python", "train.py"]

    def test_dumpenv_subcommand(self):
        result = parse_args(["dumpenv"], EMPTY_CONFIG)
        assert result.subcommand == "dumpenv"

    def test_dumpenv_with_args(self):
        result = parse_args(["dumpenv", "-o", "env.sh"], EMPTY_CONFIG)
        assert result.subcommand == "dumpenv"
        assert result.subcommand_args == ["-o", "env.sh"]

    def test_flags_before_subcommand(self):
        result = parse_args(["--triton", "dumpenv"], SAMPLE_CONFIG)
        assert result.activated_flags == {"TRITON_ALWAYS_COMPILE": "1"}
        assert result.subcommand == "dumpenv"

    def test_subcommand_after_separator_is_command(self):
        result = parse_args(["--", "trace"], SAMPLE_CONFIG)
        assert result.subcommand is None
        assert result.command == ["trace"]


class TestEdgeCases:
    def test_builtin_flags_after_command_not_consumed(self):
        result = parse_args(
            ["python", "--help", "--verbose"], EMPTY_CONFIG
        )
        assert result.command == ["python", "--help", "--verbose"]
        assert not result.help
        assert not result.verbose

    def test_config_flag_after_command_not_consumed(self):
        result = parse_args(
            ["python", "--triton", "test.py"], SAMPLE_CONFIG
        )
        assert result.command == ["python", "--triton", "test.py"]
        assert result.activated_flags == {}

    def test_flags_and_command_interleaved(self):
        result = parse_args(
            ["--verbose", "--triton", "python", "-c", "pass"], SAMPLE_CONFIG
        )
        assert result.verbose is True
        assert result.activated_flags == {"TRITON_ALWAYS_COMPILE": "1"}
        assert result.command == ["python", "-c", "pass"]
