#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Task board lifecycle management for PULSE dispatcher.
Handles stale claim recovery, dependency unblocking, capability routing, archival.
Extracted from pulse_dispatcher.py per Law 7.

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


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 apply_capability_routing(board: dict, verbose: bool = False) -> bool:
    """Mark tasks with eligible agents based on capabilities.
    If no agents match, falls back to tier-only matching (eligible_agents=None)."""
    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") != "open":
                continue

            req_caps = set(task.get("required_capabilities", []))
            if not req_caps:
                if "eligible_agents" in task:
                    task.pop("eligible_agents")
                    modified = True
                continue

            eligible = []
            for agent_id, agent in registry.items():
                # Only consider agents that have heartbeated recently
                heartbeat = agent.get("last_heartbeat", "")
                if not heartbeat:
                    continue
                try:
                    hb_time = datetime.fromisoformat(heartbeat)
                    if (datetime.now(timezone.utc) - hb_time).total_seconds() > STALE_THRESHOLD:
                        continue
                except Exception:
                    continue

                agent_caps = set(agent.get("capabilities", []))
                if req_caps.intersection(agent_caps):
                    eligible.append(agent_id)

            # Fallback: if no match, remove eligible_agents to allow any matching tier
            if not eligible:
                if "eligible_agents" in task:
                    task.pop("eligible_agents")
                    modified = True
            else:
                if task.get("eligible_agents") != sorted(eligible):
                    task["eligible_agents"] = sorted(eligible)
                    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 score_agent(agent: dict, task: dict) -> float:
    """Compute fitness score for agent-task pairing."""
    score = 1.0

    # Capability overlap
    agent_caps = set(agent.get("capabilities", []))
    task_caps = set(task.get("required_capabilities", []))
    if task_caps:
        overlap = len(agent_caps & task_caps) / len(task_caps)
        score += overlap
    else:
        score += 0.5  # No requirements = any agent can do it

    # Domain history (legacy — reads dead data until registry schema is fixed)
    domain = task.get("domain", "")
    history = agent.get("history", {}).get("domains", {}).get(domain, {})
    wins = history.get("wins", 0)
    fails = history.get("fails", 0)
    score += min(wins * 0.1, 1.0)
    score -= min(fails * 0.15, 0.5)

    # D12: Hebbian model-domain fitness boost
    model = (agent.get("model") or "").strip()
    if model and domain:
        try:
            from pulse.daemons.synthesis.hebbian_balancer import get_model_fitness
            hw = get_model_fitness(model, domain)
            # Scale: 0.5 (cold start) adds +0, 0.95 (proven) adds +0.45
            score += (hw - 0.5) * 1.0
        except Exception:
            pass  # D12 not available = no boost

    return round(score, 3)


def score_agents_for_task(task: dict) -> list[tuple[str, float]]:
    """Score all registered agents for a given task. Returns sorted list."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    scores = []
    for agent_id, agent in registry.items():
        if agent.get("status") == "offline":
            continue
        s = score_agent(agent, task)
        scores.append((agent_id, s))
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores
