"""Custom CLI argument parser for Vidhi configuration.

This module provides a CLI parser designed specifically for dataclass-based
configurations with support for:
- Polymorphic configs with variant-specific defaults
- Clean nested field naming (--scheduler.timeout)
- Container bracket-access syntax (--gpus[0].node_ip, --map[key]=value)
- Grouped help output with colors
- Type validation
"""

from __future__ import annotations

import collections.abc
import os
import re
import sys
from dataclasses import dataclass, field
from enum import Enum
from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Set,
    Tuple,
    Union,
    get_args,
    get_origin,
)

# =============================================================================
# ANSI Color Support
# =============================================================================


class Colors:
    """ANSI color codes for terminal output."""

    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"

    # Colors
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"

    # Bright colors
    BRIGHT_RED = "\033[91m"
    BRIGHT_GREEN = "\033[92m"
    BRIGHT_YELLOW = "\033[93m"
    BRIGHT_BLUE = "\033[94m"
    BRIGHT_CYAN = "\033[96m"


def _format_config_name(name: str) -> str:
    """Format a config name for display.

    Converts snake_case names like "resource_allocator_config" to
    title case "Resource Allocator Config".

    Args:
        name: The snake_case config name

    Returns:
        Properly formatted title case name
    """
    # Remove trailing _config suffix for cleaner display
    display_name = name
    if display_name.endswith("_config"):
        display_name = display_name[:-7]

    # Split by underscores and capitalize each word
    words = display_name.split("_")
    return " ".join(word.capitalize() for word in words)


def _supports_color() -> bool:
    """Check if the terminal supports colors."""
    # Check for NO_COLOR environment variable (https://no-color.org/)
    if os.environ.get("NO_COLOR"):
        return False

    # Check for FORCE_COLOR environment variable
    if os.environ.get("FORCE_COLOR"):
        return True

    # Check if stdout is a TTY
    if not hasattr(sys.stdout, "isatty") or not sys.stdout.isatty():
        return False

    # Check TERM environment variable
    term = os.environ.get("TERM", "")
    if term == "dumb":
        return False

    return True


class ColorFormatter:
    """Formatter that applies colors when supported."""

    def __init__(self, use_colors: Optional[bool] = None):
        if use_colors is None:
            use_colors = _supports_color()
        self.use_colors = use_colors

    def _c(self, text: str, *colors: str) -> str:
        """Apply colors to text if colors are enabled."""
        if not self.use_colors:
            return text
        color_code = "".join(colors)
        return f"{color_code}{text}{Colors.RESET}"

    def bold(self, text: str) -> str:
        return self._c(text, Colors.BOLD)

    def dim(self, text: str) -> str:
        return self._c(text, Colors.DIM)

    def option(self, text: str) -> str:
        """Format CLI option (e.g., --host)."""
        return self._c(text, Colors.CYAN)

    def type_hint(self, text: str) -> str:
        """Format type hint (e.g., <int>)."""
        return self._c(text, Colors.YELLOW)

    def default(self, text: str) -> str:
        """Format default value."""
        return self._c(text, Colors.GREEN)

    def required(self, text: str) -> str:
        """Format required marker."""
        return self._c(text, Colors.BRIGHT_RED, Colors.BOLD)

    def choices(self, text: str) -> str:
        """Format choices."""
        return self._c(text, Colors.MAGENTA)

    def header(self, text: str) -> str:
        """Format section header."""
        return self._c(text, Colors.BOLD, Colors.WHITE)

    def variant(self, text: str) -> str:
        """Format variant name."""
        return self._c(text, Colors.BRIGHT_CYAN, Colors.BOLD)

    def yaml_key(self, text: str) -> str:
        """Format YAML key."""
        return self._c(text, Colors.CYAN)

    def yaml_value(self, text: str) -> str:
        """Format YAML value."""
        return self._c(text, Colors.GREEN)

    def yaml_comment(self, text: str) -> str:
        """Format YAML comment."""
        return self._c(text, Colors.DIM)


# =============================================================================
# Argument Data Classes
# =============================================================================


@dataclass
class Argument:
    """Definition of a CLI argument."""

    name: str  # Internal name (e.g., "scheduler.timeout")
    cli_name: str  # CLI name (e.g., "scheduler.timeout")
    arg_type: Union[type, Callable[[str], Any]]  # Expected type or converter function
    default: Any  # Default value (or dict of variant -> default)
    help_text: str  # Help description
    required: bool = False
    is_boolean: bool = False
    is_list: bool = False  # Whether this argument accepts multiple values
    variants: Optional[Set[str]] = None  # Which variants this field belongs to
    type_field: Optional[str] = None  # The type selector field (e.g., "scheduler.type")
    group: Optional[str] = None  # Argument group for help display
    choices: Optional[List[str]] = None  # Valid choices for this argument
    aliases: Optional[List[str]] = None  # Alternative CLI names for this argument
    shortcuts: Optional[Dict[Enum, str]] = (
        None  # Enum member -> short flag string (e.g., {Action.REATTACH: "r"})
    )
    positional: bool = False  # Whether this argument accepts positional values
    is_type_selector: bool = False  # Whether this is a polymorphic type selector field
    poly_base_field: Optional[str] = (
        None  # Base field this type selector controls (e.g., "controller")
    )
    is_container: bool = False  # Whether this is a container field supporting [] access
    container_info: Optional[Any] = None  # ContainerFieldInfo for bracket-access


@dataclass
class ArgumentGroup:
    """A group of related arguments for help display."""

    name: str
    title: str
    arguments: List[Argument]


@dataclass
class HelpNode:
    """Node in the help output tree."""

    doc: str = ""
    args: List[Argument] = field(default_factory=list)
    groups: Dict[str, "HelpNode"] = field(default_factory=dict)
    # Poly config info
    type_arg: Optional[Argument] = None
    variants: Dict[str, "HelpNode"] = field(default_factory=dict)
    default_variant: str = ""


# =============================================================================
# Container Bracket-Path Parsing
# =============================================================================

# Regex to tokenize the container path portion (after the base field name).
# Matches: [key_content] or .field_name
_CONTAINER_SEGMENT_RE = re.compile(r"\[([^\]]+)\]|\.([a-zA-Z_][a-zA-Z0-9_]*)")


def parse_container_path(cli_name: str) -> Tuple[str, List[Tuple[str, str]]]:
    """Parse a CLI name into a base field name and container access segments.

    Args:
        cli_name: e.g., "gpus[0].node_ip" or "engine.rm[0].field[1].x"

    Returns:
        Tuple of (base_cli_name, segments) where segments are:
        - ("key", "0") for [0]
        - ("field", "node_ip") for .node_ip
    """
    # Find the first bracket — everything before it is the base field name
    bracket_pos = cli_name.find("[")
    if bracket_pos == -1:
        return cli_name, []

    base = cli_name[:bracket_pos]
    remainder = cli_name[bracket_pos:]

    segments: List[Tuple[str, str]] = []
    for m in _CONTAINER_SEGMENT_RE.finditer(remainder):
        bracket_content = m.group(1)
        dot_field = m.group(2)
        if bracket_content is not None:
            segments.append(("key", bracket_content))
        else:
            segments.append(("field", dot_field))

    return base, segments


def _set_nested_value(
    target: dict,
    segments: List[Tuple[str, str]],
    value: str,
) -> None:
    """Walk the path segments and set the leaf value in a nested dict.

    Intermediate dicts are created as needed. The container_info type tree
    is used later during resolution to coerce keys and instantiate types.

    Args:
        target: The dict to populate (mutated in place)
        segments: Path segments from parse_container_path
        value: The raw string value to set at the leaf
    """
    current = target
    for idx, (_seg_type, seg_key) in enumerate(segments):
        is_last = idx == len(segments) - 1
        if is_last:
            current[seg_key] = value
        else:
            if seg_key not in current:
                current[seg_key] = {}
            current = current[seg_key]


# =============================================================================
# Main Parser Class
# =============================================================================


class VidhiArgumentParser:
    """Custom argument parser for Vidhi configurations.

    Unlike argparse, this parser:
    - Handles polymorphic configs with variant-specific defaults
    - Shows all variant defaults in help text
    - Validates fields against selected variant type
    - Produces clean, grouped help output with colors
    """

    def __init__(self, prog: Optional[str] = None, description: Optional[str] = None):
        self.prog = prog or os.path.basename(sys.argv[0])
        self.description = description
        self.arguments: Dict[str, Argument] = {}
        self.groups: Dict[str, ArgumentGroup] = {}
        self._group_order: List[str] = []  # Track insertion order
        self._group_descriptions: Dict[str, str] = {}  # Group docstrings
        self._variant_descriptions: Dict[str, Dict[str, str]] = (
            {}
        )  # group -> {variant -> desc}
        self._nested_group_descriptions: Dict[str, str] = {}  # nested path -> docstring
        self._nested_groups: Dict[str, str] = {}  # field_name -> parent nested path
        self._color = ColorFormatter()
        self.help_tree: Optional[HelpNode] = None

    def set_group_description(self, group: str, description: str) -> None:
        """Set a description for a group (from class docstring)."""
        self._group_descriptions[group] = description

    def set_variant_description(
        self, group: str, variant: str, description: str
    ) -> None:
        """Set a description for a variant (from variant class docstring)."""
        if group not in self._variant_descriptions:
            self._variant_descriptions[group] = {}
        self._variant_descriptions[group][variant] = description

    def set_nested_group_description(self, group_path: str, description: str) -> None:
        """Set a description for a nested dataclass group."""
        self._nested_group_descriptions[group_path] = description

    def set_nested_groups(self, nested_groups: Dict[str, str]) -> None:
        """Set the mapping of field names to their parent nested dataclass paths."""
        self._nested_groups = nested_groups

    def add_argument(
        self,
        name: str,
        cli_name: str,
        arg_type: Union[type, Callable[[str], Any]],
        default: Any = None,
        help_text: str = "",
        required: bool = False,
        is_boolean: bool = False,
        is_list: bool = False,
        variants: Optional[Set[str]] = None,
        type_field: Optional[str] = None,
        group: Optional[str] = None,
        choices: Optional[List[str]] = None,
        aliases: Optional[List[str]] = None,
        shortcuts: Optional[Dict[Enum, str]] = None,
        positional: bool = False,
        is_type_selector: bool = False,
        poly_base_field: Optional[str] = None,
        is_container: bool = False,
        container_info: Optional[Any] = None,
    ) -> None:
        """Add an argument to the parser."""
        arg = Argument(
            name=name,
            cli_name=cli_name,
            arg_type=arg_type,
            default=default,
            help_text=help_text,
            required=required,
            is_boolean=is_boolean,
            is_list=is_list,
            variants=variants,
            type_field=type_field,
            group=group,
            choices=choices,
            aliases=aliases,
            shortcuts=shortcuts,
            positional=positional,
            is_type_selector=is_type_selector,
            poly_base_field=poly_base_field,
            is_container=is_container,
            container_info=container_info,
        )
        self.arguments[name] = arg

        # Add to group
        group_key = group or ""
        if group_key not in self.groups:
            # Use proper title case formatting for config names
            if group_key:
                formatted_name = _format_config_name(group_key)
                title = f"{formatted_name} Options"
            else:
                title = "Options"
            self.groups[group_key] = ArgumentGroup(
                name=group_key, title=title, arguments=[]
            )
            self._group_order.append(group_key)
        self.groups[group_key].arguments.append(arg)

    def parse_args(
        self,
        args: Optional[List[str]] = None,
    ) -> Tuple[Dict[str, Any], Set[str]]:
        """Parse command line arguments.

        Args:
            args: List of arguments (defaults to sys.argv[1:])

        Returns:
            Tuple of (parsed_values, provided_args) where:
            - parsed_values: Dictionary mapping argument names to parsed values
            - provided_args: Set of argument names that were explicitly provided
        """
        if args is None:
            args = sys.argv[1:]

        # Check for help
        if "-h" in args or "--help" in args:
            self.print_help()
            sys.exit(0)

        # Parse arguments
        result: Dict[str, Any] = {}
        provided: Set[str] = set()
        positional_values: List[str] = []
        i = 0

        while i < len(args):
            arg = args[i]

            # Handle single-dash shortcuts (e.g., -r, -ld, -cc)
            if arg.startswith("-") and not arg.startswith("--") and len(arg) > 1:
                shortcut_key = arg[1:]  # Strip the leading -
                match = self._find_shortcut(shortcut_key)
                if match:
                    internal_name, value = match
                    # Shortcut values are always Enum members; store as string
                    # for consistency with how the parser handles enum fields
                    result[internal_name] = value.value
                    provided.add(internal_name)
                    i += 1
                    continue
                else:
                    print(f"Error: Unknown shortcut -{shortcut_key}", file=sys.stderr)
                    sys.exit(1)

            if not arg.startswith("--"):
                # Collect as positional argument
                positional_values.append(arg)
                i += 1
                continue

            # Handle --flag or --flag=value or --flag value
            if "=" in arg:
                cli_name, value = arg[2:].split("=", 1)
            else:
                cli_name = arg[2:]
                value = None

            # Handle container bracket access: --field[key].subfield value
            if "[" in cli_name:
                base_cli_name, container_segments = parse_container_path(cli_name)
                internal_name = self._find_argument_by_cli_name(base_cli_name)
                if not internal_name:
                    print(
                        f"Error: Unknown argument --{base_cli_name}",
                        file=sys.stderr,
                    )
                    sys.exit(1)

                arg_def = self.arguments[internal_name]
                if not arg_def.is_container:
                    print(
                        f"Error: --{base_cli_name} does not support [] access",
                        file=sys.stderr,
                    )
                    sys.exit(1)

                # Get the value
                if value is None:
                    if i + 1 >= len(args):
                        print(
                            f"Error: --{cli_name} requires a value",
                            file=sys.stderr,
                        )
                        sys.exit(1)
                    value = args[i + 1]
                    i += 2
                else:
                    i += 1

                # Accumulate into a nested dict for this container field
                if internal_name not in result or not isinstance(
                    result[internal_name], dict
                ):
                    result[internal_name] = {}
                _set_nested_value(result[internal_name], container_segments, value)
                provided.add(internal_name)
                continue

            internal_name = self._find_argument_by_cli_name(cli_name)
            if not internal_name:
                print(f"Error: Unknown argument --{cli_name}", file=sys.stderr)
                sys.exit(1)

            arg_def = self.arguments[internal_name]

            # Get value(s)
            if arg_def.is_list:
                # Collect multiple values until next flag
                values = []
                if value is not None:
                    # Value was provided with = syntax
                    values.append(value)
                    i += 1
                else:
                    i += 1  # Move past the flag
                    # Collect values until next flag or end of args
                    while i < len(args) and not args[i].startswith("--"):
                        values.append(args[i])
                        i += 1

                if not values:
                    print(
                        f"Error: --{cli_name} requires at least one value",
                        file=sys.stderr,
                    )
                    sys.exit(1)

                # Convert each value to the target type
                try:
                    result[internal_name] = [arg_def.arg_type(v) for v in values]
                except (ValueError, TypeError) as e:
                    print(
                        f"Error: Invalid value for --{cli_name}: {e}", file=sys.stderr
                    )
                    sys.exit(1)
            else:
                # Single value argument
                if value is None:
                    if i + 1 >= len(args):
                        print(f"Error: --{cli_name} requires a value", file=sys.stderr)
                        sys.exit(1)
                    value = args[i + 1]
                    i += 2
                else:
                    i += 1

                # Convert type
                try:
                    if arg_def.arg_type is bool:
                        lower_val = value.lower()
                        if lower_val in ("true", "1", "yes"):
                            result[internal_name] = True
                        elif lower_val in ("false", "0", "no"):
                            result[internal_name] = False
                        else:
                            print(
                                f"Error: Invalid boolean value '{value}' for --{cli_name}. "
                                f"Use true/false, yes/no, or 1/0.",
                                file=sys.stderr,
                            )
                            sys.exit(1)
                    elif isinstance(arg_def.arg_type, type) and issubclass(
                        arg_def.arg_type, Enum
                    ):
                        result[internal_name] = value
                    else:
                        result[internal_name] = arg_def.arg_type(value)
                except (ValueError, TypeError) as e:
                    print(
                        f"Error: Invalid value for --{cli_name}: {e}", file=sys.stderr
                    )
                    sys.exit(1)

                # Validate choices if specified (case-insensitive for type selectors)
                if arg_def.choices:
                    value_to_check = result[internal_name]
                    # Check case-insensitively
                    choices_lower = {str(c).lower(): c for c in arg_def.choices}
                    if str(value_to_check).lower() in choices_lower:
                        # Normalize to the canonical case
                        result[internal_name] = choices_lower[
                            str(value_to_check).lower()
                        ]
                    elif value_to_check not in arg_def.choices:
                        print(
                            f"Error: Invalid value '{value_to_check}' for --{cli_name}. "
                            f"Choose from: {', '.join(str(c) for c in arg_def.choices)}",
                            file=sys.stderr,
                        )
                        sys.exit(1)

            provided.add(internal_name)

        # Assign positional values
        if positional_values:
            positional_arg = self._find_positional_argument()
            if positional_arg:
                if positional_arg.is_list:
                    result[positional_arg.name] = positional_values
                else:
                    if len(positional_values) > 1:
                        print(
                            f"Error: Too many positional arguments: {positional_values}",
                            file=sys.stderr,
                        )
                        sys.exit(1)
                    result[positional_arg.name] = positional_values[0]
                provided.add(positional_arg.name)
            else:
                print(
                    f"Error: Unexpected positional arguments: {positional_values}",
                    file=sys.stderr,
                )
                sys.exit(1)

        # Apply defaults and validate
        self._apply_defaults(result, provided)
        self._validate_required(result, provided)
        # Note: Variant field validation is done after merging with config files
        # in create_from_cli_args() to validate against the final type

        return result, provided

    def _find_argument_by_cli_name(self, cli_name: str) -> Optional[str]:
        """Find internal argument name by CLI name or alias."""
        for name, arg in self.arguments.items():
            if arg.cli_name == cli_name:
                return name
            # Check aliases
            if arg.aliases and cli_name in arg.aliases:
                return name
        return None

    def _find_shortcut(self, flag: str) -> Optional[Tuple[str, Any]]:
        """Find an argument and its mapped Enum value by shortcut flag.

        Args:
            flag: The shortcut flag string (e.g., "r", "ld", "cc")

        Returns:
            Tuple of (internal_name, enum_value) or None if not found
        """
        for name, arg in self.arguments.items():
            if arg.shortcuts:
                for enum_val, shortcut_flag in arg.shortcuts.items():
                    if shortcut_flag == flag:
                        return (name, enum_val)
        return None

    def _find_positional_argument(self) -> Optional[Argument]:
        """Find the argument marked as positional."""
        for arg in self.arguments.values():
            if arg.positional:
                return arg
        return None

    def _apply_defaults(self, result: Dict[str, Any], provided: Set[str]) -> None:
        """Apply default values for unprovided arguments."""
        for name, arg in self.arguments.items():
            if name not in provided:
                if isinstance(arg.default, dict):
                    # Variant-specific default - use default variant's value
                    # Get the selected type
                    if arg.type_field and arg.type_field in result:
                        selected_type = result[arg.type_field]
                        if selected_type in arg.default:
                            result[name] = arg.default[selected_type]
                        else:
                            # Use first available default
                            result[name] = next(iter(arg.default.values()))
                    else:
                        result[name] = next(iter(arg.default.values()))
                else:
                    result[name] = arg.default

    def _validate_required(self, result: Dict[str, Any], provided: Set[str]) -> None:
        """Validate that required arguments are provided."""
        for name, arg in self.arguments.items():
            if arg.required and name not in provided:
                print(f"Error: --{arg.cli_name} is required", file=sys.stderr)
                sys.exit(1)

    def print_help(self) -> None:
        """Print formatted help message with colors."""
        c = self._color
        lines = []

        # Usage line
        lines.append(self._format_usage())
        lines.append("")

        if self.description:
            lines.append(self.description)
            lines.append("")

        # Built-in options
        lines.append(c.header("Built-in Options:"))
        lines.append(
            f"  {c.option('-h')}, {c.option('--help')}"
            "            show this help message and exit"
        )
        lines.append(
            f"  {c.option('--config')} {c.type_hint('<path>')}"
            "       load configuration from YAML file"
        )
        lines.append(
            f"  {c.option('--install-shell-completions')} {c.type_hint('[shell]')}"
        )
        lines.append("                        install shell completions and exit")
        lines.append(f"  {c.option('--export-json-schema')} {c.type_hint('[path]')}")
        lines.append("                        export JSON schema for IDE autocomplete")
        lines.append("")

        # Render help tree
        if self.help_tree is not None:
            lines.extend(self._render_help_tree(self.help_tree, indent=0, is_root=True))

        # Add YAML example
        lines.extend(self._format_yaml_example())

        print("\n".join(lines))

    def _render_help_tree(
        self, node: HelpNode, indent: int, is_root: bool = False
    ) -> List[str]:
        """Render a HelpNode tree into formatted help lines."""
        c = self._color
        lines: List[str] = []
        prefix = " " * indent

        # Render direct args
        if node.args:
            if is_root:
                lines.append(c.header("Options:"))
            for arg in node.args:
                if not arg.positional:
                    lines.extend(
                        self._format_argument_help(
                            arg, show_variants=False, indent=indent
                        )
                    )
            if is_root and node.args:
                lines.append("")

        # Render child groups
        for group_name, child in node.groups.items():
            if child.type_arg is not None or child.variants:
                # Polymorphic group
                lines.extend(self._render_poly_node(child, group_name, indent, is_root))
            else:
                # Plain nested dataclass group
                if is_root:
                    formatted = _format_config_name(group_name)
                    lines.append(c.header(f"{formatted} Options:"))
                    if child.doc:
                        lines.append(f"  {child.doc}")
                    lines.extend(self._render_help_tree(child, indent=indent + 2))
                    lines.append("")
                else:
                    formatted = _format_config_name(group_name)
                    lines.append(f"{prefix}{c.header(formatted + ':')}")
                    if child.doc:
                        lines.append(f"{prefix}  {c.dim(child.doc)}")
                    lines.extend(self._render_help_tree(child, indent=indent + 2))

        return lines

    def _render_poly_node(
        self,
        node: HelpNode,
        group_name: str,
        indent: int,
        is_root: bool,
    ) -> List[str]:
        """Render a polymorphic HelpNode with variant sections."""
        c = self._color
        lines: List[str] = []
        prefix = " " * indent

        # Header
        if is_root:
            formatted = _format_config_name(group_name)
            lines.append(c.header(f"{formatted} Options:"))
        else:
            formatted = _format_config_name(group_name)
            lines.append(f"{prefix}{c.header(formatted + ' Options:')}")

        content_prefix = prefix + "  "

        # Group description
        if node.doc:
            lines.append(f"{content_prefix}{node.doc}")

        # Type selector line
        if node.type_arg:
            choices = node.type_arg.choices or sorted(node.variants.keys())
            choices_str = c.choices("{" + ",".join(choices) + "}")
            default_str = ""
            if node.type_arg.default:
                default_str = f" {c.dim('[' + str(node.type_arg.default) + ']')}"
            lines.append(
                f"{content_prefix}Select variant with "
                f"{c.option('--' + node.type_arg.cli_name)} "
                f"{choices_str}{default_str}"
            )
        lines.append("")

        # Render each variant
        for variant_name in sorted(node.variants.keys()):
            variant_node = node.variants[variant_name]

            # Variant header
            formatted_group = _format_config_name(group_name)
            default_marker = ""
            if variant_name == node.default_variant:
                default_marker = f" {c.dim('(default)')}"
            lines.append(
                f"{content_prefix}{c.dim(formatted_group + ' →')} "
                f"{c.variant(variant_name)}{default_marker}"
            )

            # Variant description
            if variant_node.doc:
                lines.append(f"{content_prefix}  {c.dim(variant_node.doc)}")

            # Variant args and sub-groups
            lines.extend(self._render_help_tree(variant_node, indent=indent + 4))
            lines.append("")

        return lines

    def _format_usage(self) -> str:
        """Format the usage line."""
        c = self._color
        parts = [c.bold(self.prog)]

        # Show a simplified usage - just indicate options are available
        required_args = [a for a in self.arguments.values() if a.required]
        optional_count = len(self.arguments) - len(required_args)

        for arg in required_args:
            parts.append(f"{c.option('--' + arg.cli_name)} {c.type_hint('<value>')}")

        if optional_count > 0:
            parts.append(c.dim("[options]"))

        # Show positional argument if any
        positional = self._find_positional_argument()
        if positional:
            parts.append(c.dim(f"[{positional.cli_name}]"))

        return f"usage: {' '.join(parts)}"

    def _format_argument_help(
        self,
        arg: Argument,
        show_variants: bool = True,
        variant: Optional[str] = None,
        indent: int = 0,
    ) -> List[str]:
        """Format help text for a single argument with colors and type info.

        Format:
          --option <type> [default]
              Help text description
        """
        c = self._color
        lines = []
        prefix = " " * indent

        # Build argument string with type hint
        type_str = self._get_type_string(arg)

        # Build the option names (excluding shortcuts — shown separately)
        option_names = [c.option("--" + arg.cli_name)]
        if arg.aliases:
            for alias in arg.aliases:
                option_names.append(c.dim("--" + alias))

        # Get default value for display
        default_display = self._get_default_display(arg, variant, show_variants)

        if default_display:
            suffix = f" {c.dim('[' + default_display + ']')}"
        else:
            suffix = ""

        if arg.is_boolean:
            bool_choices = c.choices("{true,false}")
            arg_str = f"{prefix}  {', '.join(option_names)} {bool_choices}"
            arg_str += suffix
        elif arg.choices:
            choices_str = c.choices("{" + ",".join(arg.choices) + "}")
            arg_str = f"{prefix}  {', '.join(option_names)} {choices_str}"
            arg_str += suffix
        else:
            arg_str = f"{prefix}  {', '.join(option_names)} {c.type_hint(type_str)}"
            arg_str += suffix

        # Add the argument line
        lines.append(arg_str)

        # Description indent (aligned under the option)
        desc_prefix = prefix + "      "

        # Help text line
        if arg.help_text:
            help_text = arg.help_text
            # Strip any existing "only for" annotations
            if "only for:" in help_text.lower():
                help_text = re.sub(r"\s*\(only for:[^)]+\)", "", help_text).strip()

            lines.append(f"{desc_prefix}{help_text}")

        # Container access hints
        if arg.is_container and arg.container_info:
            info = arg.container_info
            has_dc_value = info.value_is_dataclass or (
                info.nested_container is not None
                and info.nested_container.value_is_dataclass
            )

            # Determine the immediate DC value type
            dc_type = info.value_type
            if info.nested_container is not None:
                dc_type = info.nested_container.value_type

            # Drill through single-container-field DC wrappers to find the
            # leaf dataclass (e.g., ResourceAllocation -> GPULocation)
            extra_path = ""
            leaf_dc = dc_type
            if has_dc_value and hasattr(dc_type, "__dataclass_fields__"):
                extra_path, leaf_dc = self._find_leaf_dc(dc_type)

            if info.container_kind == "list":
                if has_dc_value:
                    hint = (
                        f"use --{arg.cli_name}[i]{extra_path}.field"
                        f" or --config file.yaml"
                    )
                else:
                    hint = f"use --{arg.cli_name}[i] or pass multiple values"
            else:
                if has_dc_value:
                    hint = (
                        f"use --{arg.cli_name}[key]{extra_path}.field"
                        f" or --config file.yaml"
                    )
                else:
                    hint = f"use --{arg.cli_name}[key]=value or --config file.yaml"
            lines.append(f"{desc_prefix}{c.dim(hint)}")

            # Show available fields for the leaf dataclass
            if hasattr(leaf_dc, "__dataclass_fields__"):
                field_strs = []
                for f in leaf_dc.__dataclass_fields__.values():
                    ft = f.type
                    if isinstance(ft, str):
                        field_strs.append(f"{f.name} ({ft})")
                    else:
                        field_strs.append(f"{f.name} ({self._type_name(ft)})")
                lines.append(
                    f"{desc_prefix}{c.dim(leaf_dc.__name__ + ' fields: ' + ', '.join(field_strs))}"
                )

        # Show shortcuts as a mapping table
        if arg.shortcuts:
            # Calculate padding for alignment
            max_flag_len = max(len(f) for f in arg.shortcuts.values())
            max_val_len = max(len(str(e.value)) for e in arg.shortcuts.keys())
            lines.append(f"{desc_prefix}{c.dim('shortcuts:')}")
            for enum_val, flag in sorted(arg.shortcuts.items(), key=lambda x: x[1]):
                flag_str = c.option(f"-{flag}")
                val_str = str(enum_val.value)
                padding = " " * (
                    max_flag_len - len(flag) + max_val_len - len(val_str) + 2
                )
                help_text = getattr(enum_val, "help_text", "")
                if help_text:
                    lines.append(
                        f"{desc_prefix}  {flag_str}  {val_str}{padding}"
                        f"{c.dim(help_text)}"
                    )
                else:
                    lines.append(f"{desc_prefix}  {flag_str}  {val_str}")

        # Variant info on its own line if applicable
        if show_variants and arg.variants and isinstance(arg.default, dict):
            variant_list = ", ".join(sorted(arg.variants))
            lines.append(f"{desc_prefix}{c.dim('only for: ' + variant_list)}")

        return lines

    def _get_default_display(
        self, arg: Argument, variant: Optional[str], show_variants: bool
    ) -> str:
        """Get the default value for inline display."""

        def _fmt(val: Any) -> str:
            if isinstance(val, Enum):
                return str(val.value)
            return str(val)

        if isinstance(arg.default, dict):
            if variant and variant in arg.default:
                return _fmt(arg.default[variant])
            elif not show_variants:
                return _fmt(next(iter(arg.default.values())))
            else:
                # Multiple defaults - show first one
                return _fmt(next(iter(arg.default.values())))
        elif arg.default is not None:
            return _fmt(arg.default)
        # Show None for optional fields with no default
        elif not arg.required:
            return "None"
        return ""

    def _get_type_string(self, arg: Argument) -> str:
        """Get a human-readable type string for an argument."""
        if arg.is_boolean:
            return ""

        if arg.is_container and arg.container_info:
            return f"<{self._container_type_string(arg.container_info)}>"

        if arg.is_list:
            inner = self._type_name(arg.arg_type)
            return f"<{inner}...>"

        return f"<{self._type_name(arg.arg_type)}>"

    def _container_type_string(self, info: Any) -> str:
        """Build a human-readable type string from ContainerFieldInfo."""
        val_name = self._container_value_name(info)
        if info.container_kind == "list":
            return f"List[{val_name}]"
        key_name = self._type_name(info.key_type)
        return f"Dict[{key_name}, {val_name}]"

    def _container_value_name(self, info: Any) -> str:
        """Get the display name for a container's value type."""
        if info.nested_container is not None:
            return self._container_type_string(info.nested_container)
        if hasattr(info.value_type, "__name__"):
            return info.value_type.__name__
        return str(info.value_type)

    def _type_name(self, t: Union[type, Callable]) -> str:
        """Get the name of a type."""
        if t is int:
            return "int"
        if t is float:
            return "float"
        if t is str:
            return "str"
        if t is bool:
            return "bool"
        if isinstance(t, type) and issubclass(t, Enum):
            return "choice"
        # Handle generic aliases like List[X], Dict[K, V]
        args = getattr(t, "__args__", None)
        if args and hasattr(t, "__name__"):
            inner = ", ".join(self._type_name(a) for a in args)
            return f"{t.__name__}[{inner}]"
        if hasattr(t, "__name__"):
            name = t.__name__
            if name == "loads":  # json.loads
                return "json"
            return name
        return "value"

    @staticmethod
    def _find_leaf_dc(dc_type: Any) -> Tuple[str, Any]:
        """Walk through single-container-field dataclass wrappers.

        When a dataclass value type has exactly one field and that field is
        itself a container (List/Dict) whose values are dataclasses, drill
        down to find the innermost "leaf" dataclass and build the
        bracket-access path suffix along the way.

        Returns:
            (path_suffix, leaf_dc_type) — path_suffix is the extra bracket
            path (e.g. ".resource_mapping[i]"), leaf_dc_type is the
            innermost dataclass type.  If no drilling is needed, returns
            ("", dc_type).
        """
        path = ""
        current = dc_type
        visited: set = set()

        while hasattr(current, "__dataclass_fields__"):
            if id(current) in visited:
                break
            visited.add(id(current))

            dc_fields = list(current.__dataclass_fields__.values())
            if len(dc_fields) != 1:
                break

            f = dc_fields[0]
            ft = f.type
            # Strip Optional
            if get_origin(ft) is Union:
                non_none = [a for a in get_args(ft) if a is not type(None)]
                if len(non_none) == 1:
                    ft = non_none[0]

            origin = get_origin(ft)
            args = get_args(ft)
            if origin is list or origin is collections.abc.Sequence:
                path += f".{f.name}[i]"
                inner = args[0] if args else None
            elif origin is dict or origin is collections.abc.Mapping:
                path += f".{f.name}[key]"
                inner = args[1] if len(args) > 1 else None
            else:
                break

            if inner and hasattr(inner, "__dataclass_fields__"):
                current = inner
            else:
                break

        return path, current

    def _format_yaml_example(self) -> List[str]:
        """Generate an example YAML configuration snippet."""
        c = self._color
        lines = []

        lines.append(c.header("Example YAML Config:"))
        lines.append(c.dim("  # Save as config.yaml and use with --config config.yaml"))
        lines.append("")

        # Group arguments by their top-level group
        top_level_args = []
        nested_groups: Dict[str, List[Argument]] = {}

        for arg in self.arguments.values():
            if arg.positional:
                continue
            if "." in arg.cli_name:
                group_name = arg.cli_name.split(".")[0]
                if group_name not in nested_groups:
                    nested_groups[group_name] = []
                nested_groups[group_name].append(arg)
            else:
                top_level_args.append(arg)

        # Format top-level args
        for arg in top_level_args[:5]:  # Limit to first 5
            lines.append(self._format_yaml_arg(arg, indent=2))

        # Format one nested group as example
        if nested_groups:
            group_name = next(iter(nested_groups.keys()))
            group_args = nested_groups[group_name]

            lines.append(f"  {c.yaml_key(group_name + ':')}")
            for arg in group_args[:3]:  # Limit to first 3 per group
                # Get the field name (after the dot)
                field_name = arg.cli_name.split(".", 1)[1]
                default = self._get_example_value(arg)
                comment = (
                    f"  {c.yaml_comment('# ' + arg.help_text)}" if arg.help_text else ""
                )
                lines.append(
                    f"    {c.yaml_key(field_name + ':')} {c.yaml_value(str(default))}{comment}"
                )

        lines.append("")
        return lines

    def _format_yaml_arg(self, arg: Argument, indent: int = 0) -> str:
        """Format a single argument as YAML."""
        c = self._color
        prefix = " " * indent

        default = self._get_example_value(arg)
        comment = f"  {c.yaml_comment('# ' + arg.help_text)}" if arg.help_text else ""

        return f"{prefix}{c.yaml_key(arg.cli_name + ':')} {c.yaml_value(str(default))}{comment}"

    def _get_example_value(self, arg: Argument) -> Any:
        """Get an example value for YAML output."""
        if isinstance(arg.default, dict):
            val = next(iter(arg.default.values()))
            return val.value if isinstance(val, Enum) else val
        if arg.default is not None:
            return arg.default.value if isinstance(arg.default, Enum) else arg.default
        if arg.choices:
            return arg.choices[0]
        if arg.arg_type is int:
            return 0
        if arg.arg_type is float:
            return 0.0
        if arg.arg_type is bool:
            return False
        return '""'
