"""
Marziel OS — Kernel.

The AI Kernel is the heartbeat of Marziel OS. It runs a persistent
event loop that:

  1. Listens for events (user input, data streams, timers, anomalies)
  2. Evaluates events through the AI scheduler
  3. Spawns processes via the Process Manager
  4. Manages memory consolidation
  5. Emits signals for inter-component communication
  6. Maintains system health via periodic heartbeat

Unlike a traditional OS kernel, the scheduler is AI-driven —
it uses the MarzielEngine to evaluate task priority and the
FuzzyController to adaptively manage resources.
"""

import json
import logging
import os
import queue
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Callable, Optional

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


# ══════════════════════════════════════════════════════════════
# EVENT SYSTEM
# ══════════════════════════════════════════════════════════════

class EventType(Enum):
    USER_INPUT = "user_input"       # User sent a message
    DATA_STREAM = "data_stream"     # External data arrived
    TIMER = "timer"                 # Scheduled timer fired
    ANOMALY = "anomaly"             # Anomaly detected
    SYSTEM = "system"               # System event (boot, shutdown, etc.)
    PROCESS = "process"             # Process state changed
    SIGNAL = "signal"               # Inter-component signal


@dataclass(order=True)
class Event:
    """A kernel event with priority ordering."""
    priority: int                             # Lower = higher priority
    timestamp: float = field(compare=False)
    event_type: EventType = field(compare=False)
    data: dict = field(compare=False, default_factory=dict)
    source: str = field(compare=False, default="")

    def to_dict(self) -> dict:
        return {
            "priority": self.priority,
            "type": self.event_type.value,
            "source": self.source,
            "data": self.data,
            "timestamp": self.timestamp,
        }


class EventQueue:
    """Priority event queue for the kernel."""

    def __init__(self, max_size: int = 1024):
        self._queue = queue.PriorityQueue(maxsize=max_size)
        self._total_events = 0
        self._lock = threading.Lock()

    def push(self, event: Event):
        """Push an event onto the queue."""
        try:
            self._queue.put_nowait(event)
            with self._lock:
                self._total_events += 1
        except queue.Full:
            logger.warning("⚠️ Event queue full — dropping lowest priority event")

    def pop(self, timeout: float = 1.0) -> Optional[Event]:
        """Pop the highest-priority event. Returns None on timeout."""
        try:
            return self._queue.get(timeout=timeout)
        except queue.Empty:
            return None

    @property
    def size(self) -> int:
        return self._queue.qsize()

    @property
    def total_events(self) -> int:
        with self._lock:
            return self._total_events


# ══════════════════════════════════════════════════════════════
# SIGNAL BUS — Inter-component communication
# ══════════════════════════════════════════════════════════════

class SignalBus:
    """
    Signal bus for inter-component communication.

    Components can subscribe to signal types and receive callbacks
    when those signals are emitted. Like Unix signals but semantic.
    """

    def __init__(self):
        self._handlers: dict[str, list[Callable]] = {}
        self._lock = threading.Lock()
        self._signal_count = 0

    def subscribe(self, signal_type: str, handler: Callable):
        """Subscribe a handler to a signal type."""
        with self._lock:
            if signal_type not in self._handlers:
                self._handlers[signal_type] = []
            self._handlers[signal_type].append(handler)

    def emit(self, signal_type: str, data: dict = None):
        """Emit a signal to all subscribers."""
        with self._lock:
            self._signal_count += 1
            handlers = list(self._handlers.get(signal_type, []))

        for handler in handlers:
            try:
                handler(signal_type, data or {})
            except Exception as e:
                logger.warning(f"Signal handler error ({signal_type}): {e}")

    @property
    def stats(self) -> dict:
        with self._lock:
            return {
                "total_signals": self._signal_count,
                "registered_types": list(self._handlers.keys()),
                "total_handlers": sum(
                    len(h) for h in self._handlers.values()
                ),
            }


# ══════════════════════════════════════════════════════════════
# HEARTBEAT — Periodic health check
# ══════════════════════════════════════════════════════════════

class Heartbeat:
    """
    Periodic pulse — checks system health, triggers maintenance.

    Runs every `interval` seconds and executes health checks:
    - Memory consolidation
    - Process cleanup
    - Resource monitoring
    """

    def __init__(self, interval: int = 60):
        self.interval = interval
        self._running = False
        self._thread = None
        self._tick_count = 0
        self._callbacks: list[Callable] = []
        self._started_at = 0.0

    def register(self, callback: Callable):
        """Register a callback to run on each heartbeat."""
        self._callbacks.append(callback)

    def start(self):
        """Start the heartbeat thread."""
        if self._running:
            return
        self._running = True
        self._started_at = time.monotonic()
        self._thread = threading.Thread(
            target=self._loop, name="marziel-heartbeat", daemon=True
        )
        self._thread.start()
        logger.info(f"💓 Heartbeat started (interval={self.interval}s)")

    def stop(self):
        """Stop the heartbeat."""
        self._running = False
        if self._thread:
            self._thread.join(timeout=5)
        logger.info("💓 Heartbeat stopped")

    def _loop(self):
        while self._running:
            time.sleep(self.interval)
            if not self._running:
                break
            self._tick_count += 1
            for cb in self._callbacks:
                try:
                    cb(self._tick_count)
                except Exception as e:
                    logger.warning(f"Heartbeat callback error: {e}")

    @property
    def stats(self) -> dict:
        return {
            "running": self._running,
            "interval_sec": self.interval,
            "tick_count": self._tick_count,
            "uptime_sec": round(
                time.monotonic() - self._started_at, 1
            ) if self._started_at > 0 else 0,
            "callbacks": len(self._callbacks),
        }


# ══════════════════════════════════════════════════════════════
# SCHEDULED TASKS
# ══════════════════════════════════════════════════════════════

@dataclass
class ScheduledTask:
    """A task scheduled to run at regular intervals."""
    name: str
    func: Callable
    interval_seconds: int
    last_run: float = 0.0
    run_count: int = 0
    enabled: bool = True
    args: tuple = ()
    kwargs: dict = field(default_factory=dict)

    @property
    def next_run(self) -> float:
        if self.last_run == 0:
            return 0  # Run immediately
        return self.last_run + self.interval_seconds

    @property
    def is_due(self) -> bool:
        return time.monotonic() >= self.next_run

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "interval_sec": self.interval_seconds,
            "run_count": self.run_count,
            "enabled": self.enabled,
            "next_run_in": max(0, round(self.next_run - time.monotonic(), 1)),
        }


# ══════════════════════════════════════════════════════════════
# MARZIEL KERNEL — The AI OS Core
# ══════════════════════════════════════════════════════════════

class MarzielKernel:
    """
    The Marziel OS Kernel — persistent event loop with autonomous
    decision-making.

    Coordinates:
    - Event processing (prioritized)
    - Process scheduling (AI-driven)
    - Memory management (consolidation, recall)
    - System health (heartbeat)
    - Inter-component signals

    Boot sequence:
    1. Initialize subsystems
    2. Register heartbeat callbacks
    3. Start event loop
    4. Ready for user input
    """

    VERSION = "1.0.0"

    def __init__(self, memory_manager=None, process_manager=None,
                 heartbeat_interval: int = 60):
        # Core subsystems
        self.events = EventQueue()
        self.signals = SignalBus()
        self.heartbeat = Heartbeat(interval=heartbeat_interval)
        self.memory = memory_manager
        self.processes = process_manager

        # Scheduled tasks
        self._scheduled_tasks: dict[str, ScheduledTask] = {}

        # Kernel state
        self._running = False
        self._boot_time = 0.0
        self._event_loop_thread = None
        self._scheduler_thread = None

        # Event handlers by type
        self._event_handlers: dict[str, list[Callable]] = {}

    def boot(self):
        """
        Boot the Marziel OS kernel.

        Initializes all subsystems and starts the event loop.
        """
        self._boot_time = time.monotonic()
        self._running = True

        logger.info("🔄 Marziel OS kernel booting...")

        # Register heartbeat callbacks
        self.heartbeat.register(self._on_heartbeat)
        self.heartbeat.start()

        # Start event loop
        self._event_loop_thread = threading.Thread(
            target=self._event_loop, name="marziel-kernel", daemon=True
        )
        self._event_loop_thread.start()

        # Start scheduler
        self._scheduler_thread = threading.Thread(
            target=self._scheduler_loop, name="marziel-scheduler", daemon=True
        )
        self._scheduler_thread.start()

        # Emit boot signal
        self.signals.emit("kernel.boot", {
            "version": self.VERSION,
            "timestamp": datetime.utcnow().isoformat() + "Z",
        })

        # Boot event
        self.events.push(Event(
            priority=0,
            timestamp=time.monotonic(),
            event_type=EventType.SYSTEM,
            data={"action": "boot", "version": self.VERSION},
            source="kernel",
        ))

        logger.info(
            f"✅ Marziel OS kernel v{self.VERSION} booted "
            f"(heartbeat={self.heartbeat.interval}s)"
        )

    def shutdown(self):
        """Graceful kernel shutdown."""
        logger.info("🛑 Marziel OS kernel shutting down...")
        self._running = False

        # Emit shutdown signal
        self.signals.emit("kernel.shutdown", {
            "uptime_sec": self.uptime,
        })

        # Stop heartbeat
        self.heartbeat.stop()

        # Consolidate memory before shutdown
        if self.memory:
            self.memory.consolidate()

        # Clean up processes
        if self.processes:
            self.processes.cleanup(max_age_seconds=0)

        logger.info("🛑 Marziel OS kernel shutdown complete")

    # ── Event Loop ──

    def _event_loop(self):
        """Main kernel event loop — process events as they arrive."""
        while self._running:
            event = self.events.pop(timeout=1.0)
            if event is None:
                continue

            try:
                self._handle_event(event)
            except Exception as e:
                logger.error(f"Event handler error: {e}")

    def _handle_event(self, event: Event):
        """Dispatch an event to registered handlers."""
        event_type = event.event_type.value
        handlers = self._event_handlers.get(event_type, [])
        for handler in handlers:
            try:
                handler(event)
            except Exception as e:
                logger.warning(
                    f"Event handler failed ({event_type}): {e}"
                )

        # Log to episodic memory
        if self.memory and event.event_type != EventType.SYSTEM:
            self.memory.working.store(
                f"event:{event.event_type.value}:{int(event.timestamp)}",
                event.to_dict(),
                ttl_seconds=300,  # 5-minute TTL
            )

    def on_event(self, event_type: str, handler: Callable):
        """Register a handler for a specific event type."""
        if event_type not in self._event_handlers:
            self._event_handlers[event_type] = []
        self._event_handlers[event_type].append(handler)

    def push_event(self, event_type: EventType, data: dict = None,
                   priority: int = 5, source: str = ""):
        """Push an event into the kernel event queue."""
        self.events.push(Event(
            priority=priority,
            timestamp=time.monotonic(),
            event_type=event_type,
            data=data or {},
            source=source,
        ))

    # ── Scheduler ──

    def _scheduler_loop(self):
        """Check scheduled tasks and run them when due."""
        while self._running:
            time.sleep(5)
            if not self._running:
                break

            for task in self._scheduled_tasks.values():
                if not task.enabled or not task.is_due:
                    continue

                try:
                    # Run the task via process manager if available
                    if self.processes:
                        self.processes.spawn(
                            name=f"sched:{task.name}",
                            func=task.func,
                            args=task.args,
                            kwargs=task.kwargs,
                            priority=7,  # Scheduled tasks get high priority
                        )
                    else:
                        task.func(*task.args, **task.kwargs)

                    task.last_run = time.monotonic()
                    task.run_count += 1

                except Exception as e:
                    logger.warning(f"Scheduled task '{task.name}' failed: {e}")

    def schedule(self, name: str, func: Callable, interval_seconds: int,
                 args: tuple = (), kwargs: dict = None):
        """Schedule a recurring task."""
        task = ScheduledTask(
            name=name,
            func=func,
            interval_seconds=interval_seconds,
            args=args,
            kwargs=kwargs or {},
        )
        self._scheduled_tasks[name] = task
        logger.info(
            f"⏰ Scheduled task '{name}' "
            f"(every {interval_seconds}s)"
        )
        return task

    def unschedule(self, name: str):
        """Remove a scheduled task."""
        task = self._scheduled_tasks.pop(name, None)
        if task:
            logger.info(f"⏰ Unscheduled task '{name}'")
        return task is not None

    # ── Heartbeat Callback ──

    def _on_heartbeat(self, tick: int):
        """Called on each heartbeat tick."""
        # Memory consolidation every 10 ticks
        if tick % 10 == 0 and self.memory:
            self.memory.consolidate()

        # Process cleanup every 5 ticks
        if tick % 5 == 0 and self.processes:
            self.processes.cleanup()

        # Emit heartbeat signal
        self.signals.emit("kernel.heartbeat", {
            "tick": tick,
            "uptime_sec": self.uptime,
            "events_processed": self.events.total_events,
        })

    # ── Properties ──

    @property
    def uptime(self) -> float:
        """Kernel uptime in seconds."""
        if self._boot_time == 0:
            return 0.0
        return round(time.monotonic() - self._boot_time, 1)

    @property
    def uptime_human(self) -> str:
        """Human-readable uptime."""
        s = int(self.uptime)
        if s < 60:
            return f"{s}s"
        elif s < 3600:
            return f"{s // 60}m {s % 60}s"
        else:
            h = s // 3600
            m = (s % 3600) // 60
            return f"{h}h {m}m"

    @property
    def stats(self) -> dict:
        return {
            "version": self.VERSION,
            "running": self._running,
            "uptime_sec": self.uptime,
            "uptime_human": self.uptime_human,
            "boot_time": datetime.utcfromtimestamp(
                time.time() - self.uptime
            ).isoformat() + "Z" if self._boot_time > 0 else None,
            "events": {
                "total": self.events.total_events,
                "queued": self.events.size,
            },
            "heartbeat": self.heartbeat.stats,
            "signals": self.signals.stats,
            "scheduled_tasks": {
                name: task.to_dict()
                for name, task in self._scheduled_tasks.items()
            },
            "memory": self.memory.stats if self.memory else None,
            "processes": self.processes.stats if self.processes else None,
        }

    def banner(self) -> str:
        """Generate the boot banner."""
        mem_stats = self.memory.stats if self.memory else {}
        proc_stats = self.processes.top() if self.processes else {}

        working = mem_stats.get("working", {}).get("items", 0)
        lt_kb = mem_stats.get("long_term", {}).get("vault_size_kb", 0)
        ep_kb = mem_stats.get("episodic", {}).get("storage_kb", 0)
        p_running = proc_stats.get("processes", {}).get("running", 0)
        p_queued = proc_stats.get("processes", {}).get("queued", 0)

        return (
            "\n"
            "  ╔══════════════════════════════════════════════════════════╗\n"
            "  ║  🧠 Marziel OS v1.0 — Private Intelligence OS          ║\n"
            "  ╚══════════════════════════════════════════════════════════╝\n"
            "\n"
            f"     Kernel:     active (heartbeat: {self.heartbeat.interval}s)\n"
            f"     Memory:     working={working} | long-term={lt_kb}KB | episodic={ep_kb}KB\n"
            f"     Processes:  {p_running} running, {p_queued} queued\n"
            f"     Uptime:     {self.uptime_human}\n"
        )
