"""
Marziel OS — Process Manager.

AI-managed process lifecycle:

  QUEUED → RUNNING → COMPLETED | FAILED | KILLED

Each "process" is a task the AI is working on. Processes run
as neurons (SynapsePool workers) with resource tracking,
priority scheduling, and parent-child relationships.

Like Unix ps/top/kill but for AI tasks.
"""

import logging
import threading
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Optional

logger = logging.getLogger("marziel.process")


class ProcessState(Enum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    KILLED = "killed"


@dataclass
class Process:
    """A managed AI process — the OS unit of work."""

    pid: int
    name: str
    state: ProcessState = ProcessState.QUEUED
    priority: int = 5          # 0 (lowest) → 9 (highest)
    parent_pid: int = 0        # For process trees
    created_at: float = 0.0
    started_at: float = 0.0
    finished_at: float = 0.0
    tokens_used: int = 0
    memory_mb: float = 0.0
    result: object = None
    error: str = ""
    _func: Callable = field(default=None, repr=False)
    _args: tuple = field(default=(), repr=False)
    _kwargs: dict = field(default_factory=dict, repr=False)
    _thread: threading.Thread = field(default=None, repr=False)

    @property
    def duration_ms(self) -> float:
        if self.started_at == 0:
            return 0.0
        end = self.finished_at if self.finished_at > 0 else time.monotonic()
        return round((end - self.started_at) * 1000, 1)

    @property
    def is_alive(self) -> bool:
        return self.state in (ProcessState.QUEUED, ProcessState.RUNNING)

    def to_dict(self) -> dict:
        return {
            "pid": self.pid,
            "name": self.name,
            "state": self.state.value,
            "priority": self.priority,
            "parent_pid": self.parent_pid,
            "duration_ms": self.duration_ms,
            "tokens_used": self.tokens_used,
            "memory_mb": self.memory_mb,
            "error": self.error if self.state == ProcessState.FAILED else "",
        }


class ProcessManager:
    """
    AI Process Manager — lifecycle, monitoring, resource tracking.

    Manages a process table (like /proc in Linux) and provides
    Unix-like commands: spawn, kill, ps, top, wait.
    """

    def __init__(self, max_concurrent: int = 8):
        self._processes: dict[int, Process] = {}
        self._next_pid = 1
        self._lock = threading.Lock()
        self._max_concurrent = max_concurrent

        # Stats
        self._total_spawned = 0
        self._total_completed = 0
        self._total_failed = 0
        self._total_killed = 0

        logger.info(f"⚙️ Process Manager initialized (max_concurrent={max_concurrent})")

    def spawn(self, name: str, func: Callable, args: tuple = (),
              kwargs: dict = None, priority: int = 5,
              parent_pid: int = 0) -> int:
        """
        Spawn a new process. Returns PID.

        The function runs in a background thread. When done,
        the process state transitions to COMPLETED or FAILED.
        """
        with self._lock:
            pid = self._next_pid
            self._next_pid += 1
            self._total_spawned += 1

            proc = Process(
                pid=pid,
                name=name,
                state=ProcessState.QUEUED,
                priority=priority,
                parent_pid=parent_pid,
                created_at=time.monotonic(),
                _func=func,
                _args=args,
                _kwargs=kwargs or {},
            )
            self._processes[pid] = proc

        # Check if we can run immediately or need to queue
        running = sum(
            1 for p in self._processes.values()
            if p.state == ProcessState.RUNNING
        )

        if running < self._max_concurrent:
            self._start_process(pid)
        else:
            logger.debug(
                f"📋 Process {pid} ({name}) queued "
                f"(running={running}/{self._max_concurrent})"
            )

        return pid

    def _start_process(self, pid: int):
        """Start a queued process."""
        proc = self._processes.get(pid)
        if not proc or proc.state != ProcessState.QUEUED:
            return

        proc.state = ProcessState.RUNNING
        proc.started_at = time.monotonic()

        def _run():
            try:
                result = proc._func(*proc._args, **proc._kwargs)
                proc.result = result
                proc.state = ProcessState.COMPLETED
                with self._lock:
                    self._total_completed += 1
                logger.debug(f"✅ Process {pid} ({proc.name}) completed")
            except Exception as e:
                proc.error = str(e)
                proc.state = ProcessState.FAILED
                with self._lock:
                    self._total_failed += 1
                logger.warning(f"❌ Process {pid} ({proc.name}) failed: {e}")
            finally:
                proc.finished_at = time.monotonic()
                # Try to start queued processes
                self._dequeue_next()

        thread = threading.Thread(
            target=_run, name=f"marziel-proc-{pid}", daemon=True
        )
        proc._thread = thread
        thread.start()

        logger.debug(f"▶️ Process {pid} ({proc.name}) started (priority={proc.priority})")

    def _dequeue_next(self):
        """Start the next queued process by priority."""
        queued = [
            p for p in self._processes.values()
            if p.state == ProcessState.QUEUED
        ]
        if not queued:
            return

        running = sum(
            1 for p in self._processes.values()
            if p.state == ProcessState.RUNNING
        )
        if running >= self._max_concurrent:
            return

        # Sort by priority (highest first)
        queued.sort(key=lambda p: p.priority, reverse=True)
        self._start_process(queued[0].pid)

    def kill(self, pid: int) -> bool:
        """Kill a process by PID."""
        proc = self._processes.get(pid)
        if not proc:
            return False

        if not proc.is_alive:
            return False

        proc.state = ProcessState.KILLED
        proc.finished_at = time.monotonic()

        with self._lock:
            self._total_killed += 1

        logger.info(f"🔪 Process {pid} ({proc.name}) killed")

        # Dequeue next
        self._dequeue_next()
        return True

    def ps(self, include_dead: bool = False) -> list:
        """
        List processes (like Unix `ps`).

        Args:
            include_dead: Include completed/failed/killed processes.
        """
        processes = []
        for proc in self._processes.values():
            if not include_dead and not proc.is_alive:
                continue
            processes.append(proc.to_dict())

        processes.sort(key=lambda p: p["pid"])
        return processes

    def top(self) -> dict:
        """
        Real-time system stats (like Unix `top`).

        Returns aggregate resource usage and process summary.
        """
        all_procs = list(self._processes.values())
        running = [p for p in all_procs if p.state == ProcessState.RUNNING]
        queued = [p for p in all_procs if p.state == ProcessState.QUEUED]

        total_tokens = sum(p.tokens_used for p in all_procs)
        total_memory = sum(p.memory_mb for p in running)
        avg_duration = 0.0
        completed = [
            p for p in all_procs
            if p.state == ProcessState.COMPLETED and p.duration_ms > 0
        ]
        if completed:
            avg_duration = round(
                sum(p.duration_ms for p in completed) / len(completed), 1
            )

        return {
            "processes": {
                "running": len(running),
                "queued": len(queued),
                "total": len(all_procs),
                "max_concurrent": self._max_concurrent,
            },
            "resources": {
                "total_tokens": total_tokens,
                "active_memory_mb": round(total_memory, 2),
                "avg_duration_ms": avg_duration,
            },
            "lifetime": {
                "total_spawned": self._total_spawned,
                "total_completed": self._total_completed,
                "total_failed": self._total_failed,
                "total_killed": self._total_killed,
            },
            "active_processes": [p.to_dict() for p in running],
        }

    def wait(self, pid: int, timeout: float = 60.0) -> Optional[Process]:
        """Wait for a process to complete."""
        proc = self._processes.get(pid)
        if not proc:
            return None

        deadline = time.monotonic() + timeout
        while proc.is_alive and time.monotonic() < deadline:
            time.sleep(0.1)

        return proc

    def get_process(self, pid: int) -> Optional[Process]:
        """Get a process by PID."""
        return self._processes.get(pid)

    def cleanup(self, max_age_seconds: int = 3600):
        """Remove dead processes older than max_age."""
        cutoff = time.monotonic() - max_age_seconds
        with self._lock:
            to_remove = [
                pid for pid, proc in self._processes.items()
                if not proc.is_alive and proc.finished_at < cutoff
            ]
            for pid in to_remove:
                del self._processes[pid]

        if to_remove:
            logger.debug(f"🧹 Cleaned up {len(to_remove)} dead processes")

    @property
    def stats(self) -> dict:
        return self.top()
