"""Simple argument key normalization.

Converts argument keys between alias and canonical forms using snippet metadata.
This is the ONLY transformation - no serialization, no format changes.
"""

from typing import Any, Dict, Optional


def build_alias_map(snippet_metadata: Dict[str, Any]) -> Dict[str, str]:
    """Build mapping from alias to canonical name.

    Args:
        snippet_metadata: Snippet metadata with 'arguments' list

    Returns:
        Dict mapping alias/long_alias to canonical name
        Example: {'i': 'input', 'o': 'out'}
    """
    alias_map = {}

    if not snippet_metadata or "arguments" not in snippet_metadata:
        return alias_map

    for arg in snippet_metadata["arguments"]:
        canonical = arg["name"]

        # Short alias: -i
        if arg.get("alias"):
            alias_map[arg["alias"]] = canonical

        # Long alias (rare): --sv instead of --svfile
        if arg.get("long_alias"):
            alias_map[arg["long_alias"]] = canonical

    return alias_map


def normalize_argument_key(key: str, alias_map: Dict[str, str]) -> str:
    """Normalize a single argument key to canonical form with -- prefix.

    Args:
        key: Argument key from YAML (e.g., '-i', '--input', 'input')
        alias_map: Mapping from alias to canonical name

    Returns:
        Canonical form with -- prefix for CLI args, unchanged for batch arg names
    """
    # Batch arg names have no prefix (e.g., 'input_batch') - keep as-is
    if not key.startswith("-"):
        return key

    # Strip any prefix
    stripped = key.lstrip("-")

    # Map alias to canonical if it exists
    canonical = alias_map.get(stripped, stripped)

    # Return with -- prefix (canonical long form)
    return f"--{canonical}"


def normalize_argument_keys(
    arguments: Dict[str, Any], alias_map: Dict[str, str]
) -> Dict[str, Any]:
    """Normalize all argument keys in a dict to canonical form.

    Recursively handles composite_arg result_arguments and batch_list_arg values.

    Args:
        arguments: Dict of argument key -> value from YAML
        alias_map: Mapping from alias to canonical name

    Returns:
        New dict with normalized keys
    """
    normalized = {}

    for key, value in arguments.items():
        # Normalize the key
        new_key = normalize_argument_key(key, alias_map)

        # Handle value recursively if needed
        if isinstance(value, dict):
            # Check if it's a composite_arg or batch_list_arg
            if value.get("type") == "composite_arg" and isinstance(
                value.get("value"), dict
            ):
                # Recursively normalize result_arguments
                composite_value = value["value"].copy()
                if "result_arguments" in composite_value:
                    # Need to fetch metadata for the composite's snippet to get its aliases
                    # For now, keep as-is - caller will handle nested normalization
                    composite_value["result_arguments"] = value["value"][
                        "result_arguments"
                    ]
                normalized[new_key] = {
                    "type": "composite_arg",
                    "value": composite_value,
                }
            elif value.get("type") == "batch_list_arg" and isinstance(
                value.get("value"), dict
            ):
                # Recursively normalize keys in batch list dict
                batch_value = value["value"]
                normalized_batch = normalize_argument_keys(batch_value, alias_map)
                normalized[new_key] = {
                    "type": "batch_list_arg",
                    "value": normalized_batch,
                }
            else:
                # Other dict values (batch_file_arg, boolean with action, etc.)
                normalized[new_key] = value
        else:
            # Simple value (string, list, etc.)
            normalized[new_key] = value

    return normalized


def normalize_step_arguments(step: Dict[str, Any], snippet_fetcher) -> Dict[str, Any]:
    """Normalize all argument keys in a step, including nested composites.

    Args:
        step: Step dict with 'name' and 'arguments'
        snippet_fetcher: Function that takes snippet name and returns metadata

    Returns:
        New step dict with normalized argument keys
    """
    # Build alias map for this step's snippet
    snippet_name = step["name"]
    snippet_metadata = snippet_fetcher(snippet_name)
    alias_map = build_alias_map(snippet_metadata)

    # Normalize top-level arguments
    normalized_args = normalize_argument_keys(step["arguments"], alias_map)

    # Recursively normalize composite result_arguments
    normalized_args = _normalize_composite_recursive(normalized_args, snippet_fetcher)

    return {**step, "arguments": normalized_args}


def _normalize_composite_recursive(
    arguments: Dict[str, Any], snippet_fetcher
) -> Dict[str, Any]:
    """Recursively normalize composite result_arguments.

    Args:
        arguments: Dict of arguments (possibly containing composites)
        snippet_fetcher: Function to fetch snippet metadata

    Returns:
        Dict with all nested composites normalized
    """
    result = {}

    for key, value in arguments.items():
        if isinstance(value, dict) and value.get("type") == "composite_arg":
            composite_value = value["value"]
            if "result_arguments" in composite_value:
                # Fetch metadata for the composite's snippet
                composite_snippet = composite_value.get("snippet_name")
                if composite_snippet:
                    composite_metadata = snippet_fetcher(composite_snippet)
                    composite_alias_map = build_alias_map(composite_metadata)

                    # Normalize result_arguments keys
                    normalized_result_args = normalize_argument_keys(
                        composite_value["result_arguments"], composite_alias_map
                    )

                    # Recursively normalize nested composites
                    normalized_result_args = _normalize_composite_recursive(
                        normalized_result_args, snippet_fetcher
                    )

                    result[key] = {
                        "type": "composite_arg",
                        "value": {
                            **composite_value,
                            "result_arguments": normalized_result_args,
                        },
                    }
                else:
                    result[key] = value
            else:
                result[key] = value
        else:
            result[key] = value

    return result
