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

import asyncio
import logging
import os
import re
import platform
import subprocess

try:
    import resource
except ImportError:
    resource = None  # Windows — no POSIX resource limits
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/macOS only).

        Called as preexec_fn in subprocess creation. Skipped on Windows.
        """
        if resource is None:
            return
        try:
            resource.setrlimit(resource.RLIMIT_CPU, (_MAX_CPU_TIME_S, _MAX_CPU_TIME_S))
            resource.setrlimit(resource.RLIMIT_AS, (_MAX_MEMORY_BYTES, _MAX_MEMORY_BYTES))
            resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
            resource.setrlimit(resource.RLIMIT_NPROC, (64, 64))
        except (ValueError, OSError) as e:
            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)

    extra_kwargs = {}
    if platform.system() != "Windows":
        extra_kwargs["preexec_fn"] = CommandSandbox.apply_resource_limits
    proc = await asyncio.create_subprocess_exec(
        sys.executable, "-c", job.command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env,
        **extra_kwargs,
    )
    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
