#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
PULSE API — Agents router. Endpoints for agent registry, workers, and system overview.
"""
from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from fastapi import APIRouter

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent
STATE_DIR = ROOT_DIR / "state"


from pulse.core.brain_utils import load_json_dict, AGENT_REGISTRY_FILE, TASK_BOARD_FILE

router = APIRouter(tags=["agents"])


STALE_THRESHOLD_SECONDS = 3600  # 1 hour — no heartbeat = offline

# Map internal status names → frontend-expected names
_STATUS_MAP = {"working": "running", "offline": "dead", "booting": "idle"}


def _resolve_status(data: dict) -> str:
    """Return real status based on heartbeat freshness. Maps to frontend values."""
    raw_status = data.get("status", "unknown")
    hb = data.get("last_heartbeat")
    if not hb:
        status = "offline" if raw_status in ("idle", "working", "booting") else raw_status
        return _STATUS_MAP.get(status, status)
    try:
        hb_time = datetime.fromisoformat(hb)
        if hb_time.tzinfo is None:
            hb_time = hb_time.replace(tzinfo=timezone.utc)
        age = (datetime.now(timezone.utc) - hb_time).total_seconds()
        if age > STALE_THRESHOLD_SECONDS and raw_status in ("idle", "working", "booting"):
            return _STATUS_MAP.get("offline", "dead")
    except (ValueError, TypeError):
        pass
    return _STATUS_MAP.get(raw_status, raw_status)


def _infer_type(agent_id: str) -> str:
    """Fallback for old registry entries without type field."""
    if agent_id.endswith("_daemon"):
        return "daemon"
    if agent_id.startswith("worker_"):
        return "worker"
    return "agent"


def _is_real_agent(agent_id: str, data: dict) -> bool:
    """Filter to only real agents using type field (fallback to heuristic)."""
    return data.get("type", _infer_type(agent_id)) == "agent"



@router.get("/v1/agents")
async def get_agents(include_dead: bool = False):
    """Return registered agents (excludes daemons and workers).
    By default only returns alive agents. Pass ?include_dead=true for all."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    agents = []
    for agent_id, data in registry.items():
        if not _is_real_agent(agent_id, data):
            continue
        status = _resolve_status(data)
        if not include_dead and status == "dead":
            continue
        agents.append({
            "id": agent_id,
            "model": data.get("model", "unknown"),
            "tier": data.get("tier", "unknown"),
            "type": data.get("type", _infer_type(agent_id)),
            "status": status,
            "current_task": data.get("current_task"),
            "last_heartbeat": data.get("last_heartbeat"),
            "capabilities": data.get("capabilities", []),
            "history": data.get("history", {}),
        })
    return {"count": len(agents), "agents": agents}


@router.get("/v1/agents/{agent_id}/heartbeat")
async def get_agent_heartbeat(agent_id: str):
    """Return latest heartbeat for a specific agent."""
    hb_file = STATE_DIR / "heartbeats" / f"{agent_id}.json"
    if not hb_file.exists():
        return {"agent_id": agent_id, "status": "no_heartbeat"}
    data = json.loads(hb_file.read_text(encoding="utf-8"))
    return data


@router.get("/v1/agents/{agent_id}/procedures")
async def get_agent_procedures(agent_id: str):
    """Return agent-specific procedure recommendations based on success/failure history."""
    from pulse.brain.agent_procedures import (
        get_agent_profile, match_procedures, get_weak_domains, get_strong_domains,
    )
    profile = get_agent_profile(agent_id)
    if not profile:
        return {"agent_id": agent_id, "status": "not_found"}
    return {
        "agent_id": agent_id,
        "profile": profile,
        "weak_domains": get_weak_domains(agent_id),
        "strong_domains": get_strong_domains(agent_id),
        "procedures": match_procedures(agent_id),
    }


@router.get("/v1/workers")
async def get_workers(include_dead: bool = False):
    """Return swarm workers. Uses shared get_swarm_data() — same source as CLI.
    By default only returns alive workers. Pass ?include_dead=true for all."""
    from pulse.workers.runner.worker_state import get_swarm_data
    data = get_swarm_data()
    workers = data["workers"]
    if not include_dead:
        workers = [w for w in workers if w["status"] != "dead"]
    return {"count": len(workers), "workers": workers}


@router.get("/v1/system/overview")
async def get_system_overview():
    """High-level system overview: agent count, task stats, brain health."""
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    board = load_json_dict(TASK_BOARD_FILE)

    # Count interactive agents from registry
    real_agents = {k: v for k, v in registry.items() if _is_real_agent(k, v)}
    alive_agents = {k: v for k, v in real_agents.items() if _resolve_status(v) != "dead"}
    interactive_total = len(alive_agents)
    interactive_active = sum(1 for a in alive_agents.values() if _resolve_status(a) in ("running", "idle"))
    interactive_working = sum(1 for a in alive_agents.values() if _resolve_status(a) == "running")

    # Count swarm workers
    from pulse.workers.runner.worker_state import get_swarm_data
    swarm = get_swarm_data()
    swarm_alive = swarm["alive"]
    swarm_running = sum(1 for w in swarm["workers"] if w["status"] == "running")

    # Unified workforce counts
    total_agents = interactive_total + swarm_alive
    active_agents = interactive_active + swarm_alive
    working_agents = interactive_working + swarm_running

    # Count tasks by status
    tasks_open = 0
    tasks_active = 0
    tasks_done = 0
    total_cost = 0.0
    for d in board.get("dispatches", []):
        for t in d.get("tasks", []):
            s = t.get("status", "")
            if s == "open" or s == "escalated":
                tasks_open += 1
            elif s == "active":
                tasks_active += 1
            elif s == "done":
                tasks_done += 1
            total_cost += t.get("estimated_cost_usd", 0.0)

    # Brain stats — read directly from Hippocampus
    brain_counts: dict[str, int] = {}
    hippo = ROOT_DIR / "Hippocampus"
    for label, path in [
        ("lessons", hippo / "Lessons" / "auto_lessons.md"),
        ("failures", hippo / "Failures" / "auto_fails.md"),
        ("patterns", hippo / "Patterns" / "auto_patterns.md"),
    ]:
        if path.exists():
            try:
                brain_counts[label] = path.read_text(encoding="utf-8").count("\n## ")
            except Exception:
                brain_counts[label] = 0

    evidence_count = 0
    evidence_dir = hippo / "Evidence"
    if evidence_dir.exists():
        for f in evidence_dir.rglob("*.json"):
            try:
                data = json.loads(f.read_text(encoding="utf-8"))
                if isinstance(data, list):
                    evidence_count += len(data)
            except Exception:
                pass

    wants_count = dont_wants_count = intents_count = anti_patterns_count = 0
    for label, path in [
        ("wants", evidence_dir / "Intents" / "wants.json"),
        ("dont_wants", evidence_dir / "Intents" / "dont_wants.json"),
        ("intents", evidence_dir / "Intents" / "strategic_intents.json"),
        ("anti_patterns", evidence_dir / "Anti_Patterns" / "anti_patterns_active.json"),
    ]:
        if path.exists():
            try:
                d = json.loads(path.read_text(encoding="utf-8"))
                cnt = len(d) if isinstance(d, list) else len(d.get("intents", d.get("items", []))) if isinstance(d, dict) else 0
            except Exception:
                cnt = 0
            if label == "wants": wants_count = cnt
            elif label == "dont_wants": dont_wants_count = cnt
            elif label == "intents": intents_count = cnt
            elif label == "anti_patterns": anti_patterns_count = cnt

    schemas_count = procedures_count = 0
    schemas_path = hippo / "Patterns" / "schemas.json"
    if schemas_path.exists():
        try:
            sd = json.loads(schemas_path.read_text(encoding="utf-8"))
            schemas_count = len(sd.get("schemas", [])) if isinstance(sd, dict) else (len(sd) if isinstance(sd, list) else 0)
        except Exception:
            pass
    procedures_path = hippo / "Evidence" / "Metrics" / "procedural_log.json"
    if procedures_path.exists():
        try:
            pd = json.loads(procedures_path.read_text(encoding="utf-8"))
            procedures_count = len(pd.get("procedures", [])) if isinstance(pd, dict) else (len(pd) if isinstance(pd, list) else 0)
        except Exception:
            pass

    graph_nodes = graph_edges = 0
    graph_path = hippo / "Knowledge" / "knowledge_graph.json"
    if graph_path.exists():
        try:
            gd = json.loads(graph_path.read_text(encoding="utf-8"))
            graph_nodes = len(gd.get("nodes", []))
            graph_edges = len(gd.get("edges", []))
        except Exception:
            pass

    total_knowledge = brain_counts.get("lessons", 0) + brain_counts.get("failures", 0) + brain_counts.get("patterns", 0)
    extraction_rate = (total_knowledge / evidence_count * 100) if evidence_count > 0 else 0.0

    return {
        "agents": {
            "total": total_agents,
            "active": active_agents,
            "working": working_agents,
        },
        "tasks": {
            "open": tasks_open,
            "active": tasks_active,
            "done": tasks_done,
        },
        "brain": {
            "lessons": brain_counts.get("lessons", 0),
            "failures": brain_counts.get("failures", 0),
            "patterns": brain_counts.get("patterns", 0),
            "schemas": schemas_count,
            "procedures": procedures_count,
            "wants": wants_count,
            "dont_wants": dont_wants_count,
            "strategic_intents": intents_count,
            "anti_patterns": anti_patterns_count,
            "evidence_items": evidence_count,
            "extraction_rate": round(extraction_rate, 2),
            "graph_nodes": graph_nodes,
            "graph_edges": graph_edges,
        },
        "cost_usd": round(total_cost, 4),
    }
