#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Worker Supervisor Daemon — L1.15 Phase 1

Monitors worker_pids.json every 10s, detects dead/hung workers via PID check +
heartbeat freshness, auto-respawns with exponential backoff, releases orphaned
tasks back to "open". Usage: python -m pulse.daemons.runtime.worker_supervisor [--verbose|--dry-run]
"""
from __future__ import annotations

import io
import json
import logging
import os
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

if os.name == "nt" and __name__ == "__main__":
    if hasattr(sys.stdout, 'reconfigure'):
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
    if hasattr(sys.stderr, 'reconfigure'):
        sys.stderr.reconfigure(encoding='utf-8', errors='replace')

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent.parent

if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from pulse.core.brain_utils import (
    STATE_DIR, TASK_BOARD_FILE, now_iso as _now_iso,
    is_process_running, atomic_update_json,
)
from pulse.tasks.task_ops import load_worker_pids, save_worker_pids, kill_process

log = logging.getLogger("pulse.worker_supervisor")

POLL_INTERVAL = 10
HEARTBEAT_STALE = 90        # 3x the 30s heartbeat interval
MAX_RESPAWNS = 5
BACKOFF_CAP = 60
WORKER_SCRIPT = ROOT_DIR / "pulse" / "workers" / "pulse_worker.py"
HEARTBEAT_DIR = STATE_DIR / "heartbeats"
ALERTS_FILE = STATE_DIR / "supervisor_alerts.json"
GRACE_PERIOD = 120           # seconds after start before heartbeat is required


def _is_worker_dead(record: dict) -> tuple[bool, str]:
    """Check if a worker is dead or hung. Returns (is_dead, reason)."""
    pid = record.get("pid", 0)
    wid = record.get("id", "?")

    if not pid or not is_process_running(pid):
        return True, "process_dead"

    # PID is alive — check heartbeat freshness
    hb_file = HEARTBEAT_DIR / f"{wid}.json"
    if not hb_file.exists():
        # Grace period for freshly started workers
        started = record.get("started_at", "")
        if started:
            try:
                from datetime import datetime as _dt
                start_time = _dt.fromisoformat(started.replace("Z", "+00:00"))
                age = (datetime.now(timezone.utc) - start_time).total_seconds()
                if age < GRACE_PERIOD:
                    return False, "grace_period"
            except Exception:
                pass
        return True, "no_heartbeat"

    try:
        age = time.time() - hb_file.stat().st_mtime
        if age > HEARTBEAT_STALE:
            return True, "heartbeat_stale"
    except Exception:
        return True, "heartbeat_unreadable"

    return False, "alive"


def _release_orphaned_tasks(worker_id: str, verbose: bool = False) -> int:
    """Release tasks claimed by a dead worker back to 'open'."""
    released = 0

    def _update(board):
        nonlocal released
        for dispatch in board.get("dispatches", []):
            for task in dispatch.get("tasks", []):
                if (task.get("status") == "active" and
                        task.get("claimed_by") == worker_id):
                    task["status"] = "open"
                    task["claimed_by"] = None
                    task["claimed_at"] = None
                    released += 1
        return board

    try:
        atomic_update_json(TASK_BOARD_FILE, _update, default={"dispatches": []})
    except Exception as e:
        log.warning("Failed to release tasks for %s: %s", worker_id, e)

    if released and verbose:
        log.info("Released %d orphaned task(s) from %s", released, worker_id)
    return released


def _spawn_worker(record: dict, verbose: bool = False) -> int | None:
    """Respawn a worker using its original parameters. Returns new PID or None."""
    mode = record.get("mode", "headless")
    model = record.get("model", "unknown")
    provider = record.get("provider", "google")
    tier = record.get("tier", "fast")
    worker_id = record.get("id", "unknown")

    log_dir = STATE_DIR / "worker_logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / f"{worker_id}.log"

    if mode == "cli_autopilot":
        cmd = [
            sys.executable, "-m", "pulse.cli.cmd_swarm",
            "--autopilot", "--agent", provider, "--model", model,
            "--tier", tier, "--id", worker_id,
        ]
    else:
        cmd = [
            sys.executable, str(WORKER_SCRIPT),
            "--provider", provider, "--model", model,
            "--tier", tier, "--id", worker_id,
        ]

    try:
        kwargs = {
            "stdout": open(log_file, "a", encoding="utf-8"),
            "stderr": subprocess.STDOUT,
        }
        if os.name == "nt":
            kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
        else:
            kwargs["start_new_session"] = True
        proc = subprocess.Popen(cmd, **kwargs)
        if verbose:
            log.info("Respawned %s (PID %d, mode=%s)", worker_id, proc.pid, mode)
        return proc.pid
    except Exception as e:
        log.error("Failed to respawn %s: %s", worker_id, e)
        return None


def _write_alert(worker_id: str, reason: str, respawn_count: int):
    """Write supervisor alert for a permanently dead worker."""
    try:
        alerts = []
        if ALERTS_FILE.exists():
            alerts = json.loads(ALERTS_FILE.read_text(encoding="utf-8"))
            if not isinstance(alerts, list):
                alerts = []
        alerts.append({
            "worker_id": worker_id, "reason": reason,
            "respawn_count": respawn_count, "max_respawns": MAX_RESPAWNS,
            "timestamp": _now_iso(), "status": "abandoned",
        })
        ALERTS_FILE.write_text(json.dumps(alerts[-50:], indent=2), encoding="utf-8")
    except Exception as e:
        log.warning("Failed to write supervisor alert: %s", e)


def _deregister_dead(worker_id: str):
    """Remove dead worker from agent registry + clean heartbeat."""
    try:
        from pulse.workers.runner.worker_state import deregister_worker
        deregister_worker(worker_id)
    except Exception:
        pass
    try:
        hb = HEARTBEAT_DIR / f"{worker_id}.json"
        if hb.exists():
            hb.unlink()
    except Exception:
        pass


def _kill_hung(pid: int, worker_id: str):
    """Force-kill a hung process (PID alive but heartbeat stale)."""
    try:
        kill_process(pid)
        log.info("Killed hung worker %s (PID %d)", worker_id, pid)
    except Exception as e:
        log.warning("Failed to kill hung %s: %s", worker_id, e)


def run_supervisor(verbose: bool = False, dry_run: bool = False):
    """Main supervisor loop — polls workers every POLL_INTERVAL seconds."""
    if verbose:
        log.info("Supervisor online (poll=%ds, stale=%ds, max_respawns=%d)",
                 POLL_INTERVAL, HEARTBEAT_STALE, MAX_RESPAWNS)

    respawn_counts: dict[str, int] = {}

    while True:
        try:
            records = load_worker_pids()
            if not records:
                time.sleep(POLL_INTERVAL)
                continue

            alive_records: list[dict] = []
            changed = False

            for rec in records:
                wid = rec.get("id", "?")
                pid = rec.get("pid", 0)
                is_dead, reason = _is_worker_dead(rec)

                if not is_dead:
                    alive_records.append(rec)
                    continue

                changed = True
                if verbose:
                    log.warning("Worker %s: %s (PID %d)", wid, reason, pid)

                # Kill hung processes before respawning
                if reason == "heartbeat_stale" and pid and is_process_running(pid):
                    if not dry_run:
                        _kill_hung(pid, wid)

                if not dry_run:
                    _release_orphaned_tasks(wid, verbose=verbose)
                    _deregister_dead(wid)

                count = respawn_counts.get(wid, 0)
                if count >= MAX_RESPAWNS:
                    if verbose:
                        log.error("Worker %s exceeded max respawns (%d). Abandoned.", wid, MAX_RESPAWNS)
                    _write_alert(wid, reason, count)
                    continue

                backoff = min(2 ** (count + 1), BACKOFF_CAP)
                if verbose:
                    log.info("Respawning %s in %ds (%d/%d)", wid, backoff, count + 1, MAX_RESPAWNS)

                if not dry_run:
                    time.sleep(backoff)
                    new_pid = _spawn_worker(rec, verbose=verbose)
                    if new_pid:
                        rec["pid"] = new_pid
                        rec["started_at"] = datetime.now(timezone.utc).isoformat()
                        alive_records.append(rec)
                    respawn_counts[wid] = count + 1
                else:
                    log.info("[DRY RUN] Would respawn %s", wid)

            if changed and not dry_run:
                save_worker_pids(alive_records)
                if verbose:
                    log.info("PIDs updated: %d alive", len(alive_records))

        except Exception as e:
            log.error("Supervisor scan error: %s", e, exc_info=True)

        time.sleep(POLL_INTERVAL)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [SUPERVISOR] %(message)s")

    dry_run = "--dry-run" in sys.argv
    verbose = "--verbose" in sys.argv or "-v" in sys.argv or dry_run

    if dry_run:
        print("[SUPERVISOR] Dry run mode — no respawns or writes")

    try:
        run_supervisor(verbose=verbose, dry_run=dry_run)
    except KeyboardInterrupt:
        print("\n[SUPERVISOR] Stopped.")
