"""Container wrapper script generation."""

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from jinja2 import Environment, StrictUndefined


class WrapperError(Exception):
    """Raised when wrapper generation fails."""


@dataclass(slots=True)
class WrapperConfig:
    """Configuration for wrapper script generation."""

    image_ref: str
    commands: list[str]
    runtime: str  # "singularity" or "docker"
    output_dir: Path
    image_cache: Optional[Path] = None
    extra_mounts: Optional[list[str]] = None
    env_passthrough: Optional[list[str]] = None
    gpu: bool = False

    def __post_init__(self):
        if self.extra_mounts is None:
            self.extra_mounts = []
        if self.env_passthrough is None:
            self.env_passthrough = []


def _sanitize_image_name(image_ref: str) -> str:
    """
    Convert image reference to safe filename.
    
    ghcr.io/lab/myenv:1.0 -> lab_myenv_1.0
    """
    # Remove registry prefix
    without_registry = re.sub(r'^[^/]+\.(io|com|org)/', '', image_ref)
    # Replace special chars with underscores
    sanitized = re.sub(r'[/:@]', '_', without_registry)
    return sanitized


def _singularity_wrapper_template() -> str:
    """Return Singularity wrapper script template."""
    return '''#!/bin/bash
# Auto-generated by absconda wrap
# Image: {{ image_ref }}
# Command: {{ command }}
set -euo pipefail

{% if image_cache %}
SIF_CACHE="{{ image_cache }}"
{% else %}
SIF_CACHE="${HOME}/.local/absconda/sif-cache"
{% endif %}
SIF_FILE="${SIF_CACHE}/{{ sif_filename }}"
IMAGE_REF="docker://{{ image_ref }}"

# Pull SIF if missing
if [[ ! -f "$SIF_FILE" ]]; then
    mkdir -p "$SIF_CACHE"
    echo "Pulling Singularity image to cache..." >&2
    singularity pull "$SIF_FILE" "$IMAGE_REF"
fi

# Build mount arguments
MOUNTS=()
{% for mount in mounts %}
MOUNTS+=("-B" "{{ mount }}")
{% endfor %}

{% if gpu %}
# GPU support
GPU_FLAG="--nv"
{% else %}
GPU_FLAG=""
{% endif %}

# Execute command in container
exec singularity exec \\
    $GPU_FLAG \\
    "${MOUNTS[@]}" \\
    "$SIF_FILE" \\
    {{ command }} "$@"
'''


def _docker_wrapper_template() -> str:
    """Return Docker wrapper script template."""
    return '''#!/bin/bash
# Auto-generated by absconda wrap
# Image: {{ image_ref }}
# Command: {{ command }}
set -euo pipefail

IMAGE_REF="{{ image_ref }}"

# Build mount arguments
MOUNTS=()
{% for mount in mounts %}
MOUNTS+=("-v" "{{ mount }}:{{ mount }}")
{% endfor %}

{% if gpu %}
# GPU support
GPU_FLAG="--gpus all"
{% else %}
GPU_FLAG=""
{% endif %}

# Execute command in container
exec docker run --rm -i \\
    $GPU_FLAG \\
    "${MOUNTS[@]}" \\
    -w "$PWD" \\
    -u "$(id -u):$(id -g)" \\
    "$IMAGE_REF" \\
    {{ command }} "$@"
'''


def generate_wrappers(config: WrapperConfig) -> dict[str, Path]:
    """
    Generate wrapper scripts for specified commands.
    
    Returns a dict mapping command names to their wrapper script paths.
    """
    if not config.commands:
        raise WrapperError("No commands specified for wrapping")
    
    if config.runtime not in ("singularity", "docker"):
        raise WrapperError(f"Unsupported runtime: {config.runtime}")
    
    # Create output directory
    config.output_dir.mkdir(parents=True, exist_ok=True)
    
    # Select template based on runtime
    if config.runtime == "singularity":
        template_str = _singularity_wrapper_template()
        sif_filename = f"{_sanitize_image_name(config.image_ref)}.sif"
    else:
        template_str = _docker_wrapper_template()
        sif_filename = None
    
    # Prepare Jinja2 environment
    env = Environment(
        undefined=StrictUndefined,
        trim_blocks=True,
        lstrip_blocks=True,
        keep_trailing_newline=True,
    )
    template = env.from_string(template_str)
    
    # Generate wrappers
    wrapper_paths = {}
    for command in config.commands:
        context = {
            "image_ref": config.image_ref,
            "command": command,
            "mounts": config.extra_mounts,
            "gpu": config.gpu,
            "sif_filename": sif_filename,
            "image_cache": str(config.image_cache) if config.image_cache else None,
        }
        
        wrapper_content = template.render(**context)
        wrapper_path = config.output_dir / command
        wrapper_path.write_text(wrapper_content, encoding="utf-8")
        wrapper_path.chmod(0o755)  # Make executable
        
        wrapper_paths[command] = wrapper_path
    
    return wrapper_paths


def expand_mount_paths(mount_specs: list[str]) -> list[str]:
    """
    Expand environment variables in mount path specifications.
    
    Note: Expansion happens at wrapper generation time for validation,
    but wrappers will re-expand at runtime for dynamic values like $PWD.
    """
    expanded = []
    for spec in mount_specs:
        # Keep the original spec with env vars for runtime expansion
        # But validate that required env vars exist now
        if "$" in spec:
            # Extract env var names for validation
            env_vars = re.findall(r'\$(\w+)', spec)
            for var in env_vars:
                if var not in os.environ and var not in ('PWD', 'HOME', 'USER'):
                    # Warn but don't fail - these might be set at runtime
                    pass
        expanded.append(spec)
    return expanded
