"""Zsh completion generation for tyro using embedded Python logic."""

from __future__ import annotations

from ... import _parsers
from ._base import CompletionGenerator
from ._python_completer import get_embedded_code
from ._spec import build_completion_spec


class TyroZshCompletionGenerator(CompletionGenerator):
    """Generates zsh completion scripts with embedded Python logic."""

    def generate(
        self,
        parser_spec: _parsers.ParserSpecification,
        prog: str,
        root_prefix: str,
    ) -> str:
        """Generate a zsh completion script.

        Args:
            parser_spec: Parser specification to generate completion for.
            prog: Program name.
            root_prefix: Prefix for completion function names.

        Returns:
            Zsh completion script as a string.
        """
        # Build completion spec.
        spec = build_completion_spec(parser_spec, prog)
        spec_repr = repr(spec)

        # Get embedded Python code.
        python_code = get_embedded_code()

        # Generate the zsh completion script.
        script = f"""#compdef {prog}

# AUTOMATICALLY GENERATED by tyro

# Main completion function.
_{root_prefix}() {{
  local python_cmd
  # Try to find python3 or python.
  if command -v python3 >/dev/null 2>&1; then
    python_cmd="python3"
  elif command -v python >/dev/null 2>&1; then
    python_cmd="python"
  else
    return 1
  fi

  # Get completions (format: "completion:description").
  local completions
  completions=$("$python_cmd" - "${{words[@]}}" "$CURRENT" << 'PYTHON_EOF'
# Hardcoded completion spec.
COMPLETION_SPEC = {spec_repr}

{python_code}
PYTHON_EOF
)

  # Check for special path completion marker.
  if [[ "$completions" == "__TYRO_COMPLETE_FILES__" ]]; then
    _files
    return
  fi

  # Parse completions into an array for _describe.
  local -a completion_array
  while IFS= read -r line; do
    # Convert tab-separated "completion\tdescription" to colon-separated for _describe.
    # zsh _describe expects "completion:description" format.
    local completion="${{line%%$'\\t'*}}"
    local description="${{line#*$'\\t'}}"
    # Escape colons in the completion string so they aren't treated as separators.
    local completion_escaped="${{completion//:/\\\\:}}"
    completion_array+=("$completion_escaped:$description")
  done <<< "$completions"

  # Use _describe to show completions with descriptions.
  _describe '{prog} completions' completion_array
}}

typeset -A opt_args
_{root_prefix} "$@"
"""
        return script
