"""Bash 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 TyroBashCompletionGenerator(CompletionGenerator):
    """Generates bash completion scripts with embedded Python logic."""

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

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

        Returns:
            Bash 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 bash completion script.
        script = f"""# AUTOMATICALLY GENERATED by tyro
# Bash completion for {prog}

# Helper function to trim colon-containing prefixes from completions.
# This is necessary because bash treats ':' as a word separator.
# When user types "optimizer:ad", bash splits it as ["optimizer", ":", "ad"].
# We return "optimizer:adam", but bash thinks current word is just "ad",
# so it would append and create "optimizer:optimizer:adam" without this trimming.
__tyro_{root_prefix}_ltrim_colon_completions() {{
  local cur=$1
  # Only trim if the current word contains a colon.
  if [[ "$cur" == *:* ]]; then
    # Remove everything up to and including the last colon from each completion.
    local prefix=${{cur%${{cur##*:}}}}
    local i
    for i in "${{!COMPREPLY[@]}}"; do
      COMPREPLY[$i]=${{COMPREPLY[$i]#"$prefix"}}
    done
  fi
}}

# Main completion function.
_{root_prefix}() {{
  local cur="${{COMP_WORDS[COMP_CWORD]}}"

  # Reconstruct cur if it was split by bash on colons.
  # When typing "optimizer:ad", bash splits it as ["optimizer", ":", "ad"].
  # We need to reconstruct it as "optimizer:ad" for proper completion.
  # However, we must stop at word boundaries (spaces between complete args).
  local i=$COMP_CWORD
  while [[ $i -gt 0 ]]; do
    local prev="${{COMP_WORDS[i-1]}}"
    if [[ "$prev" == ":" ]]; then
      # Previous word is a colon, prepend it and continue.
      cur=":$cur"
      i=$((i-1))
    elif [[ "$prev" != "${{COMP_WORDS[0]}}" && "$prev" != -* ]]; then
      # Check if we're at a word boundary (complete colon-separated value).
      # If the word TWO positions back is a colon, then prev is a complete value.
      # Example: ["dataset", ":", "image-net", "optimizer", ":", "ad"]
      #                                         ^---------^
      #           When at "ad", prev="optimizer", two-back=":"
      #           But optimizer is a NEW argument, not part of image-net.
      if [[ $i -ge 3 && "${{COMP_WORDS[i-2]}}" == ":" ]]; then
        # Two words back is a colon, so prev completed a colon-separated value.
        # Don't merge prev - it's a separate argument.
        break
      fi
      # Previous word is part of the current colon-separated value, prepend it.
      cur="$prev$cur"
      i=$((i-1))
    else
      # Hit program name or an option, stop.
      break
    fi
  done

  # Get completions from Python.
  local python_cmd
  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
    COMPREPLY=()
    return 1
  fi

  # Pass raw COMP_WORDS array to Python for processing.
  # Python will handle word reconstruction (e.g., merging "dataset" ":" "mnist" -> "dataset:mnist").
  #
  # We read the embedded Python into a variable via a plain heredoc rather than
  # putting the heredoc directly inside the $(...) command substitution: bash
  # scans the whole substitution body (heredoc included) for balanced quotes,
  # so a spec string with an odd number of quotes (e.g. a default like "don't")
  # would otherwise break the entire script. Feeding the heredoc to `read` has
  # no such issue. (`read -d ''` returns non-zero at EOF, hence `|| true`.)
  local _tyro_py
  IFS= read -r -d '' _tyro_py << 'PYTHON_EOF' || true
# Hardcoded completion spec.
COMPLETION_SPEC = {spec_repr}

{python_code}
PYTHON_EOF
  local completions
  completions=$(printf '%s' "$_tyro_py" | "$python_cmd" - "${{COMP_WORDS[@]}}" "$COMP_CWORD")

  # Check for special path completion marker.
  if [[ "$completions" == "__TYRO_COMPLETE_FILES__" ]]; then
    COMPREPLY=($(compgen -f -- "$cur"))
    return 0
  fi

  # Parse completions and build COMPREPLY array.
  local IFS=$'\\n'
  local -a matches
  while IFS= read -r line; do
    # Extract the completion part (before the tab).
    local completion="${{line%%$'\\t'*}}"
    matches+=("$completion")
  done <<< "$completions"

  COMPREPLY=("${{matches[@]}}")

  # Trim colon-containing prefixes to handle bash's word splitting on colons.
  __tyro_{root_prefix}_ltrim_colon_completions "$cur"

  return 0
}}

# Register completion for the full program name.
complete -F _{root_prefix} {prog}

# Also register for basename to handle ./script invocations.
_prog_basename=$(basename "{prog}")
if [ "$_prog_basename" != "{prog}" ]; then
  complete -F _{root_prefix} "$_prog_basename"
fi
"""
        return script
