"""Zsh shell completion generator."""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ..models import CommandCompletion

__all__ = ["generate_zsh"]


def _build_command_descriptions(commands: dict[str, CommandCompletion]) -> str:
    """Build the command descriptions block for zsh."""
    cmd_descs: list[str] = []
    for cmd_name, cmd in sorted(commands.items()):
        desc = cmd.description.replace("'", "'\\''") if cmd.description else cmd_name
        # Add subcommand hint if applicable
        if cmd.subcommands and not cmd.args:
            subcmds = "|".join(sorted(cmd.subcommands.keys()))
            desc = f"<{subcmds}> {desc}" if desc else f"<{subcmds}>"
        cmd_descs.append(f"        '{cmd_name}:{desc}'")
    return "\n".join(cmd_descs)


def _build_subcommand_case(cmd_name: str, cmd: CommandCompletion) -> str:
    """Build zsh case statement for a command with subcommands."""
    subcmd_descs: list[str] = []
    for subcmd_name, subcmd in sorted(cmd.subcommands.items()):
        subdesc = subcmd.description.replace("'", "'\\''") if subcmd.description else subcmd_name
        subcmd_descs.append(f"'{subcmd_name}:{subdesc}'")
    subcmd_desc_str = " ".join(subcmd_descs)

    return f"""                {cmd_name})
                    local -a subcmds=({subcmd_desc_str})
                    if [[ $CURRENT -eq 2 ]]; then
                        _describe 'subcommand' subcmds
                    fi
                    ;;"""


def _build_help_case(cmd: CommandCompletion) -> str:
    """Build zsh case for help command with hierarchical completion.

    Position 1: complete with all commands (reuses the commands array)
    Position 2+: complete subcommands based on the command at position 1
    """
    # Build case statements for commands that have subcommands
    subcmd_cases: list[str] = []
    for parent_name, parent_cmd in sorted(cmd.subcommands.items()):
        if parent_cmd.args:
            # Build subcommand descriptions for this parent
            subcmds_str = " ".join(f"'{val}'" for val in parent_cmd.args[0].values)
            subcmd_cases.append(f"                            {parent_name}) compadd {subcmds_str} ;;")

    subcmd_block = "\n".join(subcmd_cases) if subcmd_cases else "                            *) ;;"

    return f"""                help)
                    if [[ $CURRENT -eq 2 ]]; then
                        # Position 1: complete with all commands
                        _describe 'command' commands
                    else
                        # Position 2+: complete subcommands based on previous word
                        case $words[2] in
{subcmd_block}
                        esac
                    fi
                    ;;"""


def _build_args_case(cmd_name: str, cmd: CommandCompletion) -> str | None:
    """Build zsh case statement for a command with positional args."""
    arg_specs: list[str] = []
    for arg in cmd.args:
        pos = arg.position
        desc = arg.description.replace("'", "'\\''")

        if arg.completion_type in ("choices", "dynamic", "literal"):
            values_str = " ".join(arg.values)
            arg_specs.append(f"'{pos}:{desc}:({values_str})'")
        elif arg.completion_type == "file":
            arg_specs.append(f"'{pos}:{desc}:_files'")
        elif arg.completion_type == "hint":
            # Show description but no actual completions
            hint = arg.values[0] if arg.values else desc
            arg_specs.append(f"'{pos}:{hint}:'")

    if not arg_specs:
        return None

    args_line = " \\\n                        ".join(arg_specs)
    return f"""                {cmd_name})
                    _arguments \\
                        {args_line}
                    ;;"""


def generate_zsh(commands: dict[str, CommandCompletion]) -> str:
    """Generate zsh completion script content.

    Args:
        commands: Dict mapping command name -> CommandCompletion

    Returns:
        The zsh completion script content
    """
    cmd_desc_block = _build_command_descriptions(commands)

    # Build case statements for each command
    case_statements: list[str] = []
    for cmd_name, cmd in sorted(commands.items()):
        if cmd_name == "help":
            case_statements.append(_build_help_case(cmd))
        elif cmd.subcommands:
            case_statements.append(_build_subcommand_case(cmd_name, cmd))
        elif cmd.args:
            case_stmt = _build_args_case(cmd_name, cmd)
            if case_stmt:
                case_statements.append(case_stmt)

    case_block = "\n".join(case_statements) if case_statements else "                *) ;;"

    return f"""#compdef pypr
# Zsh completion for pypr
# Generated by: pypr compgen zsh

_pypr() {{
    local -a commands=(
{cmd_desc_block}
    )

    _arguments -C \\
        '1:command:->command' \\
        '*::arg:->args'

    case $state in
        command)
            _describe 'command' commands
            ;;
        args)
            case $words[1] in
{case_block}
            esac
            ;;
    esac
}}

_pypr "$@"
"""
