"""Overlay-based container building for local development.

Replicates autocont's overlay patterns in pure Python for use with hpctainers.
Instead of a full apptainer build, this creates a writable overlay on top of
a read-only base SIF, bind-mounts host repos, and runs %post scripts via
apptainer exec.
"""

from __future__ import annotations

import logging
import re
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, List, Optional

from pydantic import BaseModel

logger = logging.getLogger(__name__)


class OverlayExecutor(BaseModel, extra="forbid"):
    """Wraps apptainer exec with --fakeroot --overlay and bind args.

    Mirrors autocont's OverlayExecutor but in pure Python (no xonsh).
    """

    base_sif: Path
    overlay_img: Path
    additional_cmd_args: Optional[List[str]] = None

    def cmd_args(self) -> List[str]:
        """Create the apptainer cmd args for executing with an overlay."""
        if self.additional_cmd_args is None:
            return ["--fakeroot", "--overlay", str(self.overlay_img), str(self.base_sif)]
        return [
            "--fakeroot",
            *self.additional_cmd_args,
            "--overlay", str(self.overlay_img),
            str(self.base_sif),
        ]

    def add_bind(self, host_path: str, container_path: Optional[str] = None) -> OverlayExecutor:
        """Create a new executor with an additional bind mount.

        Args:
            host_path: Path on the host to bind
            container_path: Optional path inside the container (defaults to host_path)

        Returns:
            New OverlayExecutor with the bind added
        """
        if container_path:
            bind_arg = f"{host_path}:{container_path}"
        else:
            bind_arg = host_path
        new_args = ["--bind", bind_arg]

        return OverlayExecutor(
            base_sif=self.base_sif,
            overlay_img=self.overlay_img,
            additional_cmd_args=(
                new_args
                if self.additional_cmd_args is None
                else [*self.additional_cmd_args, *new_args]
            ),
        )


def create_overlay(path: Path, size_mb: int = 1024) -> Path:
    """Create a sparse overlay image at the given path.

    Args:
        path: Where to create the overlay image
        size_mb: Size of the overlay in megabytes

    Returns:
        Path to the created overlay image

    Raises:
        RuntimeError: If overlay creation fails
    """
    logger.info(f"Creating overlay: {path} ({size_mb} MB)")

    if path.exists():
        logger.info(f"Removing existing overlay: {path}")
        path.unlink()

    path.parent.mkdir(parents=True, exist_ok=True)

    result = subprocess.run(
        ["apptainer", "overlay", "create", "--fakeroot", "--sparse", "--size", str(size_mb), str(path)],
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode != 0:
        raise RuntimeError(
            f"Failed to create overlay: {result.stderr}"
        )

    logger.info(f"Overlay created: {path}")
    return path


def exec_in_overlay(
    executor: OverlayExecutor,
    cmd: str,
    cwd: Optional[str] = None,
) -> subprocess.CompletedProcess:
    """Execute a command inside the overlay container.

    Args:
        executor: The OverlayExecutor to use
        cmd: Shell command to execute
        cwd: Optional working directory inside the container

    Returns:
        CompletedProcess with stdout/stderr

    Raises:
        RuntimeError: If the command fails
    """
    full_cmd = ["apptainer", "exec", "-e", "-c", *executor.cmd_args()]

    if cwd:
        shell_cmd = f"cd {cwd} && {cmd}"
    else:
        shell_cmd = cmd

    full_cmd.extend(["bash", "-c", shell_cmd])

    logger.debug(f"Executing in overlay: {' '.join(full_cmd)}")

    result = subprocess.run(
        full_cmd,
        capture_output=True,
        text=True,
        check=False,
    )

    if result.returncode != 0:
        logger.error(f"Command failed (exit {result.returncode}): {cmd}")
        if result.stderr:
            logger.error(f"stderr: {result.stderr}")
        raise RuntimeError(
            f"Overlay exec failed (exit {result.returncode}): {cmd}\n"
            f"stderr: {result.stderr}"
        )

    return result


def link_repo(
    executor: OverlayExecutor,
    host_path: str,
    container_path: str,
) -> OverlayExecutor:
    """Bind-mount a host repo and symlink it inside the overlay.

    Args:
        executor: The OverlayExecutor to use
        host_path: Path to the repo on the host
        container_path: Where it should appear inside the container

    Returns:
        New OverlayExecutor with the bind mount added
    """
    logger.info(f"Linking repo: {host_path} -> {container_path}")

    # Add the bind mount
    executor_with_bind = executor.add_bind(host_path)

    # Create the symlink inside the overlay
    parent_dir = str(Path(container_path).parent)
    exec_in_overlay(
        executor_with_bind,
        f"mkdir -p {parent_dir} && ln -sf {host_path} {container_path}",
    )

    logger.info(f"Repo linked: {container_path} -> {host_path}")
    return executor_with_bind


def extract_post_section(def_file: Path) -> str:
    """Extract the %post section from an Apptainer .def file.

    Handles %post with optional flags (e.g., %post -c /bin/bash).

    Args:
        def_file: Path to the .def file

    Returns:
        Content of the %post section as a string

    Raises:
        FileNotFoundError: If def_file doesn't exist
        ValueError: If no %post section found
    """
    if not def_file.exists():
        raise FileNotFoundError(f"Definition file not found: {def_file}")

    content = def_file.read_text()
    lines = content.split('\n')

    in_post = False
    post_lines = []
    section_pattern = re.compile(r'^%\w+')

    for line in lines:
        if re.match(r'^%post\b', line):
            in_post = True
            continue
        elif in_post and section_pattern.match(line):
            break
        elif in_post:
            post_lines.append(line)

    if not post_lines:
        raise ValueError(f"No %post section found in {def_file}")

    return '\n'.join(post_lines)


def substitute_build_args(script: str, args: Dict[str, str]) -> str:
    """Replace {{ KEY }} template variables in a script with actual values.

    Handles both {{ KEY }} and {{KEY}} patterns, case-insensitive matching.

    Args:
        script: The script content with template variables
        args: Dictionary mapping arg names to values

    Returns:
        Script with substitutions applied
    """
    result = script
    for key, value in args.items():
        # Match {{ KEY }}, {{KEY}}, and variations with different spacing
        pattern = r'\{\{\s*' + re.escape(key) + r'\s*\}\}'
        result = re.sub(pattern, value, result, flags=re.IGNORECASE)
    return result


def run_post_script(
    executor: OverlayExecutor,
    def_file: Path,
    build_args: Dict[str, str],
    bind_paths: Optional[Dict[str, str]] = None,
) -> subprocess.CompletedProcess:
    """Extract %post from .def, substitute args, and execute via overlay.

    Args:
        executor: The OverlayExecutor to use
        def_file: Path to the .def file
        build_args: Build arguments for template substitution
        bind_paths: Optional dict of host_path -> container_path for bind mounts

    Returns:
        CompletedProcess from the execution
    """
    logger.info(f"Running %post from {def_file}")

    # Extract and substitute
    post_script = extract_post_section(def_file)
    post_script = substitute_build_args(post_script, build_args)

    # Add bind mounts to executor
    exec_with_binds = executor
    if bind_paths:
        for host_path, container_path in bind_paths.items():
            exec_with_binds = exec_with_binds.add_bind(host_path, container_path)

    # Write script to temp file and execute
    with tempfile.NamedTemporaryFile(mode='w', suffix='.sh', delete=False) as f:
        f.write("#!/bin/bash\nset -e\n")
        f.write(post_script)
        script_path = f.name

    try:
        # Bind the script into the container and execute
        exec_with_script = exec_with_binds.add_bind(script_path)

        full_cmd = [
            "apptainer", "exec", "-e", "-c",
            *exec_with_script.cmd_args(),
            "bash", script_path,
        ]

        logger.debug(f"Running post script: {' '.join(full_cmd)}")

        result = subprocess.run(
            full_cmd,
            text=True,
            check=False,
        )

        if result.returncode != 0:
            raise RuntimeError(
                f"Post script execution failed (exit {result.returncode})"
            )

        logger.info("Post script completed successfully")
        return result

    finally:
        Path(script_path).unlink(missing_ok=True)


def finalize_overlay(executor: OverlayExecutor, output_sif: Path) -> Path:
    """Embed the overlay into a standalone SIF file.

    Copies the base SIF, runs e2fsck + resize2fs to minimize overlay,
    then adds the overlay to the SIF via apptainer sif add.

    Args:
        executor: The OverlayExecutor containing base_sif and overlay_img
        output_sif: Path for the output SIF file

    Returns:
        Path to the finalized SIF

    Raises:
        RuntimeError: If any step fails
    """
    logger.info(f"Finalizing overlay into: {output_sif}")

    output_sif.parent.mkdir(parents=True, exist_ok=True)

    # Copy base SIF
    logger.info(f"Copying base SIF: {executor.base_sif} -> {output_sif}")
    shutil.copyfile(executor.base_sif, output_sif)

    # Check and repair filesystem
    logger.info("Running e2fsck on overlay")
    result = subprocess.run(
        ["e2fsck", "-f", "-y", str(executor.overlay_img)],
        capture_output=True,
        text=True,
        check=False,
    )
    # e2fsck returns 1 if corrections were made, which is OK
    if result.returncode > 1:
        raise RuntimeError(f"e2fsck failed: {result.stderr}")

    # Minimize overlay size
    logger.info("Resizing overlay to minimum")
    result = subprocess.run(
        ["resize2fs", "-M", str(executor.overlay_img)],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        raise RuntimeError(f"resize2fs failed: {result.stderr}")

    # Embed overlay into SIF
    logger.info("Embedding overlay into SIF")
    result = subprocess.run(
        [
            "apptainer", "sif", "add",
            "--datatype", "4",
            "--partfs", "2",
            "--parttype", "4",
            "--partarch", "2",
            "--groupid", "1",
            str(output_sif),
            str(executor.overlay_img),
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        raise RuntimeError(f"apptainer sif add failed: {result.stderr}")

    logger.info(f"Finalized SIF: {output_sif}")
    return output_sif
