#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Swarm Workers: Launch, stop, and monitor headless/autopilot workers.

Functions:
  - _launch_headless(): Start API/Ollama background workers
  - _launch_cli_autopilot(): Start CLI subscription autopilot workers
  - _build_agent_cmd(): Build CLI command for agent launch
  - _swarm_stop/status/watch(): Worker management commands
"""
from __future__ import annotations

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

from pulse.cli import ROOT_DIR, STATE_DIR, WORKERS_DIR, AGENTS_DIR
from pulse.tasks.task_ops import (
    is_process_running, kill_process,
    load_worker_pids, save_worker_pids, TASK_BOARD,
)
from .swarm_providers import _PROVIDER_META, _resolve_access

WORKER_SCRIPT = WORKERS_DIR / "pulse_worker.py"


def _is_headless_provider(provider_key: str) -> bool:
    """Headless = API key or local. CLI subscription = autopilot."""
    return _resolve_access(provider_key) in ("api", "local")


def _provider_to_worker_arg(provider_key: str) -> str:
    """Map provider key to pulse_worker.py --provider value."""
    return _PROVIDER_META.get(provider_key, {}).get("worker_provider", "google")


def _build_agent_cmd(agent: str, model: str, boot_message: str) -> list[str] | None:
    """Build CLI command to launch an agent with a boot message."""
    # Ollama cloud: route through `ollama launch claude --model X -- <claude flags>`
    if agent == "ollama-cloud":
        ollama = shutil.which("ollama")
        if not ollama:
            print("[PULSE] Ollama CLI not found. Install from https://ollama.com", flush=True)
            return None
        cmd = [ollama, "launch", "claude", "--model", model,
               "--", "-p", boot_message, "--dangerously-skip-permissions"]
        if os.name == "nt" and (ollama.lower().endswith(".cmd") or ollama.lower().endswith(".bat")):
            cmd = ["cmd.exe", "/c"] + cmd
        return cmd

    resolved = shutil.which(agent)
    if not resolved:
        print(f"[PULSE] Agent CLI not found: {agent}", flush=True)
        return None
    cmd = [resolved]
    exe = resolved.lower()
    if "gemini" in exe:
        if model:
            cmd.extend(["-m", model])
        cmd.extend(["-p", boot_message, "--yolo"])
    elif "codex" in exe:
        cmd.append("exec")
        if model:
            cmd.extend(["-m", model])
        cmd.extend(["--full-auto", boot_message])
    elif "claude" in exe:
        if model:
            cmd.extend(["--model", model])
        cmd.extend(["-p", boot_message, "--dangerously-skip-permissions"])
    if os.name == "nt" and (exe.endswith(".cmd") or exe.endswith(".bat")):
        cmd = ["cmd.exe", "/c"] + cmd
    return cmd


def _spawn_background(cmd: list, log_file: Path) -> subprocess.Popen:
    """Spawn a background process with platform-appropriate flags."""
    log_file.parent.mkdir(parents=True, exist_ok=True)
    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
    return subprocess.Popen(cmd, **kwargs)


def _launch_headless(provider_key: str, model: str, tier: str, count: int,
                     skills_used: bool = False):
    """Start headless background workers via pulse_worker.py."""
    records = load_worker_pids()
    provider_arg = _provider_to_worker_arg(provider_key)
    launched = 0
    for _ in range(count):
        model_slug = model.replace("/", "-").replace(":", "-")
        worker_id = f"{model_slug}_{len(records) + 1:02d}"
        cmd = [
            sys.executable, str(WORKER_SCRIPT),
            "--provider", provider_arg, "--model", model,
            "--tier", tier, "--id", worker_id,
        ]
        if skills_used:
            cmd.append("--skills-used")
        proc = _spawn_background(cmd, STATE_DIR / "worker_logs" / f"{worker_id}.log")
        records.append({
            "pid": proc.pid, "id": worker_id,
            "provider": provider_arg, "model": model, "tier": tier,
            "mode": "headless", "started_at": datetime.now(timezone.utc).isoformat(),
        })
        launched += 1
        print(f"  Started {worker_id} (PID {proc.pid})")
    save_worker_pids(records)
    print(f"\n[PULSE] {launched} worker(s) launched. Use 'pulse swarm status' to check.")


def _launch_cli_autopilot(provider_key: str, model: str, tier: str, count: int,
                          skills_used: bool = False):
    """Start CLI subscription autopilot as background processes."""
    meta = _PROVIDER_META.get(provider_key, {})
    agent = meta.get("cli", provider_key)
    records = load_worker_pids()
    launched = 0
    for _ in range(count):
        model_slug = model.replace("/", "-").replace(":", "-")
        worker_id = f"{model_slug}_{len(records) + 1:02d}"
        cmd = [
            sys.executable, "-m", "pulse.cli.cmd_swarm",
            "--autopilot", "--agent", agent, "--model", model,
            "--tier", tier, "--id", worker_id,
        ]
        if skills_used:
            cmd.append("--skills-used")
        proc = _spawn_background(cmd, STATE_DIR / "worker_logs" / f"{worker_id}.log")
        records.append({
            "pid": proc.pid, "id": worker_id,
            "provider": agent, "model": model, "tier": tier,
            "mode": "cli_autopilot", "started_at": datetime.now(timezone.utc).isoformat(),
        })
        launched += 1
        print(f"  Started {worker_id} (PID {proc.pid})")
    save_worker_pids(records)
    print(f"\n[PULSE] {launched} autopilot worker(s) launched. Use 'pulse swarm status' to check.")


def launch_swarm_worker(provider_key: str, model: str, tier: str, count: int,
                        skills_used: bool = False):
    """Unified swarm dispatch: routes cloud Ollama to CLI workers, else headless/autopilot."""
    # License-based worker cap
    try:
        from pulse.core.license import LicenseManager
        limits = LicenseManager().get_limits()
        max_w = limits.get("max_workers")
        if max_w is not None and count > max_w:
            if max_w == 0:
                print("[PULSE] No active license. Activate at: pulse activate <KEY>")
                return
            print(f"[PULSE] Solo plan: capping to {max_w} workers (upgrade for unlimited)")
            count = max_w
    except Exception:
        pass  # Don't block swarm on license errors

    if provider_key == "ollama":
        from pulse.cli.ollama_launch import is_cloud_model, launch_cloud_workers
        if is_cloud_model(model):
            return launch_cloud_workers(model, tier, count, skills_used=skills_used)
    if _is_headless_provider(provider_key):
        _launch_headless(provider_key, model, tier, count, skills_used=skills_used)
    else:
        _launch_cli_autopilot(provider_key, model, tier, count, skills_used=skills_used)


# ---------------------------------------------------------------------------
# Subcommands: stop, status, watch
# ---------------------------------------------------------------------------

def _swarm_stop():
    records = load_worker_pids()
    killed = 0
    if records:
        for r in records:
            pid = r.get("pid", 0)
            if is_process_running(pid):
                try:
                    kill_process(pid)
                    killed += 1
                except Exception:
                    pass
        save_worker_pids([])
    # Always clean worker/daemon entries from registry (even if no PIDs found)
    try:
        from pulse.workers.runner.worker_state import deregister_stale_workers
        cleaned = deregister_stale_workers()
    except Exception:
        cleaned = 0
    if killed == 0 and cleaned == 0:
        print("[PULSE] No workers to stop.")
        return
    print(f"[PULSE] Stopped {killed} worker(s). Cleaned {cleaned} registry entries.")


def _swarm_status():
    from pulse.workers.runner.worker_state import get_swarm_data
    data = get_swarm_data()
    workers = data["workers"]
    tasks = data["tasks"]

    print("=" * 60)
    print("  PULSE Swarm Status")
    print("=" * 60)
    if not workers:
        print("\n  Workers: none registered")
    else:
        for w in workers:
            status = "RUNNING" if w["status"] == "running" else "DEAD"
            mode = w.get("mode", "headless")[:8]
            print(f"  {w['id']:<25} PID {w['pid']:<8} {status:<8} "
                  f"[{w['tier']}] {w['model']} ({mode})")
        print(f"\n  Workers: {data['alive']} alive / {data['total']} registered")
    print(f"\n  Tasks:  {tasks['open']} open | {tasks['active']} active | {tasks['done']} done")
    if data["recent_completions"]:
        print(f"\n  Recent completions:")
        for c in data["recent_completions"]:
            print(f"    [{c['task_id']}] {c['title']} -- by {c['completed_by']}")
    act = data["last_activity"]
    if act:
        age = act["age_seconds"]
        if age < 60:
            age_str = f"{age}s ago"
        elif age < 3600:
            age_str = f"{age // 60}m ago"
        else:
            age_str = f"{age // 3600}h ago"
        print(f"\n  Last activity: {act['worker']} ({age_str})")
    print("=" * 60)
    return tasks["open"], tasks["active"], tasks["done"]


def _swarm_watch():
    """Live auto-refreshing swarm dashboard."""
    prev_done = 0
    all_done_announced = False
    try:
        while True:
            os.system("cls" if os.name == "nt" else "clear")
            tasks_open, tasks_active, tasks_done = _swarm_status()
            total_tasks = tasks_open + tasks_active + tasks_done
            if total_tasks == 0:
                print("\n  IDLE -- No tasks on board. Post tasks with: pulse post \"task\"")
            elif tasks_open > 0 or tasks_active > 0:
                progress = tasks_done / total_tasks * 100 if total_tasks > 0 else 0
                bar_len = 30
                filled = int(bar_len * tasks_done / total_tasks) if total_tasks > 0 else 0
                bar = "=" * filled + "-" * (bar_len - filled)
                print(f"\n  WORKING [{bar}] {progress:.0f}%")
                print(f"  Refreshing every 3s... Press Ctrl+C to exit.")
            else:
                if not all_done_announced:
                    all_done_announced = True
                    print("\n  " + "=" * 50)
                    print("    ALL TASKS COMPLETED!")
                    print("  " + "=" * 50)
                    if tasks_done > prev_done:
                        print(f"    {tasks_done} tasks done by your crew.")
                print(f"\n  Board clear. Post more tasks or Ctrl+C to exit.")
            prev_done = tasks_done
            time.sleep(3)
    except KeyboardInterrupt:
        print("\n\n[PULSE] Watch stopped. Your crew is still running in the background.")
        print("  pulse swarm status  -- check on them")
        print("  pulse swarm stop    -- kill all workers")
        print("  pulse post \"task\"   -- give them more work")
