#!/usr/bin/env python3
"""
╔══════════════════════════════════════════════════════════════╗
║  MITOSIS — Marziel AI Self-Healing Process Manager          ║
║  Auto-restarts crashed services, monitors resources,        ║
║  validates model integrity, and keeps everything alive.     ║
╚══════════════════════════════════════════════════════════════╝

Usage:
    python3 mitosis.py start    — Start all services with monitoring
    python3 mitosis.py stop     — Graceful shutdown
    python3 mitosis.py status   — Show service health
    python3 mitosis.py doctor   — Run diagnostics
"""

import os
import sys
import json
import time
import signal
import hashlib
import logging
import subprocess
import threading
import platform
import shutil
from pathlib import Path
from datetime import datetime, timedelta
from logging.handlers import RotatingFileHandler
from urllib.request import urlopen, Request
from urllib.error import URLError

# ══════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════

MARZIEL_HOME = Path(os.environ.get("MARZIEL_HOME", Path.home() / "marziel-ai"))
CONFIG_FILE = MARZIEL_HOME / "config" / "marziel.json"
PID_FILE = MARZIEL_HOME / ".mitosis.pid"
STATE_FILE = MARZIEL_HOME / "data" / "mitosis_state.json"
LOG_DIR = MARZIEL_HOME / "logs"

# Defaults
DEFAULT_CONFIG = {
    "llm_port": 8000,
    "api_port": 8001,
    "dashboard_port": 3001,
    "model_path": str(MARZIEL_HOME / "models" / "marziel-8b"),
    "inference_engine": "auto",
    "health_check_interval": 10,
    "max_restart_attempts": 5,
    "restart_backoff_base": 2,
    "max_backoff_seconds": 120,
    "memory_limit_percent": 85,
    "log_max_bytes": 10_485_760,  # 10 MB
    "log_backup_count": 5,
    "auto_update_check": True,
    "version": "1.0.0",
}

# ══════════════════════════════════════════════════════════════
# LOGGING
# ══════════════════════════════════════════════════════════════

LOG_DIR.mkdir(parents=True, exist_ok=True)

logger = logging.getLogger("mitosis")
logger.setLevel(logging.INFO)

# Console handler
ch = logging.StreamHandler()
ch.setFormatter(logging.Formatter(
    '\033[0;36m%(asctime)s\033[0m [\033[1m%(levelname)s\033[0m] %(message)s',
    datefmt='%H:%M:%S'
))
logger.addHandler(ch)

# File handler with rotation
fh = RotatingFileHandler(
    LOG_DIR / "mitosis.log",
    maxBytes=DEFAULT_CONFIG["log_max_bytes"],
    backupCount=DEFAULT_CONFIG["log_backup_count"],
)
fh.setFormatter(logging.Formatter(
    '%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
))
logger.addHandler(fh)


# ══════════════════════════════════════════════════════════════
# SERVICE DEFINITION
# ══════════════════════════════════════════════════════════════

class ManagedService:
    """Represents a service managed by Mitosis."""

    def __init__(self, name: str, cmd: list, port: int, health_url: str,
                 startup_timeout: int = 120, critical: bool = True):
        self.name = name
        self.cmd = cmd
        self.port = port
        self.health_url = health_url
        self.startup_timeout = startup_timeout
        self.critical = critical
        self.process = None
        self.restart_count = 0
        self.last_restart = None
        self.last_healthy = None
        self.status = "stopped"
        self.log_file = None

    def start(self, config: dict):
        """Start the service process."""
        log_path = LOG_DIR / f"{self.name}.log"
        self.log_file = open(log_path, "a")

        env = os.environ.copy()
        env["PYTHONUNBUFFERED"] = "1"

        try:
            self.process = subprocess.Popen(
                self.cmd,
                stdout=self.log_file,
                stderr=subprocess.STDOUT,
                env=env,
                cwd=str(MARZIEL_HOME),
            )
            self.status = "starting"
            logger.info(f"🚀 {self.name} started (PID {self.process.pid})")
            return True
        except Exception as e:
            logger.error(f"❌ Failed to start {self.name}: {e}")
            self.status = "error"
            return False

    def stop(self):
        """Gracefully stop the service."""
        if self.process and self.process.poll() is None:
            logger.info(f"⏹  Stopping {self.name} (PID {self.process.pid})...")
            self.process.terminate()
            try:
                self.process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                logger.warning(f"⚠️  Force killing {self.name}")
                self.process.kill()
            self.status = "stopped"
        if self.log_file:
            self.log_file.close()

    def is_running(self) -> bool:
        """Check if the process is still running."""
        return self.process is not None and self.process.poll() is None

    def health_check(self) -> bool:
        """Ping the health endpoint."""
        try:
            req = Request(self.health_url, method='GET')
            resp = urlopen(req, timeout=5)
            if resp.status == 200:
                self.last_healthy = datetime.now()
                self.status = "healthy"
                return True
        except Exception:
            pass
        return False

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "status": self.status,
            "pid": self.process.pid if self.process and self.is_running() else None,
            "port": self.port,
            "restart_count": self.restart_count,
            "last_healthy": self.last_healthy.isoformat() if self.last_healthy else None,
            "last_restart": self.last_restart.isoformat() if self.last_restart else None,
        }


# ══════════════════════════════════════════════════════════════
# MITOSIS — THE CORE WATCHDOG
# ══════════════════════════════════════════════════════════════

class Mitosis:
    """Self-healing process manager for Marziel AI."""

    def __init__(self):
        self.config = self._load_config()
        self.services: list[ManagedService] = []
        self.running = False
        self._shutdown_event = threading.Event()

    def _load_config(self) -> dict:
        """Load config, merging with defaults."""
        config = DEFAULT_CONFIG.copy()
        if CONFIG_FILE.exists():
            try:
                with open(CONFIG_FILE) as f:
                    user_config = json.load(f)
                config.update(user_config)
            except Exception as e:
                logger.warning(f"⚠️  Config read error, using defaults: {e}")
        return config

    def _save_state(self):
        """Persist current state for recovery."""
        state = {
            "timestamp": datetime.now().isoformat(),
            "services": [s.to_dict() for s in self.services],
            "system": self._get_system_info(),
        }
        STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
        with open(STATE_FILE, "w") as f:
            json.dump(state, f, indent=2)

    def _detect_inference_engine(self) -> str:
        """Auto-detect the best inference engine for this machine."""
        engine = self.config.get("inference_engine", "auto")
        if engine != "auto":
            return engine

        # Apple Silicon → MLX
        if platform.system() == "Darwin" and platform.machine() == "arm64":
            try:
                subprocess.run([sys.executable, "-c", "import mlx_lm"],
                             capture_output=True, check=True)
                return "mlx"
            except subprocess.CalledProcessError:
                pass

        # NVIDIA GPU → vLLM
        if shutil.which("nvidia-smi"):
            try:
                subprocess.run([sys.executable, "-c", "import vllm"],
                             capture_output=True, check=True)
                return "vllm"
            except subprocess.CalledProcessError:
                pass

        # CPU fallback → llama.cpp
        try:
            subprocess.run([sys.executable, "-c", "import llama_cpp"],
                         capture_output=True, check=True)
            return "llamacpp"
        except subprocess.CalledProcessError:
            pass

        return "mlx"  # Default fallback

    def _build_llm_command(self, engine: str) -> list:
        """Build the LLM server command based on engine."""
        model_path = self.config["model_path"]
        port = self.config["llm_port"]

        if engine == "mlx":
            return [sys.executable, "-m", "mlx_lm.server",
                    "--model", model_path,
                    "--port", str(port)]
        elif engine == "vllm":
            return [sys.executable, "-m", "vllm.entrypoints.openai.api_server",
                    "--model", model_path,
                    "--port", str(port),
                    "--max-model-len", "4096",
                    "--dtype", "float16"]
        else:  # llamacpp
            return [sys.executable, "-m", "llama_cpp.server",
                    "--model", model_path,
                    "--port", str(port)]

    def _verify_model(self) -> bool:
        """Verify model files exist and aren't corrupted."""
        model_path = Path(self.config["model_path"])

        if not model_path.exists():
            logger.error(f"❌ Model not found: {model_path}")
            return False

        config_file = model_path / "config.json"
        if not config_file.exists():
            logger.error(f"❌ Model config.json missing — model may be corrupted")
            return False

        # Check file sizes are reasonable
        total_size = sum(f.stat().st_size for f in model_path.rglob("*") if f.is_file())
        if total_size < 100_000_000:  # Less than 100 MB = definitely broken
            logger.error(f"❌ Model too small ({total_size / 1e6:.0f} MB) — likely corrupted")
            return False

        logger.info(f"✅ Model verified ({total_size / 1e9:.1f} GB)")
        return True

    def _get_system_info(self) -> dict:
        """Get current system resource usage."""
        info = {
            "platform": platform.platform(),
            "python": platform.python_version(),
            "cpu_count": os.cpu_count(),
        }

        # Memory
        try:
            if platform.system() == "Darwin":
                import resource
                mem_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
                info["ram_total_gb"] = round(mem_bytes / 1e9, 1)
                # Get memory pressure
                result = subprocess.run(["vm_stat"], capture_output=True, text=True)
                lines = result.stdout.split('\n')
                for line in lines:
                    if "Pages free" in line:
                        free_pages = int(line.split(':')[1].strip().rstrip('.'))
                        info["ram_free_gb"] = round(free_pages * 4096 / 1e9, 1)
                        break
            else:
                with open("/proc/meminfo") as f:
                    meminfo = f.read()
                for line in meminfo.split('\n'):
                    if line.startswith("MemTotal"):
                        info["ram_total_gb"] = round(int(line.split()[1]) / 1e6, 1)
                    elif line.startswith("MemAvailable"):
                        info["ram_free_gb"] = round(int(line.split()[1]) / 1e6, 1)
        except Exception:
            pass

        # Disk
        try:
            usage = shutil.disk_usage(MARZIEL_HOME)
            info["disk_free_gb"] = round(usage.free / 1e9, 1)
            info["disk_total_gb"] = round(usage.total / 1e9, 1)
        except Exception:
            pass

        return info

    def _check_resources(self) -> list:
        """Check if system resources are healthy. Returns list of warnings."""
        warnings = []
        info = self._get_system_info()

        # RAM check
        ram_total = info.get("ram_total_gb", 0)
        ram_free = info.get("ram_free_gb", 0)
        if ram_total > 0 and ram_free > 0:
            usage_pct = ((ram_total - ram_free) / ram_total) * 100
            if usage_pct > self.config["memory_limit_percent"]:
                warnings.append(f"RAM usage high: {usage_pct:.0f}% ({ram_free:.1f}GB free)")

        # Disk check
        disk_free = info.get("disk_free_gb", 0)
        if 0 < disk_free < 2:
            warnings.append(f"Low disk space: {disk_free:.1f}GB free")

        return warnings

    def _restart_service(self, service: ManagedService):
        """Restart a service with exponential backoff."""
        max_attempts = self.config["max_restart_attempts"]
        backoff_base = self.config["restart_backoff_base"]
        max_backoff = self.config["max_backoff_seconds"]

        if service.restart_count >= max_attempts:
            # Reset counter after cooldown period
            if (service.last_restart and
                    datetime.now() - service.last_restart > timedelta(minutes=10)):
                logger.info(f"🔄 Reset restart counter for {service.name} after cooldown")
                service.restart_count = 0
            else:
                logger.error(
                    f"🛑 {service.name} exceeded max restarts ({max_attempts}). "
                    f"Waiting for cooldown. Run 'mitosis doctor' to diagnose."
                )
                service.status = "failed"
                return

        # Exponential backoff
        backoff = min(backoff_base ** service.restart_count, max_backoff)
        service.restart_count += 1
        service.last_restart = datetime.now()

        logger.warning(
            f"🔄 Restarting {service.name} "
            f"(attempt {service.restart_count}/{max_attempts}, "
            f"wait {backoff}s)..."
        )

        service.stop()
        time.sleep(backoff)
        service.start(self.config)

    def _health_loop(self):
        """Main health monitoring loop — the heartbeat of Mitosis."""
        interval = self.config["health_check_interval"]
        resource_check_counter = 0

        while not self._shutdown_event.is_set():
            for service in self.services:
                # Check process alive
                if not service.is_running():
                    if service.status not in ("stopped", "failed"):
                        exit_code = service.process.returncode if service.process else "?"
                        logger.error(
                            f"💀 {service.name} CRASHED (exit code: {exit_code})"
                        )
                        self._restart_service(service)
                    continue

                # Health endpoint check
                if not service.health_check():
                    # Allow some grace period after start
                    if service.status == "starting":
                        elapsed = (datetime.now() - (service.last_restart or datetime.now()))
                        if elapsed.total_seconds() < service.startup_timeout:
                            continue
                        logger.warning(f"⚠️  {service.name} startup timeout exceeded")

                    # If it was healthy before, it might be a hiccup
                    if service.last_healthy:
                        unhealthy_time = datetime.now() - service.last_healthy
                        if unhealthy_time.total_seconds() > 30:
                            logger.error(
                                f"🏥 {service.name} unhealthy for "
                                f"{unhealthy_time.total_seconds():.0f}s — restarting"
                            )
                            self._restart_service(service)
                    else:
                        # Never was healthy, might still be starting
                        pass

            # Resource checks (every 60s, not every loop)
            resource_check_counter += 1
            if resource_check_counter >= (60 // interval):
                resource_check_counter = 0
                warnings = self._check_resources()
                for w in warnings:
                    logger.warning(f"⚠️  RESOURCE: {w}")

            # Save state
            self._save_state()

            # Wait for next check
            self._shutdown_event.wait(interval)

    def _signal_handler(self, signum, frame):
        """Handle shutdown signals gracefully."""
        sig_name = signal.Signals(signum).name
        logger.info(f"\n📤 Received {sig_name} — shutting down gracefully...")
        self.stop()

    # ── PUBLIC API ──

    def start(self):
        """Start all services and begin monitoring."""
        logger.info("═" * 50)
        logger.info("🧠 MITOSIS — Marziel AI Process Manager")
        logger.info("═" * 50)

        # Write PID file
        PID_FILE.write_text(str(os.getpid()))

        # Verify model
        if not self._verify_model():
            logger.error("Cannot start without a valid model. Run 'mitosis doctor'.")
            sys.exit(1)

        # Detect inference engine
        engine = self._detect_inference_engine()
        logger.info(f"🔧 Inference engine: {engine}")

        # Define services
        llm_cmd = self._build_llm_command(engine)
        api_port = self.config["api_port"]
        llm_port = self.config["llm_port"]

        self.services = [
            ManagedService(
                name="llm-server",
                cmd=llm_cmd,
                port=llm_port,
                health_url=f"http://localhost:{llm_port}/health",
                startup_timeout=120,
                critical=True,
            ),
            ManagedService(
                name="marziel-backend",
                cmd=[
                    sys.executable, str(MARZIEL_HOME / "marziel_server.py"),
                ],
                port=api_port,
                health_url=f"http://localhost:{api_port}/health",
                startup_timeout=15,
                critical=True,
            ),
        ]

        # Set env for backend
        os.environ["VLLM_URL"] = f"http://localhost:{llm_port}/v1/chat/completions"
        os.environ["VLLM_MODEL"] = self.config["model_path"]
        os.environ["MEMORY_DIR"] = str(MARZIEL_HOME / "data")

        # Register signal handlers
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)

        # Start services
        self.running = True
        for service in self.services:
            service.start(self.config)
            time.sleep(1)

        # Wait for LLM to be ready (it's slow)
        logger.info("⏳ Waiting for LLM model to load...")
        llm = self.services[0]
        for i in range(llm.startup_timeout):
            if llm.health_check():
                logger.info(f"✅ LLM server ready in {i}s")
                break
            if not llm.is_running():
                logger.error("❌ LLM server crashed during startup")
                self._restart_service(llm)
                break
            time.sleep(1)

        # Print status
        self._print_status()

        # Start health monitoring loop
        logger.info("💓 Health monitoring active")
        self._health_loop()

    def stop(self):
        """Graceful shutdown of all services."""
        self._shutdown_event.set()
        self.running = False

        logger.info("📤 Saving state...")
        self._save_state()

        for service in reversed(self.services):
            service.stop()

        if PID_FILE.exists():
            PID_FILE.unlink()

        logger.info("✅ All services stopped. Goodbye.")
        sys.exit(0)

    def status(self):
        """Show current service status."""
        if STATE_FILE.exists():
            with open(STATE_FILE) as f:
                state = json.load(f)

            print("\n  🧠 MARZIEL AI — System Status")
            print("  " + "─" * 46)

            for svc in state.get("services", []):
                status_icon = {
                    "healthy": "🟢",
                    "starting": "🟡",
                    "stopped": "⚪",
                    "failed": "🔴",
                    "error": "🔴",
                }.get(svc["status"], "⚪")

                print(f"  {status_icon} {svc['name']:20} port:{svc['port']}  "
                      f"pid:{svc.get('pid', '-'):>6}  "
                      f"restarts:{svc['restart_count']}")

            sys_info = state.get("system", {})
            print("  " + "─" * 46)
            print(f"  RAM: {sys_info.get('ram_free_gb', '?')}GB free / "
                  f"{sys_info.get('ram_total_gb', '?')}GB total")
            print(f"  Disk: {sys_info.get('disk_free_gb', '?')}GB free")
            print(f"  Last check: {state.get('timestamp', '?')}")
            print()
        else:
            print("\n  ⚪ Marziel AI is not running.")
            print("  Start with: python3 mitosis.py start\n")

    def doctor(self):
        """Run diagnostics and suggest fixes."""
        print("\n  🩺 MARZIEL AI — Diagnostics")
        print("  " + "─" * 46)

        issues = []

        # Check Python
        py_ver = platform.python_version()
        major, minor = int(py_ver.split('.')[0]), int(py_ver.split('.')[1])
        if major < 3 or minor < 10:
            issues.append(f"  ❌ Python {py_ver} — need 3.10+")
        else:
            print(f"  ✅ Python {py_ver}")

        # Check model
        model_path = Path(self.config["model_path"])
        if model_path.exists() and (model_path / "config.json").exists():
            total_size = sum(f.stat().st_size for f in model_path.rglob("*") if f.is_file())
            print(f"  ✅ Model ({total_size / 1e9:.1f} GB)")
        else:
            issues.append(f"  ❌ Model not found at {model_path}")
            issues.append(f"     Fix: re-run install.sh to re-download")

        # Check inference engine
        engine = self._detect_inference_engine()
        print(f"  ✅ Inference engine: {engine}")

        # Check ports
        for port in [self.config["llm_port"], self.config["api_port"]]:
            import socket
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            result = sock.connect_ex(('localhost', port))
            if result == 0:
                print(f"  ✅ Port {port} — service running")
            else:
                print(f"  ⚪ Port {port} — not in use")
            sock.close()

        # Resources
        sys_info = self._get_system_info()
        ram = sys_info.get("ram_total_gb", 0)
        if ram < 8:
            issues.append(f"  ⚠️  Only {ram}GB RAM — 8GB+ recommended")
        else:
            print(f"  ✅ RAM: {ram}GB")

        disk = sys_info.get("disk_free_gb", 0)
        if disk < 2:
            issues.append(f"  ⚠️  Low disk: {disk}GB free")
        else:
            print(f"  ✅ Disk: {disk}GB free")

        # Check logs for recent errors
        log_file = LOG_DIR / "mitosis.log"
        if log_file.exists():
            with open(log_file) as f:
                last_lines = f.readlines()[-20:]
            errors = [l.strip() for l in last_lines if "[ERROR]" in l]
            if errors:
                print(f"\n  📋 Recent errors ({len(errors)}):")
                for e in errors[-3:]:
                    print(f"     {e}")

        if issues:
            print(f"\n  ⚠️  {len(issues)} issue(s) found:")
            for i in issues:
                print(i)
        else:
            print("\n  ✅ All systems healthy!")

        print()

    def _print_status(self):
        """Print the startup status banner."""
        print()
        print("  \033[0;36m╔══════════════════════════════════════════╗\033[0m")
        print("  \033[0;36m║     🧠  Marziel AI is running!          ║\033[0m")
        print("  \033[0;36m╠══════════════════════════════════════════╣\033[0m")
        print(f"  \033[0;36m║\033[0m  Dashboard:  http://localhost:{self.config['dashboard_port']}     \033[0;36m║\033[0m")
        print(f"  \033[0;36m║\033[0m  API:        http://localhost:{self.config['api_port']}         \033[0;36m║\033[0m")
        print(f"  \033[0;36m║\033[0m  LLM:        http://localhost:{self.config['llm_port']}         \033[0;36m║\033[0m")
        print("  \033[0;36m╠══════════════════════════════════════════╣\033[0m")
        print("  \033[0;36m║\033[0m  💓 Health checks every "
              f"{self.config['health_check_interval']}s          \033[0;36m║\033[0m")
        print("  \033[0;36m║\033[0m  🔄 Auto-restart on crash              \033[0;36m║\033[0m")
        print("  \033[0;36m║\033[0m  📊 Resource monitoring active         \033[0;36m║\033[0m")
        print("  \033[0;36m║\033[0m  📝 Logs: ~/marziel-ai/logs/           \033[0;36m║\033[0m")
        print("  \033[0;36m╠══════════════════════════════════════════╣\033[0m")
        print("  \033[0;36m║\033[0m  Press Ctrl+C to stop                  \033[0;36m║\033[0m")
        print("  \033[0;36m╚══════════════════════════════════════════╝\033[0m")
        print()


# ══════════════════════════════════════════════════════════════
# CLI
# ══════════════════════════════════════════════════════════════

def main():
    mitosis = Mitosis()

    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(0)

    cmd = sys.argv[1].lower()

    if cmd == "start":
        mitosis.start()
    elif cmd == "stop":
        if PID_FILE.exists():
            pid = int(PID_FILE.read_text().strip())
            try:
                os.kill(pid, signal.SIGTERM)
                print("  ✅ Sent shutdown signal to Mitosis")
            except ProcessLookupError:
                print("  ⚪ Mitosis not running (stale PID file)")
                PID_FILE.unlink()
        else:
            print("  ⚪ Mitosis not running")
    elif cmd == "status":
        mitosis.status()
    elif cmd == "doctor":
        mitosis.doctor()
    else:
        print(f"  Unknown command: {cmd}")
        print(__doc__)
        sys.exit(1)


if __name__ == "__main__":
    main()
