"""Execute routed jobs on this node."""
from __future__ import annotations

import asyncio
import logging
import os
import re
import resource
import subprocess
import sys
import time
from typing import Any, Dict, List, Optional, Set

from mesh.models import JobInfo, JobResult, JobStatus, JobType

logger = logging.getLogger(__name__)

# Active jobs tracked by this executor
_active_jobs: Dict[str, asyncio.Task] = {}

# ── Command Sandbox ─────────────────────────────────────────────────────────

# Dangerous imports/calls that must never appear in command strings
_BLOCKED_PATTERNS: List[re.Pattern] = [
    re.compile(r'\bos\.system\b'),
    re.compile(r'\bos\.popen\b'),
    re.compile(r'\bos\.exec[a-z]*\b'),
    re.compile(r'\bos\.spawn[a-z]*\b'),
    re.compile(r'\bos\.remove\b'),
    re.compile(r'\bos\.unlink\b'),
    re.compile(r'\bos\.rmdir\b'),
    re.compile(r'\bos\.rename\b'),
    re.compile(r'\bsubprocess\b'),
    re.compile(r'\bshutil\.rmtree\b'),
    re.compile(r'\bshutil\.move\b'),
    re.compile(r'\b__import__\b'),
    re.compile(r'\beval\s*\('),
    re.compile(r'\bexec\s*\('),
    re.compile(r'\bcompile\s*\('),
    re.compile(r'\bopen\s*\(.*(w|a|x)'),  # write-mode file open
    re.compile(r'\bsocket\b'),
    re.compile(r'\bctypes\b'),
]

# Default allowed command prefixes (module paths). Extend via register_allowed_command().
_ALLOWED_COMMANDS: Set[str] = {
    "from mesh.",
    "import mesh.",
    "from atlas.",
    "import atlas.",
    "from kernels.",
    "import kernels.",
    "from optimization.",
    "import optimization.",
    "from profiling.",
    "import profiling.",
    "from product_runners.",
    "import product_runners.",
}

# Resource limits for sandboxed commands
_MAX_CPU_TIME_S = 7200      # 2 hours
_MAX_MEMORY_BYTES = 8 * 1024 * 1024 * 1024  # 8 GB


def register_allowed_command(prefix: str) -> None:
    """Register an additional allowed command prefix."""
    _ALLOWED_COMMANDS.add(prefix)


class CommandSandbox:
    """Validates commands against allowlist and blocks dangerous patterns."""

    @staticmethod
    def validate(command: str, node_id: str = "unknown") -> None:
        """Validate a command string. Raises ValueError if rejected.

        Checks:
        1. Command starts with an allowed prefix
        2. No blocked dangerous patterns
        3. No null bytes
        """
        if not command or not command.strip():
            raise ValueError("Empty command")

        if "\x00" in command:
            raise ValueError("Command contains null bytes")

        # Check allowed prefixes
        cmd_stripped = command.strip()
        allowed = any(cmd_stripped.startswith(prefix) for prefix in _ALLOWED_COMMANDS)
        if not allowed:
            logger.warning(
                "SECURITY: Blocked command from node %s — no matching allowlist prefix: %.100s",
                node_id, cmd_stripped,
            )
            raise ValueError(
                f"Command not in allowlist. Must start with one of: "
                f"{sorted(_ALLOWED_COMMANDS)}"
            )

        # Check blocked patterns
        for pattern in _BLOCKED_PATTERNS:
            if pattern.search(command):
                logger.warning(
                    "SECURITY: Blocked dangerous pattern '%s' in command from node %s: %.100s",
                    pattern.pattern, node_id, cmd_stripped,
                )
                raise ValueError(f"Command contains blocked pattern: {pattern.pattern}")

    @staticmethod
    def apply_resource_limits():
        """Apply resource limits to the child process (Linux only).

        Called as preexec_fn in subprocess creation.
        """
        try:
            # CPU time limit
            resource.setrlimit(resource.RLIMIT_CPU, (_MAX_CPU_TIME_S, _MAX_CPU_TIME_S))
            # Virtual memory limit
            resource.setrlimit(resource.RLIMIT_AS, (_MAX_MEMORY_BYTES, _MAX_MEMORY_BYTES))
            # Prevent core dumps
            resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
            # Limit number of child processes
            resource.setrlimit(resource.RLIMIT_NPROC, (64, 64))
        except (ValueError, OSError) as e:
            # Non-fatal — log and continue
            logger.debug("Could not set resource limits: %s", e)


_sandbox = CommandSandbox()


async def execute_job(job: JobInfo, node_id: str = "unknown") -> JobResult:
    """Execute a job and return the result."""
    t0 = time.time()
    logger.info("Executing job %s: type=%s node=%s", job.job_id, job.job_type, node_id)

    try:
        if job.job_type == JobType.PROBE:
            result = await _run_probe_job(job)
        elif job.job_type == JobType.BENCHMARK:
            result = await _run_command_job(job, node_id)
        elif job.job_type == JobType.TRAINING:
            result = await _run_command_job(job, node_id)
        elif job.job_type == JobType.INFERENCE:
            result = await _run_command_job(job, node_id)
        else:
            result = await _run_command_job(job, node_id)

        duration = time.time() - t0
        return JobResult(
            job_id=job.job_id,
            status=JobStatus.COMPLETED,
            result=result,
            duration_s=round(duration, 2),
        )
    except Exception as e:
        duration = time.time() - t0
        logger.error("Job %s failed: %s", job.job_id, e)
        return JobResult(
            job_id=job.job_id,
            status=JobStatus.FAILED,
            error=str(e),
            duration_s=round(duration, 2),
        )


async def _run_probe_job(job: JobInfo) -> Dict[str, Any]:
    """Run probe job using probe_runner."""
    from mesh.agent.probe_runner import run_probes
    iterations = job.args.get("iterations", 20)
    kernel_types = job.args.get("kernel_types")
    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(
        None, lambda: run_probes(iterations=iterations, kernel_types=kernel_types)
    )
    return result


async def _run_command_job(job: JobInfo, node_id: str = "unknown") -> Dict[str, Any]:
    """Run a shell command job with sandboxing."""
    if not job.command:
        return {"error": "No command specified"}

    # Validate command against sandbox
    _sandbox.validate(job.command, node_id=node_id)

    logger.info("AUDIT: Command execution node=%s job=%s cmd=%.200s",
                node_id, job.job_id, job.command)

    timeout = job.args.get("timeout_s", 3600)
    env = None
    if job.args.get("use_slento_cuda", False):
        env = _build_slento_cuda_env(job)
        logger.info("Slento CUDA enabled for job %s", job.job_id)

    proc = await asyncio.create_subprocess_exec(
        sys.executable, "-c", job.command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env,
        preexec_fn=CommandSandbox.apply_resource_limits,
    )
    try:
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError:
        proc.kill()
        return {"error": f"Job timed out after {timeout}s"}

    return {
        "returncode": proc.returncode,
        "stdout": stdout.decode()[-10000:],
        "stderr": stderr.decode()[-5000:],
    }


def get_active_job_count() -> int:
    return sum(1 for t in _active_jobs.values() if not t.done())


def _build_slento_cuda_env(job: JobInfo) -> Dict[str, str]:
    from mesh.agent.hardware_scanner import _detect_slento_cuda

    slento_config = {"slento_cuda_path": job.args.get("slento_cuda_path")}
    slento = _detect_slento_cuda(config=slento_config)
    if not slento["has_slento_cuda"]:
        raise ValueError("Slento CUDA requested but no compatible runtime was detected")

    lib_path = slento["slento_cuda_lib_path"]
    base_dir = os.path.dirname(lib_path)
    compat_libs = [
        os.path.join(lib_path, "libcuda.so.1"),
        os.path.join(lib_path, "libcudart.so.12"),
        os.path.join(lib_path, "libcublas.so.12"),
        os.path.join(lib_path, "libcublasLt.so.12"),
        os.path.join(lib_path, "libcudnn.so.9"),
    ]

    env = os.environ.copy()
    env["LD_PRELOAD"] = ":".join(compat_libs)
    ld_library_entries = [
        lib_path,
        os.path.join(base_dir, "lib"),
    ]
    existing_ld_library_path = env.get("LD_LIBRARY_PATH", "")
    if existing_ld_library_path:
        ld_library_entries.append(existing_ld_library_path)
    env["LD_LIBRARY_PATH"] = ":".join(ld_library_entries)
    env["HIP_VISIBLE_DEVICES"] = "0"
    env["ROCR_VISIBLE_DEVICES"] = "0"
    env["CUDA_VISIBLE_DEVICES"] = "0"
    env["MIOPEN_LOG_LEVEL"] = "3"
    return env
