import json
import ctypes
import os
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from collections import deque
from datetime import datetime, timezone
from threading import RLock


def _now_iso():
    return datetime.now(timezone.utc).isoformat()


class SandboxRuntimeState:
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    CANCELED = "canceled"
    TIMEOUT = "timeout"
    FAILED = "failed"


class SandboxRuntimeManager:
    """
    Sandboxed evolution runtime manager with subprocess isolation.
    Applies platform-specific hard limits where available.
    """

    def __init__(self, max_runs=500, max_runtime_seconds=20, max_memory_mb=512, max_cpu_seconds=10):
        self._max_runs = max_runs
        self._max_runtime_seconds = max_runtime_seconds
        self._max_memory_mb = max_memory_mb
        self._max_cpu_seconds = max_cpu_seconds
        self._lock = RLock()
        self._queue = deque()
        self._runs = {}
        self._active = None  # {"sandbox_run_id", "proc", "started_monotonic", "timeout_s", "cwd", "job_handle"}
        self._events = deque(maxlen=2000)
        self._stats = {
            "started": 0,
            "canceled": 0,
            "completed": 0,
            "failed": 0,
            "timeouts": 0,
            "isolation_applied": 0,
            "isolation_degraded": 0,
        }

    def start_run(self, run_spec):
        if not isinstance(run_spec, dict):
            raise ValueError("run_spec must be an object")
        constraints = self._normalize_constraints(run_spec.get("constraints") or {})
        timeout_s = self._resolve_timeout(constraints)

        sandbox_run_id = f"sbox-{uuid.uuid4().hex[:10]}"
        now = _now_iso()
        run = {
            "sandbox_run_id": sandbox_run_id,
            "state": SandboxRuntimeState.QUEUED,
            "created_at": now,
            "started_at": None,
            "finished_at": None,
            "objective": run_spec.get("objective", ""),
            "task_id": run_spec.get("task_id"),
            "chunk_id": run_spec.get("chunk_id"),
            "agent_id": run_spec.get("agent_id"),
            "constraints": dict(constraints),
            "timeout_seconds": timeout_s,
            "snapshot_ref": run_spec.get("snapshot_ref"),
            "result": None,
            "error": None,
        }

        with self._lock:
            self._runs[sandbox_run_id] = run
            self._queue.append(sandbox_run_id)
            self._stats["started"] += 1
            self._trim_old_runs_locked()
        return sandbox_run_id

    def get_result(self, sandbox_run_id):
        with self._lock:
            run = self._runs.get(sandbox_run_id)
            return dict(run) if run else None

    def cancel(self, sandbox_run_id):
        with self._lock:
            run = self._runs.get(sandbox_run_id)
            if not run:
                return False
            if run["state"] in {
                SandboxRuntimeState.COMPLETED,
                SandboxRuntimeState.CANCELED,
                SandboxRuntimeState.TIMEOUT,
                SandboxRuntimeState.FAILED,
            }:
                return False
            if self._active and self._active["sandbox_run_id"] == sandbox_run_id:
                proc = self._active["proc"]
                if proc.poll() is None:
                    proc.terminate()
            run["state"] = SandboxRuntimeState.CANCELED
            run["finished_at"] = _now_iso()
            self._stats["canceled"] += 1
            self._events.append({"type": "canceled", "sandbox_run_id": sandbox_run_id, "timestamp": _now_iso()})
            return True

    def process_next(self):
        """
        Cooperative non-blocking step:
        - polls active subprocess-backed sandbox run
        - enforces run timeout
        - launches queued run if slot is available
        """
        self._poll_active()
        self._launch_next_if_idle()

    def drain_events(self, max_events=100):
        with self._lock:
            events = []
            for _ in range(min(max_events, len(self._events))):
                events.append(self._events.popleft())
            return events

    def get_stats(self):
        with self._lock:
            queued = sum(1 for r in self._runs.values() if r["state"] == SandboxRuntimeState.QUEUED)
            running = sum(1 for r in self._runs.values() if r["state"] == SandboxRuntimeState.RUNNING)
            timeouts = sum(1 for r in self._runs.values() if r["state"] == SandboxRuntimeState.TIMEOUT)
            return {
                "total_runs": len(self._runs),
                "queued_runs": queued,
                "running_runs": running,
                "timeout_runs": timeouts,
                "active_process": bool(self._active),
                **self._stats,
            }

    def _trim_old_runs_locked(self):
        if len(self._runs) <= self._max_runs:
            return
        ordered = sorted(self._runs.values(), key=lambda r: r.get("created_at") or "")
        drop_count = len(self._runs) - self._max_runs
        for run in ordered[:drop_count]:
            self._runs.pop(run["sandbox_run_id"], None)

    def _resolve_timeout(self, constraints):
        requested = constraints.get("max_runtime_seconds", self._max_runtime_seconds)
        try:
            timeout_s = int(requested)
        except Exception as exc:
            raise ValueError(f"invalid max_runtime_seconds: {exc}") from exc
        if timeout_s < 1:
            timeout_s = 1
        if timeout_s > self._max_runtime_seconds:
            timeout_s = self._max_runtime_seconds
        return timeout_s

    def _normalize_constraints(self, constraints):
        if not isinstance(constraints, dict):
            raise ValueError("constraints must be an object")

        allowed = {
            "max_runtime_seconds",
            "max_memory_mb",
            "max_cpu_seconds",
            "max_open_files",
            "no_network",
            "read_only_fs",
            "simulate_work_ms",
        }
        unknown = [k for k in constraints.keys() if k not in allowed]
        if unknown:
            raise ValueError(f"unsupported constraints: {unknown}")

        no_network = constraints.get("no_network", True)
        read_only_fs = constraints.get("read_only_fs", True)
        if no_network is not True:
            raise ValueError("no_network must be true in sandbox policy")
        if read_only_fs is not True:
            raise ValueError("read_only_fs must be true in sandbox policy")

        memory = constraints.get("max_memory_mb", self._max_memory_mb)
        cpu = constraints.get("max_cpu_seconds", self._max_cpu_seconds)
        max_open_files = constraints.get("max_open_files", 64)
        runtime = constraints.get("max_runtime_seconds", self._max_runtime_seconds)
        simulate_ms = constraints.get("simulate_work_ms", 0)

        try:
            memory = int(memory)
            cpu = int(cpu)
            max_open_files = int(max_open_files)
            runtime = int(runtime)
            simulate_ms = int(simulate_ms)
        except Exception as exc:
            raise ValueError(f"invalid numeric constraint: {exc}") from exc

        if memory < 64:
            memory = 64
        if memory > self._max_memory_mb:
            memory = self._max_memory_mb
        if cpu < 1:
            cpu = 1
        if cpu > self._max_cpu_seconds:
            cpu = self._max_cpu_seconds
        if runtime < 1:
            runtime = 1
        if runtime > self._max_runtime_seconds:
            runtime = self._max_runtime_seconds
        if max_open_files < 16:
            max_open_files = 16
        if max_open_files > 256:
            max_open_files = 256
        if simulate_ms < 0:
            simulate_ms = 0

        return {
            "max_runtime_seconds": runtime,
            "max_memory_mb": memory,
            "max_cpu_seconds": cpu,
            "max_open_files": max_open_files,
            "no_network": True,
            "read_only_fs": True,
            "simulate_work_ms": simulate_ms,
        }

    def _launch_next_if_idle(self):
        with self._lock:
            if self._active is not None:
                return
            sandbox_run_id = self._queue.popleft() if self._queue else None
            if not sandbox_run_id:
                return
            run = self._runs.get(sandbox_run_id)
            if not run or run["state"] != SandboxRuntimeState.QUEUED:
                return

            payload = {
                "sandbox_run_id": sandbox_run_id,
                "objective": run.get("objective"),
                "constraints": run.get("constraints", {}),
                "snapshot_ref": run.get("snapshot_ref"),
                "task_id": run.get("task_id"),
                "chunk_id": run.get("chunk_id"),
                "agent_id": run.get("agent_id"),
            }
            sandbox_cwd = tempfile.mkdtemp(prefix=f"hive_sbox_{sandbox_run_id}_")
            env = _build_worker_env()
            env["PYTHONNOUSERSITE"] = "1"
            env["HIVE_SANDBOX"] = "1"
            env["PYTHONIOENCODING"] = "utf-8"
            env["PYTHONUNBUFFERED"] = "1"
            worker_path = os.path.join(os.path.dirname(__file__), "sandbox_worker.py")
            popen_kwargs = {
                "stdin": subprocess.PIPE,
                "stdout": subprocess.PIPE,
                "stderr": subprocess.PIPE,
                "text": True,
                "cwd": sandbox_cwd,
                "env": env,
            }
            if os.name == "posix":
                popen_kwargs["preexec_fn"] = self._build_posix_preexec(run.get("constraints", {}))
            if os.name == "nt":
                creationflags = 0
                creationflags |= getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
                creationflags |= getattr(subprocess, "CREATE_NO_WINDOW", 0)
                creationflags |= 0x01000000  # CREATE_BREAKAWAY_FROM_JOB
                popen_kwargs["creationflags"] = creationflags
            proc = subprocess.Popen(
                [sys.executable, "-I", worker_path],
                **popen_kwargs,
            )
            proc.stdin.write(json.dumps(payload, separators=(",", ":")))
            proc.stdin.close()
            job_handle = None
            if os.name == "nt":
                job_handle = self._assign_windows_job_limits(proc, run.get("constraints", {}))
                if job_handle:
                    self._stats["isolation_applied"] += 1
                    self._events.append(
                        {
                            "type": "isolation_applied",
                            "sandbox_run_id": sandbox_run_id,
                            "timestamp": _now_iso(),
                            "backend": "windows_job_object",
                        }
                    )
                else:
                    self._stats["isolation_degraded"] += 1
                    self._events.append(
                        {
                            "type": "isolation_degraded",
                            "sandbox_run_id": sandbox_run_id,
                            "timestamp": _now_iso(),
                            "backend": "windows_job_object",
                        }
                    )
            elif os.name == "posix":
                self._stats["isolation_applied"] += 1
                self._events.append(
                    {
                        "type": "isolation_applied",
                        "sandbox_run_id": sandbox_run_id,
                        "timestamp": _now_iso(),
                        "backend": "posix_setrlimit",
                    }
                )
            run["state"] = SandboxRuntimeState.RUNNING
            run["started_at"] = _now_iso()
            self._active = {
                "sandbox_run_id": sandbox_run_id,
                "proc": proc,
                "started_monotonic": time.monotonic(),
                "timeout_s": int(run.get("timeout_seconds", self._max_runtime_seconds)),
                "cwd": sandbox_cwd,
                "job_handle": job_handle,
            }
            self._events.append({"type": "started", "sandbox_run_id": sandbox_run_id, "timestamp": _now_iso()})

    def _poll_active(self):
        with self._lock:
            active = dict(self._active) if self._active else None
        if not active:
            return

        sandbox_run_id = active["sandbox_run_id"]
        proc = active["proc"]
        run = self.get_result(sandbox_run_id)
        if not run:
            with self._lock:
                self._active = None
            return

        if proc.poll() is None:
            elapsed = time.monotonic() - active["started_monotonic"]
            if elapsed > active["timeout_s"]:
                proc.kill()
                proc.communicate(timeout=0.2)
                with self._lock:
                    live = self._runs.get(sandbox_run_id)
                    if live and live["state"] == SandboxRuntimeState.RUNNING:
                        live["state"] = SandboxRuntimeState.TIMEOUT
                        live["finished_at"] = _now_iso()
                        live["error"] = f"timeout>{active['timeout_s']}s"
                        self._stats["timeouts"] += 1
                        self._events.append(
                            {"type": "timeout", "sandbox_run_id": sandbox_run_id, "timestamp": _now_iso()}
                        )
                    self._active = None
                self._close_windows_job_handle(active.get("job_handle"))
                self._cleanup_sandbox_dir(active.get("cwd"))
            return

        stdout, stderr = proc.communicate(timeout=0.2)
        if proc.returncode != 0:
            with self._lock:
                live = self._runs.get(sandbox_run_id)
                if live and live["state"] == SandboxRuntimeState.RUNNING:
                    live["state"] = SandboxRuntimeState.FAILED
                    live["finished_at"] = _now_iso()
                    live["error"] = (stderr or f"worker_exit_{proc.returncode}").strip()[:1000]
                    self._stats["failed"] += 1
                    self._events.append({"type": "failed", "sandbox_run_id": sandbox_run_id, "timestamp": _now_iso()})
                self._active = None
            self._close_windows_job_handle(active.get("job_handle"))
            self._cleanup_sandbox_dir(active.get("cwd"))
            return

        try:
            raw_result = json.loads((stdout or "").strip() or "{}")
        except Exception:
            raw_result = {"artifact_only": True, "message": "worker returned non-json output"}

        sanitized = self._sanitize_worker_result(raw_result)
        with self._lock:
            live = self._runs.get(sandbox_run_id)
            if live and live["state"] == SandboxRuntimeState.RUNNING:
                live["state"] = SandboxRuntimeState.COMPLETED
                live["finished_at"] = _now_iso()
                live["result"] = sanitized
                self._stats["completed"] += 1
                self._events.append({"type": "completed", "sandbox_run_id": sandbox_run_id, "timestamp": _now_iso()})
            self._active = None
        self._close_windows_job_handle(active.get("job_handle"))
        self._cleanup_sandbox_dir(active.get("cwd"))

    def _sanitize_worker_result(self, raw):
        if not isinstance(raw, dict):
            raw = {}
        allowed_keys = {"artifact_only", "proposal_candidates", "metrics", "traces", "message", "meta"}
        sanitized = {k: raw.get(k) for k in allowed_keys if k in raw}
        sanitized["artifact_only"] = True
        if "proposal_candidates" not in sanitized or not isinstance(sanitized["proposal_candidates"], list):
            sanitized["proposal_candidates"] = []
        if "metrics" not in sanitized or not isinstance(sanitized["metrics"], dict):
            sanitized["metrics"] = {}
        if "traces" not in sanitized or not isinstance(sanitized["traces"], list):
            sanitized["traces"] = []
        return sanitized

    def _cleanup_sandbox_dir(self, path):
        if not path:
            return
        try:
            shutil.rmtree(path, ignore_errors=True)
        except Exception:
            pass

    def _build_posix_preexec(self, constraints):
        def _preexec():
            import os as _os
            import resource as _resource

            mem_bytes = int(constraints.get("max_memory_mb", self._max_memory_mb)) * 1024 * 1024
            cpu_s = int(constraints.get("max_cpu_seconds", self._max_cpu_seconds))
            nofile = int(constraints.get("max_open_files", 64))

            _os.setsid()
            if hasattr(_resource, "RLIMIT_AS"):
                _resource.setrlimit(_resource.RLIMIT_AS, (mem_bytes, mem_bytes))
            if hasattr(_resource, "RLIMIT_CPU"):
                _resource.setrlimit(_resource.RLIMIT_CPU, (cpu_s, cpu_s))
            if hasattr(_resource, "RLIMIT_NOFILE"):
                _resource.setrlimit(_resource.RLIMIT_NOFILE, (nofile, nofile))
            if hasattr(_resource, "RLIMIT_NPROC"):
                _resource.setrlimit(_resource.RLIMIT_NPROC, (1, 1))

        return _preexec

    def _assign_windows_job_limits(self, proc, constraints):
        try:
            kernel32 = ctypes.windll.kernel32
        except Exception:
            return None

        class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure):
            _fields_ = [
                ("PerProcessUserTimeLimit", ctypes.c_longlong),
                ("PerJobUserTimeLimit", ctypes.c_longlong),
                ("LimitFlags", ctypes.c_uint32),
                ("MinimumWorkingSetSize", ctypes.c_size_t),
                ("MaximumWorkingSetSize", ctypes.c_size_t),
                ("ActiveProcessLimit", ctypes.c_uint32),
                ("Affinity", ctypes.c_size_t),
                ("PriorityClass", ctypes.c_uint32),
                ("SchedulingClass", ctypes.c_uint32),
            ]

        class IO_COUNTERS(ctypes.Structure):
            _fields_ = [
                ("ReadOperationCount", ctypes.c_uint64),
                ("WriteOperationCount", ctypes.c_uint64),
                ("OtherOperationCount", ctypes.c_uint64),
                ("ReadTransferCount", ctypes.c_uint64),
                ("WriteTransferCount", ctypes.c_uint64),
                ("OtherTransferCount", ctypes.c_uint64),
            ]

        class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure):
            _fields_ = [
                ("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION),
                ("IoInfo", IO_COUNTERS),
                ("ProcessMemoryLimit", ctypes.c_size_t),
                ("JobMemoryLimit", ctypes.c_size_t),
                ("PeakProcessMemoryUsed", ctypes.c_size_t),
                ("PeakJobMemoryUsed", ctypes.c_size_t),
            ]

        JOB_OBJECT_LIMIT_PROCESS_TIME = 0x00000002
        JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x00000008
        JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x00000100
        JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000
        JOB_OBJECT_EXTENDED_LIMIT_INFORMATION_CLASS = 9

        job = kernel32.CreateJobObjectW(None, None)
        if not job:
            return None

        info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
        info.BasicLimitInformation.LimitFlags = (
            JOB_OBJECT_LIMIT_PROCESS_TIME
            | JOB_OBJECT_LIMIT_ACTIVE_PROCESS
            | JOB_OBJECT_LIMIT_PROCESS_MEMORY
            | JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
        )
        cpu_s = int(constraints.get("max_cpu_seconds", self._max_cpu_seconds))
        info.BasicLimitInformation.PerProcessUserTimeLimit = int(cpu_s * 10_000_000)
        info.BasicLimitInformation.ActiveProcessLimit = 1
        info.ProcessMemoryLimit = int(constraints.get("max_memory_mb", self._max_memory_mb)) * 1024 * 1024

        ok = kernel32.SetInformationJobObject(
            job,
            JOB_OBJECT_EXTENDED_LIMIT_INFORMATION_CLASS,
            ctypes.byref(info),
            ctypes.sizeof(info),
        )
        if not ok:
            kernel32.CloseHandle(job)
            return None

        ok = kernel32.AssignProcessToJobObject(job, ctypes.c_void_p(int(proc._handle)))
        if not ok:
            kernel32.CloseHandle(job)
            return None
        return job

    def _close_windows_job_handle(self, job_handle):
        if not job_handle:
            return
        try:
            ctypes.windll.kernel32.CloseHandle(job_handle)
        except Exception:
            pass


def _build_worker_env():
    # Build a minimal environment for the sandbox worker to reduce ambient authority.
    allowed_keys = {
        "SYSTEMROOT",
        "WINDIR",
        "PATH",
        "PATHEXT",
        "COMSPEC",
        "TMP",
        "TEMP",
        "HOME",
        "USERPROFILE",
        "LANG",
        "LC_ALL",
    }
    src = os.environ
    env = {}
    for key in allowed_keys:
        value = src.get(key)
        if value is not None:
            env[key] = value
    return env
