import asyncio
import logging
import sys
from argparse import ArgumentParser, Namespace, _SubParsersAction
from contextlib import ExitStack
from pathlib import Path
from tempfile import TemporaryDirectory
from threading import Event as ThreadEvent
from time import perf_counter
from types import SimpleNamespace
from unittest import IsolatedAsyncioTestCase, TestCase
from unittest.mock import AsyncMock, MagicMock, call, patch

from avalan.cli import CommandAbortException
from avalan.cli.__main__ import (
    CLI,
    _consume_task_input_field_args,
    _task_run_json_stdout,
)


def _collect_progs(parser: ArgumentParser) -> list[str]:
    progs: list[str] = []
    stack = [parser]
    while stack:
        p = stack.pop()
        progs.append(p.prog)
        for action in p._actions:
            if isinstance(action, _SubParsersAction):
                for sub in action.choices.values():
                    stack.append(sub)
    return progs


def _find_parser(parser: ArgumentParser, prog: str) -> ArgumentParser:
    stack = [parser]
    while stack:
        candidate = stack.pop()
        if candidate.prog == prog:
            return candidate
        for action in candidate._actions:
            if isinstance(action, _SubParsersAction):
                stack.extend(action.choices.values())
    raise AssertionError(f"Parser {prog!r} was not found.")


def _find_parser_with_suffix(
    parser: ArgumentParser,
    prog_suffix: str,
) -> ArgumentParser:
    stack = [parser]
    while stack:
        candidate = stack.pop()
        if candidate.prog.endswith(prog_suffix):
            return candidate
        for action in candidate._actions:
            if isinstance(action, _SubParsersAction):
                stack.extend(action.choices.values())
    raise AssertionError(f"Parser ending with {prog_suffix!r} was not found.")


class CliInitTestCase(TestCase):
    def test_constructor_creates_parser_and_sets_help_full(self):
        logger = MagicMock()
        with (
            patch.object(sys, "argv", ["prog"]),
            patch(
                "avalan.cli.__main__.getlocale",
                return_value=("en_US", "UTF-8"),
            ),
            patch(
                "avalan.cli.__main__.TransformerModel.get_default_device",
                return_value="cpu",
            ),
            patch.object(
                CLI, "_create_parser", wraps=CLI._create_parser
            ) as cp,
        ):
            cli = CLI(logger)

        locales_path = str(
            Path(sys.modules[CLI.__module__].__file__).resolve().parents[3]
            / "locale"
        )
        cp.assert_called_once_with(
            "cpu",
            "~/.cache/huggingface/hub",
            locales_path,
            "en_US",
        )
        args = cli._parser.parse_args(["--help-full"])
        self.assertTrue(hasattr(args, "help_full"))


class CliParallelOptionTestCase(TestCase):
    def test_default_flow_parallel_count_uses_cpu_count(self) -> None:
        with patch("avalan.cli.__main__.cpu_count", return_value=12):
            self.assertEqual(CLI._default_flow_parallel_count(), 12)

        with patch("avalan.cli.__main__.cpu_count", return_value=None):
            self.assertEqual(CLI._default_flow_parallel_count(), 1)

    def test_parallel_argument(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)
        args = cli._parser.parse_args(["--parallel", "colwise"])
        self.assertEqual(args.parallel, "colwise")

    def test_parallel_count_argument(self) -> None:
        logger = MagicMock()
        with (
            patch.object(sys, "argv", ["prog"]),
            patch.object(CLI, "_default_parallel_count", return_value=3),
        ):
            cli = CLI(logger)
        args = cli._parser.parse_args(
            [
                "--parallel",
                "colwise",
                "--parallel-count",
                "5",
            ]
        )
        self.assertEqual(args.parallel_count, 5)

    def test_flow_run_parallel_argument(self) -> None:
        logger = MagicMock()
        with (
            patch.object(sys, "argv", ["prog"]),
            patch.object(CLI, "_default_flow_parallel_count", return_value=8),
        ):
            cli = CLI(logger)

        default_args = cli._parser.parse_args(["flow", "run", "flow.toml"])
        explicit_args = cli._parser.parse_args(
            ["flow", "run", "flow.toml", "--flow-parallel", "4"]
        )

        self.assertIsNone(getattr(default_args, "parallel", None))
        self.assertEqual(default_args.flow_parallel, 8)
        self.assertEqual(explicit_args.flow_parallel, 4)

    def test_flow_run_parallel_argument_must_be_positive(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        for value in ("0", "many"):
            with self.subTest(value=value), self.assertRaises(SystemExit):
                cli._parser.parse_args(
                    ["flow", "run", "flow.toml", "--flow-parallel", value]
                )


class CliTaskOptionTestCase(TestCase):
    def test_task_validate_argument(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "task",
                "validate",
                "tasks/example.task.toml",
            ]
        )

        self.assertEqual(args.command, "task")
        self.assertEqual(args.task_command, "validate")
        self.assertEqual(args.definition, "tasks/example.task.toml")

    def test_task_pgsql_migrate_argument(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "task",
                "pgsql",
                "migrate",
                "--dsn",
                "postgresql://localhost/tasks",
                "head",
            ]
        )

        self.assertEqual(args.command, "task")
        self.assertEqual(args.task_command, "pgsql")
        self.assertEqual(args.task_pgsql_command, "migrate")
        self.assertEqual(args.dsn, "postgresql://localhost/tasks")
        self.assertEqual(args.migration_revision, "head")

    def test_task_shell_subcommand_arguments(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        cases = [
            (
                ["task", "run", "tasks/example.task.toml"],
                {
                    "task_command": "run",
                    "definition": "tasks/example.task.toml",
                },
            ),
            (
                ["task", "enqueue", "tasks/example.task.toml"],
                {
                    "task_command": "enqueue",
                    "definition": "tasks/example.task.toml",
                },
            ),
            (
                ["task", "inspect", "run-123"],
                {"task_command": "inspect", "run_id": "run-123"},
            ),
            (
                [
                    "task",
                    "usage",
                    "run-123",
                    "--attempt-id",
                    "attempt-1",
                    "--source",
                    "exact",
                ],
                {
                    "task_command": "usage",
                    "run_id": "run-123",
                    "attempt_id": "attempt-1",
                    "source": "exact",
                },
            ),
            (
                ["task", "output", "run-123"],
                {"task_command": "output", "run_id": "run-123"},
            ),
            (
                ["task", "events", "run-123"],
                {"task_command": "events", "run_id": "run-123"},
            ),
            (
                ["task", "artifacts", "run-123"],
                {"task_command": "artifacts", "run_id": "run-123"},
            ),
            (
                ["task", "worker", "--queue", "documents"],
                {"task_command": "worker", "queue": "documents"},
            ),
            (
                [
                    "task",
                    "retention-sweep",
                    "--store-dsn",
                    "postgresql://localhost/tasks",
                    "--purpose",
                    "input",
                    "--purpose",
                    "output",
                    "--limit",
                    "25",
                ],
                {
                    "task_command": "retention-sweep",
                    "store_dsn": "postgresql://localhost/tasks",
                    "purpose": ["input", "output"],
                    "limit": 25,
                },
            ),
        ]

        for argv, expected in cases:
            with self.subTest(argv=argv):
                args = cli._parser.parse_args(argv)

            self.assertEqual(args.command, "task")
            for name, value in expected.items():
                self.assertEqual(getattr(args, name), value)

    def test_flow_run_arguments(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "flow",
                "run",
                "flows/report.toml",
                "--input-json",
                '{"answer":"ok"}',
                "--file",
                "document=sample.pdf",
                "--file-mime",
                "document=application/pdf",
                "--json",
                "--output",
                "result.json",
            ]
        )

        self.assertEqual(args.command, "flow")
        self.assertEqual(args.flow_command, "run")
        self.assertEqual(args.flow, "flows/report.toml")
        self.assertEqual(args.task_input_json, '{"answer":"ok"}')
        self.assertEqual(args.task_files, ["document=sample.pdf"])
        self.assertEqual(
            args.task_file_mime_types,
            ["document=application/pdf"],
        )
        self.assertTrue(args.task_run_json)
        self.assertEqual(args.task_output_path, "result.json")

    def test_flow_non_executing_arguments(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        cases = [
            (
                [
                    "flow",
                    "validate",
                    "flows/report.toml",
                    "--json",
                ],
                {
                    "flow_command": "validate",
                    "flow": "flows/report.toml",
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "inspect",
                    "run-123",
                    "--store-dsn",
                    "postgresql://localhost/tasks",
                    "--store-schema",
                    "workflow",
                    "--after-sequence",
                    "4",
                    "--json",
                ],
                {
                    "flow_command": "inspect",
                    "run_id": "run-123",
                    "store_dsn": "postgresql://localhost/tasks",
                    "store_schema": "workflow",
                    "after_sequence": 4,
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "trace",
                    "run-123",
                    "--store-dsn",
                    "postgresql://localhost/tasks",
                    "--after-sequence",
                    "5",
                    "--json",
                ],
                {
                    "flow_command": "trace",
                    "run_id": "run-123",
                    "store_dsn": "postgresql://localhost/tasks",
                    "after_sequence": 5,
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "cancel",
                    "run-123",
                    "--store-dsn",
                    "postgresql://localhost/tasks",
                    "--json",
                ],
                {
                    "flow_command": "cancel",
                    "run_id": "run-123",
                    "store_dsn": "postgresql://localhost/tasks",
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "resume",
                    "flows/report.toml",
                    "run-123",
                    "--decision-json",
                    '{"review":{"decision":"approved"}}',
                    "--store-dsn",
                    "postgresql://localhost/tasks",
                    "--store-schema",
                    "workflow",
                    "--json",
                ],
                {
                    "flow_command": "resume",
                    "flow": "flows/report.toml",
                    "run_id": "run-123",
                    "decision_json": '{"review":{"decision":"approved"}}',
                    "store_dsn": "postgresql://localhost/tasks",
                    "store_schema": "workflow",
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "compile",
                    "flows/report.toml",
                    "--output",
                    "build/report.flow.toml",
                    "--json",
                ],
                {
                    "flow_command": "compile",
                    "flow": "flows/report.toml",
                    "output": "build/report.flow.toml",
                    "check": False,
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "compile",
                    "flows/report.toml",
                    "--check",
                ],
                {
                    "flow_command": "compile",
                    "flow": "flows/report.toml",
                    "output": None,
                    "check": True,
                },
            ),
            (
                [
                    "flow",
                    "graph",
                    "inspect",
                    "flows/report.toml",
                    "--json",
                ],
                {
                    "flow_command": "graph",
                    "flow_graph_command": "inspect",
                    "flow": "flows/report.toml",
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "mermaid",
                    "parse",
                    "topology.mmd",
                    "--mode",
                    "executable",
                    "--json",
                ],
                {
                    "flow_command": "mermaid",
                    "flow_mermaid_command": "parse",
                    "diagram": "topology.mmd",
                    "mode": "executable",
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "mermaid",
                    "render",
                    "topology.mmd",
                    "--mode",
                    "presentation",
                    "--json",
                ],
                {
                    "flow_command": "mermaid",
                    "flow_mermaid_command": "render",
                    "diagram": "topology.mmd",
                    "mode": "presentation",
                    "flow_json": True,
                },
            ),
            (
                [
                    "flow",
                    "mermaid",
                    "compare",
                    "topology.mmd",
                    "flow.toml",
                    "--mode",
                    "presentation",
                ],
                {
                    "flow_command": "mermaid",
                    "flow_mermaid_command": "compare",
                    "diagram": "topology.mmd",
                    "flow": "flow.toml",
                    "mode": "presentation",
                    "flow_json": False,
                },
            ),
            (
                [
                    "flow",
                    "mermaid",
                    "skeleton",
                    "topology.mmd",
                    "--mode",
                    "presentation",
                    "--name",
                    "topology",
                    "--flow-version",
                    "1",
                ],
                {
                    "flow_command": "mermaid",
                    "flow_mermaid_command": "skeleton",
                    "diagram": "topology.mmd",
                    "mode": "presentation",
                    "name": "topology",
                    "version": "1",
                },
            ),
        ]

        for argv, expected in cases:
            with self.subTest(argv=argv):
                args = cli._parser.parse_args(argv)

            self.assertEqual(args.command, "flow")
            for name, value in expected.items():
                self.assertEqual(getattr(args, name), value)

    def test_flow_run_consumes_dynamic_input_fields(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "flow",
                "run",
                "flow.toml",
                "--input-name",
                "Ada",
            ]
        )

        self.assertEqual(args.command, "flow")
        self.assertEqual(args.flow_command, "run")
        self.assertEqual(args.task_input_fields, ("name=Ada",))

    def test_flow_non_run_does_not_consume_dynamic_input_fields(self) -> None:
        namespace = Namespace(command="flow", flow_command="inspect")

        self.assertFalse(
            _consume_task_input_field_args(
                namespace,
                ["--input-name", "Ada"],
            )
        )


class CliModelRunOptionTestCase(TestCase):
    def test_backend_help_includes_ds4_choice(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        help_text = cli._parser.format_help()

        self.assertIn("{transformers,mlx,vllm,ds4}", help_text)

    def test_input_file_argument(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "model",
                "run",
                "model-id",
                "--input-file",
                "doc-1.pdf",
                "--input-file",
                "doc-2.pdf",
            ]
        )

        self.assertEqual(args.input_file, ["doc-1.pdf", "doc-2.pdf"])

    def test_ds4_backend_options(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "model",
                "run",
                "model-id",
                "--backend",
                "ds4",
                "--ds4-ctx",
                "4096",
                "--ds4-native-backend",
                "metal",
                "--ds4-mtp",
                "draft.gguf",
                "--ds4-mtp-draft",
                "2",
                "--ds4-mtp-margin",
                "0.25",
                "--ds4-warm-weights",
                "--ds4-quality",
                "--with-ds4-native-log",
            ]
        )

        self.assertEqual(args.backend, "ds4")
        self.assertEqual(args.ds4_ctx, 4096)
        self.assertEqual(args.ds4_native_backend, "metal")
        self.assertEqual(args.ds4_mtp, "draft.gguf")
        self.assertEqual(args.ds4_mtp_draft, 2)
        self.assertEqual(args.ds4_mtp_margin, 0.25)
        self.assertTrue(args.ds4_warm_weights)
        self.assertTrue(args.ds4_quality)
        self.assertTrue(args.ds4_native_log)

    def test_agent_run_ds4_backend_options(self) -> None:
        logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        args = cli._parser.parse_args(
            [
                "agent",
                "run",
                "--backend",
                "ds4",
                "--engine-uri",
                "ai://local/./ds4flash.gguf",
                "--ds4-ctx",
                "4096",
                "--ds4-native-backend",
                "metal",
                "--ds4-mtp",
                "draft.gguf",
                "--ds4-mtp-draft",
                "2",
                "--ds4-mtp-margin",
                "0.25",
                "--ds4-warm-weights",
                "--ds4-quality",
                "--no-ds4-native-log",
            ]
        )

        self.assertEqual(args.backend, "ds4")
        self.assertEqual(args.engine_uri, "ai://local/./ds4flash.gguf")
        self.assertEqual(args.ds4_ctx, 4096)
        self.assertEqual(args.ds4_native_backend, "metal")
        self.assertEqual(args.ds4_mtp, "draft.gguf")
        self.assertEqual(args.ds4_mtp_draft, 2)
        self.assertEqual(args.ds4_mtp_margin, 0.25)
        self.assertTrue(args.ds4_warm_weights)
        self.assertTrue(args.ds4_quality)
        self.assertFalse(args.ds4_native_log)

    def test_documented_ds4_options_are_in_model_run_help(self) -> None:
        logger = MagicMock()
        docs = (
            Path(__file__).resolve().parents[2] / "docs" / "DS4.md"
        ).read_text(encoding="utf-8")
        with patch.object(sys, "argv", ["prog"]):
            cli = CLI(logger)

        model_run_help = _find_parser_with_suffix(
            cli._parser, " model run"
        ).format_help()
        agent_run_help = _find_parser_with_suffix(
            cli._parser, " agent run"
        ).format_help()

        for option in (
            "--ds4-ctx",
            "--ds4-native-backend",
            "--ds4-mtp",
            "--ds4-mtp-draft",
            "--ds4-mtp-margin",
            "--ds4-warm-weights",
            "--ds4-quality",
            "--with-ds4-native-log",
            "--ds4-native-log",
            "--no-ds4-native-log",
            "--reasoning-effort",
        ):
            self.assertIn(option, docs)
            self.assertIn(option, model_run_help)
            self.assertIn(option, agent_run_help)

        self.assertIn("--no-reasoning", docs)
        self.assertIn("--no-reasoning", model_run_help)

    def test_ds4_docs_call_out_limitations(self) -> None:
        docs = (
            Path(__file__).resolve().parents[2] / "docs" / "DS4.md"
        ).read_text(encoding="utf-8")
        normalized_docs = docs.lower()

        self.assertIn("not a generic gguf", normalized_docs)
        self.assertIn("native ds4 tool calls use dsml", normalized_docs)
        self.assertIn("renders tool schemas", normalized_docs)
        self.assertIn("parses completed dsml tool blocks", normalized_docs)
        self.assertIn("streams\n  argument deltas", normalized_docs)
        self.assertIn("exact raw dsml replay metadata", normalized_docs)
        self.assertIn("cpu backend exists for diagnostics", normalized_docs)
        self.assertNotIn("generic gguf models are supported", normalized_docs)
        self.assertNotIn(
            "native ds4 tool calls are implemented", normalized_docs
        )
        self.assertNotIn(
            "argument deltas are still in progress", normalized_docs
        )


class CliAgentPortTestCase(TestCase):
    def setUp(self) -> None:
        self.logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            self.cli = CLI(self.logger)

    def test_agent_serve_default_port(self) -> None:
        args = self.cli._parser.parse_args(["agent", "serve"])
        self.assertEqual(args.port, 9001)

    def test_agent_proxy_default_port(self) -> None:
        args = self.cli._parser.parse_args(["agent", "proxy"])
        self.assertEqual(args.port, 9001)


class CliAgentCorsTestCase(TestCase):
    def setUp(self) -> None:
        self.logger = MagicMock()
        with patch.object(sys, "argv", ["prog"]):
            self.cli = CLI(self.logger)

    def test_agent_serve_cors_options(self) -> None:
        args = self.cli._parser.parse_args(
            [
                "agent",
                "serve",
                "--cors-origin",
                "https://a",
                "--cors-origin",
                "https://b",
                "--cors-origin-regex",
                "^https://example.com$",
                "--cors-method",
                "GET",
                "--cors-method",
                "POST",
                "--cors-header",
                "X-Test",
                "--cors-header",
                "X-Other",
                "--cors-credentials",
            ]
        )
        self.assertEqual(args.cors_origin, ["https://a", "https://b"])
        self.assertEqual(args.cors_origin_regex, "^https://example.com$")
        self.assertEqual(args.cors_method, ["GET", "POST"])
        self.assertEqual(args.cors_header, ["X-Test", "X-Other"])
        self.assertTrue(args.cors_credentials)

    def test_agent_proxy_cors_options(self) -> None:
        args = self.cli._parser.parse_args(
            [
                "agent",
                "proxy",
                "--engine-uri",
                "uri",
                "--cors-origin",
                "https://a",
                "--cors-method",
                "GET",
            ]
        )
        self.assertEqual(args.cors_origin, ["https://a"])
        self.assertEqual(args.cors_method, ["GET"])


class CliCallTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        from logging import getLogger

        self.logger = getLogger("cli-test")
        with patch.object(sys, "argv", ["prog"]):
            self.cli = CLI(self.logger)
        self.translator = SimpleNamespace(
            gettext=lambda s: s, ngettext=lambda s, p, n: s if n == 1 else p
        )

    async def test_call_runs_main(self):
        console = MagicMock(export_text=lambda: "")
        with (
            patch.object(sys, "argv", ["prog"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=console),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch.object(CLI, "_help") as help_mock,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        main_mock.assert_awaited_once()
        help_mock.assert_not_called()

    async def test_call_raises_when_huggingface_hub_is_missing(self):
        with (
            patch.object(sys, "argv", ["prog"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch(
                "avalan.cli.__main__._huggingface_hub_class",
                side_effect=ModuleNotFoundError("huggingface_hub"),
            ),
        ):
            with self.assertRaises(ModuleNotFoundError):
                await self.cli()

    async def test_call_help_full_outputs_all_commands(self):
        records_console = None

        def create_console(*args, **kwargs):
            nonlocal records_console
            from rich.console import Console

            records_console = Console(record=True)
            return records_console

        with (
            patch.object(sys, "argv", ["prog", "--help-full"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", side_effect=create_console),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
            patch.object(CLI, "_help", wraps=self.cli._help) as help_mock,
        ):
            await self.cli()

        main_mock.assert_not_called()
        self.assertEqual(
            help_mock.call_count, len(_collect_progs(self.cli._parser))
        )
        output = records_console.export_text()
        for prog in _collect_progs(self.cli._parser):
            self.assertIn(prog, output)


class CliModuleUtilitiesTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        from logging import getLogger

        self.logger = getLogger("cli-test")
        with patch.object(sys, "argv", ["prog"]):
            self.cli = CLI(self.logger)
        self.translator = SimpleNamespace(
            gettext=lambda s: s, ngettext=lambda s, p, n: s if n == 1 else p
        )

    def test_huggingface_hub_class_loads_from_module(self) -> None:
        hub_class = type("HubClass", (), {})
        module = SimpleNamespace(HuggingfaceHub=hub_class)

        with patch("avalan.cli.__main__.import_module", return_value=module):
            from avalan.cli.__main__ import _huggingface_hub_class

            self.assertIs(_huggingface_hub_class(), hub_class)

    async def test_call_version_outputs_version_and_exits(self):
        with (
            patch.object(sys, "argv", ["prog", "--version"]),
            patch("builtins.print") as print_patch,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
            patch.object(CLI, "_help") as help_mock,
        ):
            await self.cli()

        print_patch.assert_called_once_with(
            f"{self.cli._name} {self.cli._version}"
        )
        main_mock.assert_not_called()
        help_mock.assert_not_called()

    async def test_call_parallel_invokes_torchrun(self):
        with (
            patch.object(sys, "argv", ["prog", "--parallel", "colwise"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch("avalan.cli.__main__.run") as run_patch,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        run_patch.assert_called_once()
        main_mock.assert_not_called()

    async def test_call_parallel_child_destroys_group(self):
        with (
            patch.dict("os.environ", {"LOCAL_RANK": "0"}, clear=False),
            patch.object(sys, "argv", ["prog", "--parallel", "colwise"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch("avalan.cli.__main__.destroy_process_group") as destroy,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        main_mock.assert_awaited_once()
        destroy.assert_called_once()

    async def test_call_parallel_child_ignores_destroy_assertion_error(self):
        with (
            patch.dict("os.environ", {"LOCAL_RANK": "0"}, clear=False),
            patch.object(sys, "argv", ["prog", "--parallel", "colwise"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch(
                "avalan.cli.__main__.destroy_process_group",
                side_effect=AssertionError,
            ) as destroy,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        main_mock.assert_awaited_once()
        destroy.assert_called_once()

    async def test_call_parallel_child_keeps_non_cuda_device(self):
        with (
            patch.dict("os.environ", {"LOCAL_RANK": "1"}, clear=False),
            patch.object(
                sys,
                "argv",
                ["prog", "--parallel", "colwise", "--device", "cpu"],
            ),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch("avalan.cli.__main__.destroy_process_group"),
            patch("avalan.cli.__main__.set_device") as set_device_mock,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        main_mock.assert_awaited_once()
        self.assertEqual(main_mock.await_args.args[0].device, "cpu")
        set_device_mock.assert_not_called()

    async def test_call_parallel_child_sets_rank_on_cuda_device(self):
        rank = 2
        with (
            patch.dict("os.environ", {"LOCAL_RANK": str(rank)}, clear=False),
            patch.object(
                sys,
                "argv",
                ["prog", "--parallel", "colwise", "--device", "cuda"],
            ),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch("avalan.cli.__main__.destroy_process_group"),
            patch("avalan.cli.__main__.set_device") as set_device_mock,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        main_mock.assert_awaited_once()
        self.assertEqual(main_mock.await_args.args[0].device, f"cuda:{rank}")
        set_device_mock.assert_called_once_with(rank)

    async def test_call_parallel_child_leaves_explicit_cuda_device(self):
        with (
            patch.dict("os.environ", {"LOCAL_RANK": "0"}, clear=False),
            patch.object(
                sys,
                "argv",
                ["prog", "--parallel", "colwise", "--device", "cuda:1"],
            ),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch("avalan.cli.__main__.destroy_process_group"),
            patch("avalan.cli.__main__.set_device") as set_device_mock,
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        main_mock.assert_awaited_once()
        self.assertEqual(main_mock.await_args.args[0].device, "cuda:1")
        set_device_mock.assert_not_called()

    async def test_call_adds_mask_spec_embed_filter(self):
        hf_logger = MagicMock()
        added_filters: list = []

        def capture_filter(f: logging.Filter) -> None:
            added_filters.append(f)

        hf_logger.addFilter.side_effect = capture_filter

        with (
            patch.object(sys, "argv", ["prog"]),
            patch(
                "avalan.cli.__main__.translation", return_value=self.translator
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch(
                "avalan.cli.__main__.hf_logging.get_logger",
                return_value=hf_logger,
            ),
            patch("avalan.cli.__main__.find_spec", return_value=None),
            patch("avalan.cli.__main__.logger_replace"),
            patch("avalan.cli.__main__.filterwarnings"),
            patch("avalan.cli.__main__.has_input", return_value=True),
            patch("avalan.cli.__main__.Confirm.ask", return_value=False),
        ):
            await self.cli()

        self.assertEqual(len(added_filters), 1)
        filt = added_filters[0]
        ok_record = logging.LogRecord(
            name="t",
            level=logging.INFO,
            pathname=__file__,
            lineno=1,
            msg="ok",
            args=(),
            exc_info=None,
        )
        bad_record = logging.LogRecord(
            name="t",
            level=logging.INFO,
            pathname=__file__,
            lineno=1,
            msg="wav2vec2.masked_spec_embed warning",
            args=(),
            exc_info=None,
        )
        self.assertTrue(filt.filter(ok_record))
        self.assertFalse(filt.filter(bad_record))


class CliLazyUtilityTestCase(IsolatedAsyncioTestCase):
    def test_anonymous_hub_records_cache_dir_and_rejects_user_calls(
        self,
    ) -> None:
        cli_module = sys.modules[CLI.__module__]
        hub = cli_module._AnonymousHub("cache")

        self.assertEqual(hub.cache_dir, "cache")
        with self.assertRaises(NotImplementedError):
            hub.login()
        with self.assertRaises(NotImplementedError):
            hub.user()

    def test_fancy_theme_loads_real_theme_lazily(self) -> None:
        cli_module = sys.modules[CLI.__module__]

        class LoadedTheme:
            def __init__(self, *args, **kwargs):
                self.args = args
                self.kwargs = kwargs

        module = SimpleNamespace(FancyTheme=LoadedTheme)
        with patch.object(cli_module, "import_module", return_value=module):
            result = cli_module.FancyTheme("gettext", key="value")

        self.assertIsInstance(result, LoadedTheme)
        self.assertEqual(result.args, ("gettext",))
        self.assertEqual(result.kwargs, {"key": "value"})

    def test_load_command_imports_target_function(self) -> None:
        cli_module = sys.modules[CLI.__module__]
        command = MagicMock(return_value="ok")
        module = SimpleNamespace(run=command)

        with patch.object(cli_module, "import_module", return_value=module):
            result = cli_module._load_command("module.path", "run")(
                "arg", key="value"
            )

        command.assert_called_once_with("arg", key="value")
        self.assertEqual(result, "ok")

    async def test_async_lazy_command_proxies_forward_calls(self) -> None:
        cli_module = sys.modules[CLI.__module__]

        async def command(*args, **kwargs):
            return args, kwargs

        commands = [
            (
                "agent_message_search",
                "avalan.cli.commands.agent",
                "agent_message_search",
            ),
            ("agent_run", "avalan.cli.commands.agent", "agent_run"),
            ("agent_serve", "avalan.cli.commands.agent", "agent_serve"),
            ("agent_proxy", "avalan.cli.commands.agent", "agent_proxy"),
            ("agent_init", "avalan.cli.commands.agent", "agent_init"),
            (
                "memory_document_index",
                "avalan.cli.commands.memory",
                "memory_document_index",
            ),
            (
                "memory_search",
                "avalan.cli.commands.memory",
                "memory_search",
            ),
            (
                "memory_embeddings",
                "avalan.cli.commands.memory",
                "memory_embeddings",
            ),
            ("model_run", "avalan.cli.commands.model", "model_run"),
            ("model_search", "avalan.cli.commands.model", "model_search"),
            ("deploy_run", "avalan.cli.commands.deploy", "deploy_run"),
            ("tokenize", "avalan.cli.commands.tokenizer", "tokenize"),
        ]
        for wrapper_name, module_name, function_name in commands:
            with (
                self.subTest(wrapper=wrapper_name),
                patch.object(
                    cli_module, "_load_command", return_value=command
                ) as load_command,
            ):
                result = await getattr(cli_module, wrapper_name)(
                    "arg", key="value"
                )

            self.assertEqual(result, (("arg",), {"key": "value"}))
            load_command.assert_called_once_with(module_name, function_name)

    def test_sync_lazy_command_proxies_forward_calls(self) -> None:
        cli_module = sys.modules[CLI.__module__]
        command = MagicMock(return_value="ok")
        commands = [
            ("cache_delete", "avalan.cli.commands.cache", "cache_delete"),
            (
                "cache_download",
                "avalan.cli.commands.cache",
                "cache_download",
            ),
            ("cache_list", "avalan.cli.commands.cache", "cache_list"),
            (
                "model_display",
                "avalan.cli.commands.model",
                "model_display",
            ),
            (
                "model_install",
                "avalan.cli.commands.model",
                "model_install",
            ),
            (
                "model_uninstall",
                "avalan.cli.commands.model",
                "model_uninstall",
            ),
            (
                "task_artifacts",
                "avalan.cli.commands.task",
                "task_artifacts",
            ),
            (
                "task_enqueue",
                "avalan.cli.commands.task",
                "task_enqueue",
            ),
            ("task_events", "avalan.cli.commands.task", "task_events"),
            ("task_inspect", "avalan.cli.commands.task", "task_inspect"),
            ("task_usage", "avalan.cli.commands.task", "task_usage"),
            ("task_output", "avalan.cli.commands.task", "task_output"),
            ("task_validate", "avalan.cli.commands.task", "task_validate"),
            (
                "task_pgsql_check",
                "avalan.cli.commands.task",
                "task_pgsql_check",
            ),
            (
                "task_pgsql_diagnose",
                "avalan.cli.commands.task",
                "task_pgsql_diagnose",
            ),
            (
                "task_pgsql_migrate",
                "avalan.cli.commands.task",
                "task_pgsql_migrate",
            ),
            (
                "task_pgsql_stamp",
                "avalan.cli.commands.task",
                "task_pgsql_stamp",
            ),
            (
                "task_pgsql_status",
                "avalan.cli.commands.task",
                "task_pgsql_status",
            ),
            (
                "task_retention_sweep",
                "avalan.cli.commands.task",
                "task_retention_sweep",
            ),
            ("task_run", "avalan.cli.commands.task", "task_run"),
            ("task_worker", "avalan.cli.commands.task", "task_worker"),
            ("flow_cancel", "avalan.cli.commands.flow", "flow_cancel"),
            ("flow_compile", "avalan.cli.commands.flow", "flow_compile"),
            ("flow_graph", "avalan.cli.commands.flow", "flow_graph"),
            ("flow_inspect", "avalan.cli.commands.flow", "flow_inspect"),
            ("flow_run", "avalan.cli.commands.flow", "flow_run"),
            ("flow_resume", "avalan.cli.commands.flow", "flow_resume"),
            ("flow_trace", "avalan.cli.commands.flow", "flow_trace"),
            (
                "flow_validate",
                "avalan.cli.commands.flow",
                "flow_validate",
            ),
            (
                "flow_mermaid",
                "avalan.cli.commands.flow",
                "flow_mermaid",
            ),
        ]
        for wrapper_name, module_name, function_name in commands:
            command.reset_mock()
            with (
                self.subTest(wrapper=wrapper_name),
                patch.object(
                    cli_module, "_load_command", return_value=command
                ) as load_command,
            ):
                result = getattr(cli_module, wrapper_name)("arg", key="value")

            self.assertEqual(result, "ok")
            load_command.assert_called_once_with(module_name, function_name)
            command.assert_called_once_with("arg", key="value")


class CliMainDispatchTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        from logging import getLogger

        self.logger = getLogger("cli-dispatch-test")
        with patch.object(sys, "argv", ["prog"]):
            self.cli = CLI(self.logger)

    async def test_main_dispatch(self):
        args_base = Namespace(
            verbose=None,
            quiet=True,
            login=False,
            hf_token="tok",
        )
        theme = MagicMock()
        theme._ = lambda s: s
        console = MagicMock()
        hub = MagicMock(domain="hf")
        with ExitStack() as stack:
            stack.enter_context(
                patch("avalan.cli.__main__.has_input", return_value=True)
            )
            stack.enter_context(
                patch("avalan.cli.__main__.Confirm.ask", return_value=False)
            )
            stack.enter_context(
                patch("avalan.cli.__main__.find_spec", return_value=None)
            )
            stack.enter_context(patch("avalan.cli.__main__.logger_replace"))
            stack.enter_context(patch("avalan.cli.__main__.filterwarnings"))
            run_mock = stack.enter_context(
                patch("avalan.cli.__main__.agent_run", AsyncMock())
            )
            msg_mock = stack.enter_context(
                patch("avalan.cli.__main__.agent_message_search", AsyncMock())
            )
            serve_mock = stack.enter_context(
                patch("avalan.cli.__main__.agent_serve", AsyncMock())
            )
            proxy_mock = stack.enter_context(
                patch("avalan.cli.__main__.agent_proxy", AsyncMock())
            )
            init_mock = stack.enter_context(
                patch("avalan.cli.__main__.agent_init", AsyncMock())
            )
            cache_del = stack.enter_context(
                patch("avalan.cli.__main__.cache_delete")
            )
            cache_down = stack.enter_context(
                patch("avalan.cli.__main__.cache_download")
            )
            cache_list = stack.enter_context(
                patch("avalan.cli.__main__.cache_list")
            )
            mem_doc = stack.enter_context(
                patch("avalan.cli.__main__.memory_document_index", AsyncMock())
            )
            mem_search = stack.enter_context(
                patch("avalan.cli.__main__.memory_search", AsyncMock())
            )
            mem_emb = stack.enter_context(
                patch("avalan.cli.__main__.memory_embeddings", AsyncMock())
            )
            model_disp = stack.enter_context(
                patch("avalan.cli.__main__.model_display")
            )
            model_install = stack.enter_context(
                patch("avalan.cli.__main__.model_install")
            )
            model_run = stack.enter_context(
                patch("avalan.cli.__main__.model_run", AsyncMock())
            )
            model_search = stack.enter_context(
                patch("avalan.cli.__main__.model_search", AsyncMock())
            )
            model_uninstall = stack.enter_context(
                patch("avalan.cli.__main__.model_uninstall")
            )
            deploy_run_mock = stack.enter_context(
                patch("avalan.cli.__main__.deploy_run", AsyncMock())
            )
            flow_run_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_run")
            )
            flow_cancel_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_cancel")
            )
            flow_compile_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_compile")
            )
            flow_graph_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_graph")
            )
            flow_inspect_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_inspect")
            )
            flow_resume_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_resume")
            )
            flow_trace_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_trace")
            )
            flow_validate_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_validate")
            )
            flow_mermaid_mock = stack.enter_context(
                patch("avalan.cli.__main__.flow_mermaid")
            )
            task_artifacts_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_artifacts")
            )
            task_enqueue_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_enqueue")
            )
            task_events_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_events")
            )
            task_inspect_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_inspect")
            )
            task_usage_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_usage")
            )
            task_output_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_output")
            )
            task_validate_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_validate")
            )
            task_pgsql_check_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_pgsql_check")
            )
            task_pgsql_diagnose_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_pgsql_diagnose")
            )
            task_pgsql_migrate_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_pgsql_migrate")
            )
            task_pgsql_stamp_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_pgsql_stamp")
            )
            task_pgsql_status_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_pgsql_status")
            )
            task_retention_sweep_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_retention_sweep")
            )
            task_run_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_run")
            )
            task_worker_mock = stack.enter_context(
                patch("avalan.cli.__main__.task_worker")
            )
            tokenize_mock = stack.enter_context(
                patch("avalan.cli.__main__.tokenize", AsyncMock())
            )
            scenarios = [
                ("agent", "run", run_mock),
                ("agent", "message", msg_mock),
                ("agent", "serve", serve_mock),
                ("agent", "proxy", proxy_mock),
                ("agent", "init", init_mock),
                ("cache", "delete", cache_del),
                ("cache", "download", cache_down),
                ("cache", "list", cache_list),
                ("memory", "document", mem_doc),
                ("memory", "search", mem_search),
                ("memory", "embeddings", mem_emb),
                ("model", "display", model_disp),
                ("model", "install", model_install),
                ("model", "run", model_run),
                ("model", "search", model_search),
                ("model", "uninstall", model_uninstall),
                ("deploy", "run", deploy_run_mock),
                ("flow", "cancel", flow_cancel_mock),
                ("flow", "compile", flow_compile_mock),
                ("flow", "graph", flow_graph_mock),
                ("flow", "inspect", flow_inspect_mock),
                ("flow", "resume", flow_resume_mock),
                ("flow", "run", flow_run_mock),
                ("flow", "trace", flow_trace_mock),
                ("flow", "validate", flow_validate_mock),
                ("flow", "mermaid", flow_mermaid_mock),
                ("task", "artifacts", task_artifacts_mock),
                ("task", "enqueue", task_enqueue_mock),
                ("task", "events", task_events_mock),
                ("task", "inspect", task_inspect_mock),
                ("task", "usage", task_usage_mock),
                ("task", "output", task_output_mock),
                ("task", "validate", task_validate_mock),
                ("task", "pgsql:check", task_pgsql_check_mock),
                ("task", "pgsql:diagnose", task_pgsql_diagnose_mock),
                ("task", "pgsql:migrate", task_pgsql_migrate_mock),
                ("task", "pgsql:stamp", task_pgsql_stamp_mock),
                ("task", "pgsql:status", task_pgsql_status_mock),
                ("task", "retention-sweep", task_retention_sweep_mock),
                ("task", "run", task_run_mock),
                ("task", "worker", task_worker_mock),
                ("tokenizer", None, tokenize_mock),
            ]
            for cmd, subcmd, fn in scenarios:
                args = Namespace(**vars(args_base))
                args.command = cmd
                if cmd == "agent":
                    if subcmd == "message":
                        args.agent_command = "message"
                        args.agent_message_command = "search"
                    else:
                        args.agent_command = subcmd
                elif cmd == "cache":
                    args.cache_command = subcmd
                elif cmd == "memory":
                    if subcmd == "document":
                        args.memory_command = "document"
                        args.memory_document_command = "index"
                    else:
                        args.memory_command = subcmd
                elif cmd == "model":
                    args.model_command = subcmd
                elif cmd == "deploy":
                    args.deploy_command = subcmd
                elif cmd == "flow":
                    args.flow_command = subcmd
                elif cmd == "task":
                    if subcmd and subcmd.startswith("pgsql:"):
                        args.task_command = "pgsql"
                        args.task_pgsql_command = subcmd.split(":", 1)[1]
                    else:
                        args.task_command = subcmd
                await self.cli._main(args, theme, console, hub)
                self.assertTrue(fn.called)
                fn.reset_mock()

    async def test_task_validate_failure_exits_nonzero(self):
        args = Namespace(
            command="task",
            task_command="validate",
            verbose=None,
            quiet=True,
            login=False,
            hf_token=None,
        )
        theme = MagicMock()
        theme._ = lambda s: s
        console = MagicMock()
        hub = MagicMock(domain="hf")
        with (
            patch("avalan.cli.__main__.has_input", return_value=True),
            patch("avalan.cli.__main__.Confirm.ask", return_value=False),
            patch("avalan.cli.__main__.find_spec", return_value=None),
            patch("avalan.cli.__main__.logger_replace"),
            patch("avalan.cli.__main__.filterwarnings"),
            patch(
                "avalan.cli.__main__.task_validate",
                return_value=False,
            ) as task_validate_mock,
        ):
            with self.assertRaises(SystemExit) as exit_context:
                await self.cli._main(args, theme, console, hub)

        self.assertEqual(exit_context.exception.code, 1)
        task_validate_mock.assert_called_once_with(args, console, theme)

    async def test_task_run_failure_exits_nonzero(self):
        args = Namespace(
            command="task",
            task_command="run",
            verbose=None,
            quiet=True,
            login=False,
            hf_token=None,
            task_run_json=False,
        )
        theme = MagicMock()
        theme._ = lambda s: s
        console = MagicMock()
        hub = MagicMock(domain="hf")
        with (
            patch("avalan.cli.__main__.has_input", return_value=True),
            patch("avalan.cli.__main__.Confirm.ask", return_value=False),
            patch("avalan.cli.__main__.find_spec", return_value=None),
            patch("avalan.cli.__main__.logger_replace"),
            patch("avalan.cli.__main__.filterwarnings"),
            patch(
                "avalan.cli.__main__.task_run",
                return_value=False,
            ) as task_run_mock,
        ):
            with self.assertRaises(SystemExit) as exit_context:
                await self.cli._main(args, theme, console, hub)

        self.assertEqual(exit_context.exception.code, 1)
        task_run_mock.assert_called_once_with(
            args,
            console,
            theme,
            hub,
            self.cli._logger,
        )

    async def test_flow_run_failure_exits_nonzero(self):
        args = Namespace(
            command="flow",
            flow_command="run",
            verbose=None,
            quiet=True,
            login=False,
            hf_token=None,
            task_run_json=False,
        )
        theme = MagicMock()
        theme._ = lambda s: s
        console = MagicMock()
        hub = MagicMock(domain="hf")
        with (
            patch("avalan.cli.__main__.has_input", return_value=True),
            patch("avalan.cli.__main__.Confirm.ask", return_value=False),
            patch("avalan.cli.__main__.find_spec", return_value=None),
            patch("avalan.cli.__main__.logger_replace"),
            patch("avalan.cli.__main__.filterwarnings"),
            patch(
                "avalan.cli.__main__.flow_run",
                return_value=False,
            ) as flow_run_mock,
        ):
            with self.assertRaises(SystemExit) as exit_context:
                await self.cli._main(args, theme, console, hub)

        self.assertEqual(exit_context.exception.code, 1)
        flow_run_mock.assert_called_once_with(
            args,
            console,
            theme,
            hub,
            self.cli._logger,
        )

    async def test_flow_non_executing_failure_exits_nonzero(self):
        cases = (
            ("cancel", "flow_cancel"),
            ("compile", "flow_compile"),
            ("graph", "flow_graph"),
            ("inspect", "flow_inspect"),
            ("resume", "flow_resume"),
            ("trace", "flow_trace"),
            ("validate", "flow_validate"),
            ("mermaid", "flow_mermaid"),
        )
        theme = MagicMock()
        theme._ = lambda s: s
        console = MagicMock()
        hub = MagicMock(domain="hf")

        for flow_command, function_name in cases:
            with self.subTest(flow_command=flow_command):
                args = Namespace(
                    command="flow",
                    flow_command=flow_command,
                    verbose=None,
                    quiet=True,
                    login=False,
                    hf_token=None,
                    flow_json=False,
                )
                with (
                    patch("avalan.cli.__main__.has_input", return_value=True),
                    patch(
                        "avalan.cli.__main__.Confirm.ask",
                        return_value=False,
                    ),
                    patch("avalan.cli.__main__.find_spec", return_value=None),
                    patch("avalan.cli.__main__.logger_replace"),
                    patch("avalan.cli.__main__.filterwarnings"),
                    patch(
                        f"avalan.cli.__main__.{function_name}",
                        return_value=False,
                    ) as command_mock,
                ):
                    with self.assertRaises(SystemExit) as exit_context:
                        await self.cli._main(args, theme, console, hub)

                self.assertEqual(exit_context.exception.code, 1)
                command_mock.assert_called_once_with(args, console, theme)

    def test_task_run_json_stdout_predicate(self):
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="task",
                    task_command="run",
                    task_run_json=True,
                )
            )
        )
        self.assertFalse(
            _task_run_json_stdout(
                Namespace(
                    command="task",
                    task_command="run",
                    task_run_json=False,
                )
            )
        )
        self.assertFalse(
            _task_run_json_stdout(
                Namespace(
                    command="task",
                    task_command="validate",
                    task_run_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="run",
                    task_run_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="cancel",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="compile",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="graph",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="inspect",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="validate",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="resume",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="trace",
                    flow_json=True,
                )
            )
        )
        self.assertTrue(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="mermaid",
                    flow_json=True,
                )
            )
        )
        self.assertFalse(
            _task_run_json_stdout(
                Namespace(
                    command="flow",
                    flow_command="validate",
                    flow_json=False,
                )
            )
        )


class CliMainLoginTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        from logging import getLogger

        self.logger = getLogger("cli-login-test")
        with patch.object(sys, "argv", ["prog"]):
            self.cli = CLI(self.logger)

    async def test_main_login_and_logger_replace(self):
        args = Namespace(
            command="cache",
            cache_command="list",
            verbose=1,
            quiet=False,
            login=True,
            hf_token="tok",
        )
        theme = MagicMock()
        theme._ = lambda s: s
        theme.welcome.return_value = "welcome"
        console = MagicMock()
        status_cm = MagicMock()
        status_cm.__enter__.return_value = None
        status_cm.__exit__.return_value = False
        console.status.return_value = status_cm
        hub = MagicMock(domain="hf")
        with (
            patch("avalan.cli.__main__.has_input", return_value=True),
            patch("avalan.cli.__main__.Confirm.ask", return_value=True),
            patch("avalan.cli.__main__.find_spec", return_value=True),
            patch("avalan.cli.__main__.logger_replace") as lr,
            patch("avalan.cli.__main__.filterwarnings"),
            patch("avalan.cli.__main__.cache_list") as cache_list,
        ):
            await self.cli._main(args, theme, console, hub, suggest_login=True)
        lr.assert_any_call(self.cli._logger, ["sentence_transformers"])
        lr.assert_any_call(self.cli._logger, ["httpx"])
        hub.login.assert_called_once()
        hub.user.assert_called_once()
        cache_list.assert_called_once()
        console.print.assert_any_call("welcome")


class CliMainAdditionalTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        from logging import getLogger

        self.logger = getLogger("cli-additional-test")
        with (
            patch.dict("os.environ", {"HF_TOKEN": ""}),
            patch.object(sys, "argv", ["prog"]),
        ):
            self.cli = CLI(self.logger)

    def test_add_tool_settings_arguments(self):
        from dataclasses import dataclass

        @dataclass
        class Settings:
            flag: bool = False
            count: int = 0
            ratio: float = 0.0
            name: str = ""

        parser = ArgumentParser()
        CLI._add_tool_settings_arguments(
            parser, prefix="x", settings_cls=Settings
        )
        args = parser.parse_args(
            [
                "--tool-x-flag",
                "--tool-x-count",
                "5",
                "--tool-x-ratio",
                "0.5",
                "--tool-x-name",
                "Bob",
            ]
        )
        self.assertTrue(args.tool_x_flag)
        self.assertEqual(args.tool_x_count, 5)
        self.assertEqual(args.tool_x_ratio, 0.5)
        self.assertEqual(args.tool_x_name, "Bob")

    def test_add_tool_settings_arguments_list(self):
        from dataclasses import dataclass

        @dataclass
        class Settings:
            names: list[str]

        parser = ArgumentParser()
        CLI._add_tool_settings_arguments(
            parser, prefix="y", settings_cls=Settings
        )
        args = parser.parse_args(["--tool-y-names", "Alice"])
        self.assertEqual(args.tool_y_names, "Alice")

    async def test_call_prompts_for_token_and_handles_exception(self):
        with (
            patch.object(
                sys,
                "argv",
                [
                    "prog",
                    "agent",
                    "run",
                    "--engine-uri",
                    "e",
                    "--role",
                    "r",
                    "--run-chat-x",
                    "--hf-token",
                    "",
                ],
            ),
            patch(
                "avalan.cli.__main__.translation",
                side_effect=FileNotFoundError(),
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=True),
            patch("avalan.cli.__main__.has_input", return_value=True),
            patch("builtins.open", side_effect=OSError()),
            patch(
                "avalan.cli.__main__.Prompt.ask", return_value="tok"
            ) as ask_patch,
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch.object(CLI, "_help"),
            patch.object(CLI, "_main", AsyncMock()) as main_mock,
        ):
            await self.cli()
        ask_patch.assert_called_once()
        main_mock.assert_awaited_once()
        self.assertTrue(hasattr(self.cli._parser.parse_args([]), "help_full"))

    async def test_call_handles_abort_exception(self):
        with (
            patch.object(sys, "argv", ["prog"]),
            patch(
                "avalan.cli.__main__.translation",
                return_value=SimpleNamespace(
                    gettext=lambda s: s, ngettext=lambda s, p, n: s
                ),
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch(
                "avalan.cli.__main__.Console", return_value=MagicMock()
            ) as console_patch,
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch.object(
                CLI, "_main", AsyncMock(side_effect=CommandAbortException)
            ),
        ):
            with self.assertRaises(CommandAbortException):
                await self.cli()
        console_patch.return_value.print.assert_called()

    async def test_call_handles_keyboard_interrupt(self):
        with (
            patch.object(sys, "argv", ["prog"]),
            patch(
                "avalan.cli.__main__.translation",
                return_value=SimpleNamespace(
                    gettext=lambda s: s, ngettext=lambda s, p, n: s
                ),
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch(
                "avalan.cli.__main__.Console", return_value=MagicMock()
            ) as console_patch,
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class"),
            patch.object(
                CLI, "_main", AsyncMock(side_effect=KeyboardInterrupt)
            ),
        ):
            with self.assertRaises(KeyboardInterrupt):
                await self.cli()
        console_patch.return_value.print.assert_called()

    def test_direct_keyboard_interrupts_installs_default_sigint_handler(self):
        cli_module = sys.modules[CLI.__module__]
        thread = object()
        with (
            patch.object(cli_module, "current_thread", return_value=thread),
            patch.object(cli_module, "main_thread", return_value=thread),
            patch.object(
                cli_module, "set_signal_handler", return_value="previous"
            ) as set_signal_handler,
        ):
            with cli_module._direct_keyboard_interrupts():
                pass

        set_signal_handler.assert_has_calls(
            [
                call(cli_module.SIGINT, cli_module.default_int_handler),
                call(cli_module.SIGINT, "previous"),
            ]
        )

    def test_direct_keyboard_interrupts_restores_after_interrupt(self):
        cli_module = sys.modules[CLI.__module__]
        thread = object()

        with (
            patch.object(cli_module, "current_thread", return_value=thread),
            patch.object(cli_module, "main_thread", return_value=thread),
            patch.object(
                cli_module, "set_signal_handler", return_value="previous"
            ) as set_signal_handler,
        ):
            with self.assertRaises(KeyboardInterrupt):
                with cli_module._direct_keyboard_interrupts():
                    raise KeyboardInterrupt()

        set_signal_handler.assert_has_calls(
            [
                call(cli_module.SIGINT, cli_module.default_int_handler),
                call(cli_module.SIGINT, "previous"),
            ]
        )

    def test_direct_keyboard_interrupts_skips_non_main_thread(self):
        cli_module = sys.modules[CLI.__module__]
        with (
            patch.object(cli_module, "current_thread", return_value=object()),
            patch.object(cli_module, "main_thread", return_value=object()),
            patch.object(
                cli_module, "set_signal_handler"
            ) as set_signal_handler,
        ):
            with cli_module._direct_keyboard_interrupts():
                pass

        set_signal_handler.assert_not_called()

    def test_print_bye_skips_when_already_printed(self):
        console = MagicMock()
        self.cli._abort_printed = True

        self.cli._print_bye(console)

        console.print.assert_not_called()

    def test_print_bye_skips_when_quiet(self):
        console = MagicMock()

        self.cli._print_bye(console, quiet=True)

        console.print.assert_not_called()
        self.assertFalse(self.cli._abort_printed)

    def test_can_use_anonymous_hub_for_remote_model_run(self):
        args = Namespace(
            command="model",
            model_command="run",
            model="ai://openai/gpt-4o-mini",
        )

        self.assertTrue(CLI._can_use_anonymous_hub(args))

    def test_can_use_anonymous_hub_for_local_ds4_model_run(self):
        args = Namespace(
            command="model",
            model_command="run",
            model="ai://local/../pyds4/.local/ds4/ds4flash.gguf",
            backend="ds4",
        )

        self.assertTrue(CLI._can_use_anonymous_hub(args))

    def test_can_use_anonymous_hub_for_task_command(self):
        args = Namespace(command="task", task_command="validate")

        self.assertTrue(CLI._can_use_anonymous_hub(args))

    async def test_call_uses_anonymous_hub_for_remote_model_run(self):
        captured_hub = None

        async def capture_main(args, theme, console, hub):
            nonlocal captured_hub
            del args, theme, console
            captured_hub = hub

        with (
            patch.object(
                sys,
                "argv",
                ["prog", "model", "run", "ai://openai/gpt-4o-mini"],
            ),
            patch(
                "avalan.cli.__main__.translation",
                return_value=SimpleNamespace(
                    gettext=lambda s: s, ngettext=lambda s, p, n: s
                ),
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch.object(CLI, "_needs_hf_token", return_value=False),
            patch("avalan.cli.__main__._huggingface_hub_class") as hub_class,
            patch.object(CLI, "_main", AsyncMock(side_effect=capture_main)),
        ):
            await self.cli()

        self.assertIsNotNone(captured_hub)
        self.assertEqual(captured_hub.cache_dir, "~/.cache/huggingface/hub")
        hub_class.assert_not_called()

    async def test_call_uses_anonymous_hub_for_local_ds4_model_run(self):
        captured_hub = None

        async def capture_main(args, theme, console, hub):
            nonlocal captured_hub
            del args, theme, console
            captured_hub = hub

        with (
            patch.object(
                sys,
                "argv",
                [
                    "prog",
                    "model",
                    "run",
                    "ai://local/../pyds4/.local/ds4/ds4flash.gguf",
                    "--backend",
                    "ds4",
                ],
            ),
            patch(
                "avalan.cli.__main__.translation",
                return_value=SimpleNamespace(
                    gettext=lambda s: s, ngettext=lambda s, p, n: s
                ),
            ),
            patch(
                "avalan.cli.__main__.FancyTheme",
                return_value=MagicMock(get_styles=lambda: {}),
            ),
            patch("avalan.cli.__main__.Console", return_value=MagicMock()),
            patch("avalan.cli.__main__._huggingface_hub_class") as hub_class,
            patch.object(CLI, "_main", AsyncMock(side_effect=capture_main)),
        ):
            await self.cli()

        self.assertIsNotNone(captured_hub)
        self.assertEqual(captured_hub.cache_dir, "~/.cache/huggingface/hub")
        hub_class.assert_not_called()

    async def test_call_preserves_explicit_hf_token_for_task_command(self):
        captured_hub = None

        async def capture_main(args, theme, console, hub):
            nonlocal captured_hub
            del args, theme, console
            captured_hub = hub

        with TemporaryDirectory() as tmp:
            task = Path(tmp) / "task.toml"
            task.write_text("", encoding="utf-8")
            with (
                patch.object(
                    sys,
                    "argv",
                    [
                        "prog",
                        "task",
                        "validate",
                        str(task),
                        "--hf-token",
                        "task-token",
                    ],
                ),
                patch(
                    "avalan.cli.__main__.translation",
                    return_value=SimpleNamespace(
                        gettext=lambda s: s, ngettext=lambda s, p, n: s
                    ),
                ),
                patch(
                    "avalan.cli.__main__.FancyTheme",
                    return_value=MagicMock(get_styles=lambda: {}),
                ),
                patch(
                    "avalan.cli.__main__.Console",
                    return_value=MagicMock(),
                ),
                patch.object(CLI, "_needs_hf_token", return_value=False),
                patch(
                    "avalan.cli.__main__._huggingface_hub_class"
                ) as hub_class,
                patch.object(
                    CLI,
                    "_main",
                    AsyncMock(side_effect=capture_main),
                ),
            ):
                await self.cli()

        self.assertIsNotNone(captured_hub)
        hub_class.return_value.assert_called_once_with(
            "task-token",
            "~/.cache/huggingface/hub",
            self.logger,
        )
        self.assertEqual(captured_hub, hub_class.return_value.return_value)


class CliMainFunctionTestCase(TestCase):
    def test_main_invokes_run(self):
        with (
            patch("avalan.cli.__main__.run_in_loop") as run,
            patch("avalan.cli.__main__.CLI") as cli,
        ):
            from avalan.cli.__main__ import main

            main()
        run.assert_called_once()
        cli.assert_called_once()

    def test_main_sets_tokenizers_parallelism_default(self):
        with (
            patch.dict("avalan.cli.__main__.environ", {}, clear=True),
            patch("avalan.cli.__main__.run_in_loop"),
            patch("avalan.cli.__main__.CLI"),
        ):
            from avalan.cli.__main__ import main

            main()
            self.assertEqual(
                sys.modules["avalan.cli.__main__"].environ[
                    "TOKENIZERS_PARALLELISM"
                ],
                "false",
            )

    def test_main_handles_keyboard_interrupt(self):
        cli = MagicMock()
        with (
            patch(
                "avalan.cli.__main__.run_in_loop",
                side_effect=KeyboardInterrupt,
            ),
            patch("avalan.cli.__main__.CLI", return_value=cli),
        ):
            from avalan.cli.__main__ import main

            main()
        cli._print_bye.assert_called_once_with()

    def test_main_handles_cancelled_error(self):
        cli = MagicMock()
        with (
            patch(
                "avalan.cli.__main__.run_in_loop",
                side_effect=asyncio.CancelledError,
            ),
            patch("avalan.cli.__main__.CLI", return_value=cli),
        ):
            from avalan.cli.__main__ import main

            main()
        cli._print_bye.assert_called_once_with()

    def test_main_handles_abort_exception(self):
        cli = MagicMock()
        with (
            patch(
                "avalan.cli.__main__.run_in_loop",
                side_effect=CommandAbortException,
            ),
            patch("avalan.cli.__main__.CLI", return_value=cli),
        ):
            from avalan.cli.__main__ import main

            main()
        cli._print_bye.assert_called_once_with()

    def test_run_in_loop_bounds_interrupt_cleanup(self):
        cli_module = sys.modules[CLI.__module__]

        async def stubborn_task() -> None:
            try:
                await asyncio.Future()
            except asyncio.CancelledError:
                await asyncio.Event().wait()

        async def interrupted() -> None:
            asyncio.create_task(stubborn_task())
            raise KeyboardInterrupt()

        with (
            patch.object(cli_module, "_INTERRUPT_DRAIN_TIMEOUT", 0.01),
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(interrupted())

    def test_run_in_loop_signal_handler_interrupts_immediately(self):
        cli_module = sys.modules[CLI.__module__]
        thread = object()
        captured: dict[str, object] = {}

        def set_handler(_signal: int, handler: object) -> str:
            if callable(handler):
                captured["handler"] = handler
            return "previous"

        async def interrupted() -> None:
            handler = captured["handler"]
            assert callable(handler)
            handler(cli_module.SIGINT, None)
            self.fail("signal handler should raise")

        with (
            patch.object(cli_module, "current_thread", return_value=thread),
            patch.object(cli_module, "main_thread", return_value=thread),
            patch.object(
                cli_module, "set_signal_handler", side_effect=set_handler
            ) as set_signal_handler,
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(interrupted())

        self.assertEqual(set_signal_handler.call_count, 2)
        self.assertEqual(
            set_signal_handler.call_args_list[-1],
            call(cli_module.SIGINT, "previous"),
        )

    def test_run_in_loop_signal_callback_interrupts_waiter(self):
        cli_module = sys.modules[CLI.__module__]
        thread = object()
        captured: dict[str, object] = {}

        def set_handler(_signal: int, handler: object) -> str:
            if callable(handler):
                captured["handler"] = handler
            return "previous"

        async def interrupted() -> None:
            handler = captured["handler"]
            assert callable(handler)
            try:
                handler(cli_module.SIGINT, None)
            except KeyboardInterrupt:
                pass
            await asyncio.Future()

        with (
            patch.object(cli_module, "current_thread", return_value=thread),
            patch.object(cli_module, "main_thread", return_value=thread),
            patch.object(
                cli_module, "set_signal_handler", side_effect=set_handler
            ),
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(interrupted())

    def test_run_in_loop_cancels_main_task_after_interrupt_signal(self):
        cli_module = sys.modules[CLI.__module__]

        async def fake_wait(tasks, return_when):
            interrupt_signal = next(
                task for task in tasks if not isinstance(task, asyncio.Task)
            )
            return {interrupt_signal}, set()

        async def never_finishes() -> None:
            await asyncio.Future()

        with (
            patch.object(cli_module, "wait", side_effect=fake_wait),
            patch.object(cli_module, "current_thread", return_value=object()),
            patch.object(cli_module, "main_thread", return_value=object()),
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(never_finishes())

    def test_run_in_loop_signal_handler_falls_back_after_runtime_error(self):
        cli_module = sys.modules[CLI.__module__]
        thread = object()
        loop = asyncio.new_event_loop()
        captured: dict[str, object] = {}

        def set_handler(_signal: int, handler: object) -> str:
            if callable(handler):
                captured["handler"] = handler
            return "previous"

        async def interrupted() -> None:
            handler = captured["handler"]
            assert callable(handler)
            handler(cli_module.SIGINT, None)
            self.fail("signal handler should raise")

        with (
            patch.object(cli_module, "new_event_loop", return_value=loop),
            patch.object(cli_module, "current_thread", return_value=thread),
            patch.object(cli_module, "main_thread", return_value=thread),
            patch.object(
                cli_module, "set_signal_handler", side_effect=set_handler
            ),
            patch.object(
                loop,
                "call_soon_threadsafe",
                side_effect=RuntimeError("closed"),
            ),
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(interrupted())

    def test_run_in_loop_skips_executor_shutdown_on_interrupt(self):
        cli_module = sys.modules[CLI.__module__]
        ready = ThreadEvent()
        release = ThreadEvent()

        def blocked_worker() -> None:
            ready.set()
            release.wait(timeout=5)

        async def interrupted() -> None:
            asyncio.create_task(asyncio.to_thread(blocked_worker))
            while not ready.is_set():
                await asyncio.sleep(0)
            raise KeyboardInterrupt()

        start = perf_counter()
        try:
            with self.assertRaises(KeyboardInterrupt):
                cli_module.run_in_loop(interrupted())
        finally:
            release.set()

        self.assertLess(perf_counter() - start, 1)

    def test_run_in_loop_drains_pending_tasks_on_success(self):
        cli_module = sys.modules[CLI.__module__]
        cleaned_up: list[str] = []

        async def pending_task() -> None:
            try:
                await asyncio.Future()
            except asyncio.CancelledError:
                cleaned_up.append("cancelled")

        async def complete_with_pending_task() -> None:
            asyncio.create_task(pending_task())
            await asyncio.sleep(0)

        cli_module.run_in_loop(complete_with_pending_task())

        self.assertEqual(cleaned_up, ["cancelled"])

    def test_run_in_loop_bounds_pending_cleanup_after_interrupt(self):
        cli_module = sys.modules[CLI.__module__]

        async def interrupted() -> None:
            started = asyncio.Event()

            async def stubborn_task() -> None:
                try:
                    started.set()
                    await asyncio.Future()
                except asyncio.CancelledError:
                    await asyncio.Event().wait()

            asyncio.create_task(stubborn_task())
            await started.wait()
            raise KeyboardInterrupt()

        with (
            patch.object(cli_module, "_INTERRUPT_DRAIN_TIMEOUT", 0.01),
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(interrupted())

    def test_run_in_loop_bounds_async_generator_shutdown_after_interrupt(self):
        cli_module = sys.modules[CLI.__module__]

        async def hanging_generator():
            try:
                yield "started"
            finally:
                await asyncio.Event().wait()

        async def interrupted() -> None:
            generator = hanging_generator()
            await generator.__anext__()
            raise KeyboardInterrupt()

        with (
            patch.object(cli_module, "_INTERRUPT_DRAIN_TIMEOUT", 0.01),
            self.assertRaises(KeyboardInterrupt),
        ):
            cli_module.run_in_loop(interrupted())
