"""Docker execution engine — runs jobs as containers with GPU isolation."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import Optional

from radix_edge.config import get_config
from radix_edge.db import get_db
from radix_edge.executor.gpu_allocator import release

logger = logging.getLogger("radix_edge.executor")


def _build_docker_command(
    image: str,
    command: Optional[list[str]],
    args: Optional[list[str]],
    env: Optional[dict[str, str]],
    gpu_indices: list[int],
    num_gpus: int,
    job_dir: str,
    job_id: str,
    timeout_seconds: int,
) -> list[str]:
    """Build the full docker run command with security hardening."""
    config = get_config()

    docker_cmd = [
        "docker", "run",
        "--rm",
        # --- Workspace ---
        "-v", f"{job_dir}:/workspace",
        "-w", "/workspace",
        # --- Security hardening ---
        "--user", "1000:1000",
        "--cap-drop=ALL",
        "--read-only",
        "--tmpfs", "/tmp:size=1g,noexec",
        "--security-opt=no-new-privileges",
        f"--pids-limit={config.job_pids_limit}",
        # --- Resource limits ---
        f"--memory={config.job_memory_limit}",
        f"--cpus={config.job_cpu_limit}",
        # --- Network isolation ---
        f"--network={config.job_network_mode}",
        # --- Labels for lifecycle management ---
        "--label", "radix.managed=true",
        "--label", f"radix.job_id={job_id}",
    ]

    # GPU isolation via CUDA_VISIBLE_DEVICES
    if gpu_indices:
        cuda_devices = ",".join(str(i) for i in gpu_indices)
        docker_cmd.extend(["--gpus", "all"])
        docker_cmd.extend(["-e", f"CUDA_VISIBLE_DEVICES={cuda_devices}"])

    # Environment variables
    if env:
        for key, value in env.items():
            docker_cmd.extend(["-e", f"{key}={value}"])

    docker_cmd.append(image)

    # Build the actual command to run inside the container
    run_cmd = list(command) if command else []

    # Wrap with torchrun for multi-GPU jobs
    if num_gpus > 1 and run_cmd:
        run_cmd = [
            "torchrun",
            f"--nproc_per_node={num_gpus}",
            "--standalone",
        ] + run_cmd

    if args:
        run_cmd.extend(args)

    docker_cmd.extend(run_cmd)
    return docker_cmd


def validate_image(image: str) -> None:
    """Validate that the Docker image matches allowed patterns.

    Raises ValueError if the image is not in the allowlist.
    """
    config = get_config()
    if not config.allowed_image_patterns:
        return  # No allowlist configured — permit all

    patterns = [p.strip() for p in config.allowed_image_patterns.split(",") if p.strip()]
    if not patterns:
        return

    for pattern in patterns:
        if re.match(pattern, image):
            return

    raise ValueError(
        f"Image '{image}' does not match any allowed pattern. "
        f"Allowed: {config.allowed_image_patterns}"
    )


async def run_job(job_id: str, gpu_indices: list[int]):
    """Execute a job as a Docker container. Updates job status throughout lifecycle.

    This is meant to be run as an asyncio task.
    """
    config = get_config()
    def now():
        return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")

    # Load job details
    with get_db() as db:
        row = db.execute("SELECT * FROM jobs WHERE job_id = ?", (job_id,)).fetchone()
        if not row:
            logger.error("Job %s not found in database", job_id)
            return

        job = dict(row)

    image = job["image"]
    command = json.loads(job["command"]) if job["command"] else None
    args = json.loads(job["args"]) if job["args"] else None
    env = json.loads(job["env"]) if job["env"] else None
    num_gpus = job["num_gpus"]
    timeout_seconds = job["timeout_seconds"] or config.default_timeout

    # Validate image against allowlist
    try:
        validate_image(image)
    except ValueError as e:
        with get_db() as db:
            db.execute(
                "UPDATE jobs SET status = 'failed', error_message = ?, completed_at = strftime('%Y-%m-%dT%H:%M:%fZ','now') WHERE job_id = ?",
                (str(e), job_id),
            )
        release(job_id)
        logger.warning("Job %s rejected: %s", job_id, e)
        return

    # Create job workspace
    job_dir = os.path.join(config.jobs_workdir, job_id)
    os.makedirs(job_dir, exist_ok=True)

    # Mark as running
    start_time = now()
    with get_db() as db:
        db.execute(
            "UPDATE jobs SET status = 'running', started_at = ?, scheduled_gpus = ? WHERE job_id = ?",
            (start_time, json.dumps(gpu_indices), job_id),
        )

    docker_cmd = _build_docker_command(
        image=image,
        command=command,
        args=args,
        env=env,
        gpu_indices=gpu_indices,
        num_gpus=num_gpus,
        job_dir=job_dir,
        job_id=job_id,
        timeout_seconds=timeout_seconds,
    )

    logger.info("Executing job %s: %s", job_id, " ".join(docker_cmd))

    exit_code = None
    stdout_data = ""
    stderr_data = ""
    error_message = None

    try:
        proc = await asyncio.create_subprocess_exec(
            *docker_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        try:
            stdout_bytes, stderr_bytes = await asyncio.wait_for(
                proc.communicate(),
                timeout=timeout_seconds,
            )
            exit_code = proc.returncode
            stdout_data = stdout_bytes.decode("utf-8", errors="replace")
            stderr_data = stderr_bytes.decode("utf-8", errors="replace")
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            exit_code = -1
            error_message = f"Job timed out after {timeout_seconds}s"
            logger.warning("Job %s timed out after %ds", job_id, timeout_seconds)

    except Exception as e:
        exit_code = -1
        error_message = str(e)
        logger.exception("Job %s execution error", job_id)

    # Compute runtime
    end_time = now()
    try:
        started = datetime.fromisoformat(start_time.replace("Z", "+00:00"))
        ended = datetime.fromisoformat(end_time.replace("Z", "+00:00"))
        runtime_seconds = (ended - started).total_seconds()
    except Exception:
        runtime_seconds = 0.0

    # Determine status
    status = "succeeded" if exit_code == 0 else "failed"
    if error_message is None and exit_code != 0:
        error_message = f"Container exited with code {exit_code}"

    # Truncate output tails (last 10KB)
    max_tail = 10240
    stdout_tail = stdout_data[-max_tail:] if stdout_data else None
    stderr_tail = stderr_data[-max_tail:] if stderr_data else None

    # Update job in database
    with get_db() as db:
        db.execute(
            """UPDATE jobs SET
                status = ?, exit_code = ?, error_message = ?,
                stdout_tail = ?, stderr_tail = ?,
                runtime_seconds = ?, completed_at = ?
            WHERE job_id = ?""",
            (status, exit_code, error_message, stdout_tail, stderr_tail, runtime_seconds, end_time, job_id),
        )

    # Release GPUs
    freed = release(job_id)
    logger.info(
        "Job %s %s (exit=%s, runtime=%.1fs, gpus=%s freed)",
        job_id, status, exit_code, runtime_seconds, freed,
    )

    # Clean up workspace (best effort)
    try:
        import shutil
        shutil.rmtree(job_dir, ignore_errors=True)
    except Exception:
        pass
