#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""State tracking and heartbeat operations for background workers."""
from __future__ import annotations

import json
import logging
import time
from pathlib import Path

from pulse.core.brain_utils import load_json_dict, save_json, AGENT_REGISTRY_FILE, TASK_BOARD_FILE, now_iso as _now_iso
from pulse.tasks.task_ops import update_registry, load_worker_pids, is_process_running


def register_worker(worker_id: str, model: str, tier: str, log: logging.Logger, tool_names: list[str]):
    """Register this worker in the system registry."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    caps = list(tool_names) + ["code"]
    m = model.lower()
    if "pro" in m or "opus" in m: caps.extend(["architecture", "security"])
    if "flash" in m or "haiku" in m: caps.extend(["docs", "testing"])
    if tier == "standard": caps.append("refactoring")

    registry[worker_id] = {
        "agent_id": worker_id, "model": model, "tier": tier,
        "type": "worker", "status": "idle", "capabilities": sorted(set(caps)),
        "last_heartbeat": _now_iso(), "current_task": None,
        "history": registry.get(worker_id, {}).get("history", {"tasks_completed": 0, "domains": {}}),
    }
    save_json(AGENT_REGISTRY_FILE, registry)
    log.info("Registered (tier=%s, model=%s)", tier, model)


def deregister_worker(worker_id: str):
    """Remove this worker from the registry on shutdown."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    if worker_id in registry:
        del registry[worker_id]
        save_json(AGENT_REGISTRY_FILE, registry)
    # Also remove its daemon shadow entry if present
    daemon_id = f"{worker_id}_daemon"
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    if daemon_id in registry:
        del registry[daemon_id]
        save_json(AGENT_REGISTRY_FILE, registry)


def deregister_stale_workers():
    """Remove ALL worker and daemon entries from the registry. Called by swarm stop."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    to_remove = [k for k, v in registry.items() if v.get("type") in ("worker", "daemon")]
    if not to_remove:
        # Fallback for entries without type field
        to_remove = [k for k in registry if k.endswith("_daemon") or k.startswith("worker_")]
        # Also catch swarm workers (qwen3-coder-*, gpt-5.3-codex_*)
        to_remove += [k for k, v in registry.items()
                      if k not in to_remove and v.get("type") not in ("agent",)
                      and not any(k == name for name in ("claude", "gemini", "codex", "grok"))]
    for k in to_remove:
        registry.pop(k, None)
    save_json(AGENT_REGISTRY_FILE, registry)
    return len(to_remove)


def send_heartbeat(worker_id: str, model: str, tier: str, last_heartbeat: float, heartbeat_interval: int, state_dir: Path, task_id: str | None = None, iteration: int | None = None, status: str = "idle") -> float:
    """Send a heartbeat if interval has elapsed. Returns new last_heartbeat time."""
    now = time.time()
    if now - last_heartbeat < heartbeat_interval:
        return last_heartbeat

    update_registry(worker_id, status=status, task_id=task_id)
    hb_dir = state_dir / "heartbeats"
    hb_dir.mkdir(parents=True, exist_ok=True)
    try:
        (hb_dir / f"{worker_id}.json").write_text(json.dumps({
            "agent_id": worker_id, "current_task_id": task_id,
            "iteration": iteration, "timestamp": _now_iso(),
            "status": status, "tier": tier, "model": model,
        }, indent=2), encoding="utf-8")
    except Exception:
        pass
    return now


def get_swarm_data() -> dict:
    """Return structured swarm state — single source of truth for CLI and API.

    Returns dict with keys: workers (list), tasks (dict), last_activity (dict|None).
    """
    records = load_worker_pids()
    board = load_json_dict(TASK_BOARD_FILE)

    # Find which workers are actively working on a task
    active_workers: set[str] = set()
    for d in board.get("dispatches", []):
        for t in d.get("tasks", []):
            if t.get("status") == "active":
                claimed = t.get("claimed_by", "")
                if claimed:
                    active_workers.add(claimed)

    # Worker list with liveness + task activity check
    workers = []
    for r in records:
        pid = r.get("pid", 0)
        alive = is_process_running(pid)
        wid = r.get("id", "?")
        if not alive:
            status = "dead"
        elif wid in active_workers:
            status = "running"
        else:
            status = "idle"
        workers.append({
            "id": wid,
            "pid": pid,
            "model": r.get("model", "unknown"),
            "tier": r.get("tier", "unknown"),
            "mode": r.get("mode", "headless"),
            "status": status,
            "started_at": r.get("started_at"),
        })

    # Task counts + recent completions
    tasks_open = tasks_active = tasks_done = 0
    recent: list[dict] = []
    for d in board.get("dispatches", []):
        for t in d.get("tasks", []):
            s = t.get("status", "")
            if s == "open":
                tasks_open += 1
            elif s == "active":
                tasks_active += 1
            elif s == "done":
                tasks_done += 1
                recent.append({
                    "task_id": t.get("task_id", t.get("id", "?")),
                    "title": t.get("title", "?")[:40],
                    "completed_by": t.get("completed_by") or t.get("claimed_by", "?"),
                })

    # Last log activity
    import time as _time
    from pulse.core.brain_utils import STATE_DIR as _state_dir
    last_activity = None
    log_dir = _state_dir / "worker_logs"
    if log_dir.exists():
        latest_log = None
        latest_mtime = 0.0
        for lf in log_dir.glob("*.log"):
            mt = lf.stat().st_mtime
            if mt > latest_mtime:
                latest_mtime = mt
                latest_log = lf
        if latest_log and latest_mtime > 0:
            last_activity = {
                "worker": latest_log.stem,
                "age_seconds": int(_time.time() - latest_mtime),
            }

    return {
        "workers": workers,
        "alive": sum(1 for w in workers if w["status"] != "dead"),
        "total": len(workers),
        "tasks": {"open": tasks_open, "active": tasks_active, "done": tasks_done},
        "recent_completions": recent[-5:],
        "last_activity": last_activity,
    }
