#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Task board maintenance for PULSE dispatcher.
Handles stale claim recovery, dependency unblocking, archival, and GC.
Split from dispatcher_lifecycle.py per Law 7.

Dependencies: Python stdlib only (Law 4)
"""
import json
import re
from datetime import datetime, timezone
from pathlib import Path

from pulse.core.brain_utils import (
    TASK_BOARD_FILE,
    AGENT_REGISTRY_FILE,
    STATE_DIR,
    load_json_dict,
    now_iso as _now_iso,
)

STALE_THRESHOLD = 300  # 5 minutes with no heartbeat
COMPLETED_DIR = STATE_DIR / "completed"

# H2: Agent whitelist — daemons may register but NEVER claim tasks (SoC)
DAEMON_AGENTS = frozenset({
    "claude_daemon", "gemini_daemon", "grok_daemon",
    "codex_daemon", "bridge_daemon", "neural_daemon",
    "compliance_daemon", "dispatcher_daemon",
})
# Workers + daemons are allowed for registration
ALLOWED_AGENTS = DAEMON_AGENTS | frozenset({"claude", "gemini", "codex", "grok"})
# Swarm workers use dynamic IDs like "worker_google_12345" — match by prefix
ALLOWED_AGENT_PREFIXES = ("worker_", "claude-", "gemini-", "grok-", "codex-",
                          "minimax-", "kimi-", "glm-", "qwen", "ollama",
                          "devstral-", "gpt-oss", "nemotron-", "rnj-", "ministral-")


def check_stale_claims(board: dict, verbose: bool = False) -> bool:
    """Release claims with no heartbeat > STALE_THRESHOLD. Returns True if modified."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    modified = False

    for dispatch in board.get("dispatches", []):
        for task in dispatch.get("tasks", []):
            if task.get("status") not in ("claimed", "active"):
                continue

            agent_id = task.get("claimed_by")
            if not agent_id:
                continue

            # SoC: Daemons must NEVER hold task claims — force-release
            if agent_id in DAEMON_AGENTS:
                if verbose:
                    print(f"[LIFECYCLE] SoC: Releasing daemon claim '{agent_id}' "
                          f"on task {task.get('task_id', '?')}")
                task["status"] = "open"
                task["claimed_by"] = None
                task["claimed_at"] = None
                modified = True
                continue

            # H2: Reject claims from unknown agents (allow swarm prefixes)
            if agent_id not in ALLOWED_AGENTS and not agent_id.startswith(ALLOWED_AGENT_PREFIXES):
                if verbose:
                    print(f"[LIFECYCLE] SECURITY: Rejecting claim from unknown agent '{agent_id}' "
                          f"on task {task.get('task_id', '?')}")
                task["status"] = "open"
                task["claimed_by"] = None
                task["claimed_at"] = None
                modified = True
                continue

            agent = registry.get(agent_id, {})
            heartbeat = agent.get("last_heartbeat", "")
            claimed_at = task.get("claimed_at", "")

            # Check heartbeat first, fall back to claimed_at for swarm agents
            # (swarm agents don't register heartbeats — they're short-lived CLI calls)
            check_time = heartbeat or claimed_at
            if not check_time:
                stale = True
            else:
                try:
                    t = datetime.fromisoformat(check_time.replace("Z", "+00:00"))
                    now = datetime.now(timezone.utc)
                    stale = (now - t).total_seconds() > STALE_THRESHOLD
                except (ValueError, TypeError):
                    stale = True

            if stale:
                if verbose:
                    print(f"[LIFECYCLE] Releasing stale claim: {task['task_id']} "
                          f"(agent: {agent_id})")
                task["status"] = "open"
                task["claimed_by"] = None
                task["claimed_at"] = None
                modified = True

    return modified


def unblock_dependents(board: dict, verbose: bool = False) -> bool:
    """Promote blocked -> open when all dependencies are done. Returns True if modified."""
    modified = False

    # Build map of task_id -> status across all dispatches
    status_map = {}
    for dispatch in board.get("dispatches", []):
        for task in dispatch.get("tasks", []):
            status_map[task["task_id"]] = task.get("status", "open")

    for dispatch in board.get("dispatches", []):
        for task in dispatch.get("tasks", []):
            if task.get("status") != "blocked":
                continue

            deps = task.get("depends_on", [])
            if not deps:
                task["status"] = "open"
                modified = True
                continue

            all_done = all(status_map.get(dep) == "done" for dep in deps)
            if all_done:
                if verbose:
                    print(f"[LIFECYCLE] Unblocking: {task['task_id']} "
                          f"(deps satisfied: {deps})")
                task["status"] = "open"
                modified = True

    return modified


def archive_completed(board: dict, verbose: bool = False) -> bool:
    """Move fully-completed dispatches to completed/ dir. Returns True if modified."""
    modified = False
    remaining = []

    for dispatch in board.get("dispatches", []):
        tasks = dispatch.get("tasks", [])
        if not tasks:
            remaining.append(dispatch)
            continue

        all_terminal = all(
            t.get("status") in ("done", "failed") for t in tasks
        )

        if all_terminal:
            # Archive this dispatch
            COMPLETED_DIR.mkdir(parents=True, exist_ok=True)
            # G3: Sanitize dispatch_id to prevent path traversal
            safe_id = re.sub(r'[^a-zA-Z0-9_\-]', '_', dispatch.get('dispatch_id', 'unknown'))
            archive_path = COMPLETED_DIR / f"{safe_id}.json"
            # Verify resolved path stays within COMPLETED_DIR
            if not archive_path.resolve().is_relative_to(COMPLETED_DIR.resolve()):
                if verbose:
                    print(f"[LIFECYCLE] SECURITY: Path traversal blocked for '{safe_id}'")
                remaining.append(dispatch)
                continue
            dispatch["archived_at"] = _now_iso()
            try:
                with open(archive_path, "w", encoding="utf-8") as f:
                    json.dump(dispatch, f, indent=2, ensure_ascii=False)
                if verbose:
                    print(f"[LIFECYCLE] Archived: {dispatch['dispatch_id']}")
                modified = True
            except OSError as e:
                if verbose:
                    print(f"[LIFECYCLE] Archive failed: {e}")
                remaining.append(dispatch)
        else:
            remaining.append(dispatch)

    if modified:
        board["dispatches"] = remaining

    return modified


def gc_stale_dispatches(board: dict, max_age_hours: int = 6, verbose: bool = False) -> bool:
    """Remove dispatches where ALL tasks are still 'open' and dispatch is older than max_age_hours.
    Archives them to completed/ with gc_reason. Returns True if modified."""
    modified = False
    remaining = []
    now = datetime.now(timezone.utc)

    for dispatch in board.get("dispatches", []):
        tasks = dispatch.get("tasks", [])
        dispatched_at = dispatch.get("dispatched_at", "")

        # Only GC if there are tasks, all are open, and dispatch is old enough
        if not tasks or not dispatched_at:
            remaining.append(dispatch)
            continue

        all_open = all(t.get("status") == "open" for t in tasks)
        if not all_open:
            remaining.append(dispatch)
            continue

        try:
            t = datetime.fromisoformat(dispatched_at.replace("Z", "+00:00"))
            age_hours = (now - t).total_seconds() / 3600
        except (ValueError, TypeError):
            remaining.append(dispatch)
            continue

        if age_hours < max_age_hours:
            remaining.append(dispatch)
            continue

        # Archive stale dispatch
        COMPLETED_DIR.mkdir(parents=True, exist_ok=True)
        safe_id = re.sub(r'[^a-zA-Z0-9_\-]', '_', dispatch.get('dispatch_id', 'unknown'))
        archive_path = COMPLETED_DIR / f"{safe_id}.json"
        if not archive_path.resolve().is_relative_to(COMPLETED_DIR.resolve()):
            if verbose:
                print(f"[LIFECYCLE] SECURITY: Path traversal blocked for '{safe_id}'")
            remaining.append(dispatch)
            continue

        dispatch["archived_at"] = _now_iso()
        dispatch["gc_reason"] = "stale_open"
        try:
            with open(archive_path, "w", encoding="utf-8") as f:
                json.dump(dispatch, f, indent=2, ensure_ascii=False)
            if verbose:
                print(f"[LIFECYCLE] GC stale dispatch: {dispatch.get('dispatch_id', '?')} "
                      f"({len(tasks)} tasks, {age_hours:.1f}h old)")
            modified = True
        except OSError as e:
            if verbose:
                print(f"[LIFECYCLE] GC archive failed: {e}")
            remaining.append(dispatch)

    if modified:
        board["dispatches"] = remaining

    return modified
