"""
Phase 78: Swarm Telemetry and Heatmaps

Records fine-grained telemetry events for every swarm: task starts/completions,
idle periods, handoffs, failures.  Provides timeline views, bottleneck
detection, heatmaps, throughput metrics, efficiency scores, and cross-swarm
comparison.

Storage:
  state/broker/telemetry/{swarm_id}.jsonl
"""

import json
import uuid
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from .server import BROKER_DIR, _append_jsonl_sync, _read_jsonl_sync

TELEMETRY_DIR = BROKER_DIR / "telemetry"

VALID_EVENT_TYPES = {
    "task_start",
    "task_complete",
    "task_fail",
    "idle",
    "blocked",
    "message",
    "handoff",
}


def _ensure_dir() -> None:
    TELEMETRY_DIR.mkdir(parents=True, exist_ok=True)


def _swarm_path(swarm_id: str) -> Path:
    return TELEMETRY_DIR / f"{swarm_id}.jsonl"


def _load_events(swarm_id: str) -> list[dict]:
    return _read_jsonl_sync(_swarm_path(swarm_id))


# ── Public API ────────────────────────────────────────────────────────────


def record_event(
    swarm_id: str,
    agent_id: str,
    event_type: str,
    data: Optional[dict] = None,
) -> dict:
    """Log a telemetry event for an agent within a swarm.

    Valid *event_type* values: task_start, task_complete, task_fail, idle,
    blocked, message, handoff.
    """
    if event_type not in VALID_EVENT_TYPES:
        return {"ok": False, "error": f"Invalid event_type '{event_type}'. "
                f"Must be one of: {', '.join(sorted(VALID_EVENT_TYPES))}"}

    _ensure_dir()
    event = {
        "event_id": uuid.uuid4().hex[:12],
        "swarm_id": swarm_id,
        "agent_id": agent_id,
        "event_type": event_type,
        "data": data or {},
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    _append_jsonl_sync(_swarm_path(swarm_id), event)
    return {"ok": True, "event": event}


def get_agent_timeline(swarm_id: str, agent_id: str) -> list[dict]:
    """Chronological event list for one agent in a swarm."""
    events = _load_events(swarm_id)
    return [e for e in events if e.get("agent_id") == agent_id]


def get_swarm_timeline(swarm_id: str) -> list[dict]:
    """All events for a swarm, sorted chronologically."""
    return _load_events(swarm_id)


def get_bottlenecks(swarm_id: str) -> dict:
    """Detect bottlenecks in a swarm.

    Returns agents with long idle periods, abnormally slow tasks, and
    frequent handoffs.
    """
    events = _load_events(swarm_id)
    if not events:
        return {"idle_agents": [], "slow_tasks": [], "handoff_heavy": []}

    # Track per-agent stats
    agent_idle_count: dict[str, int] = defaultdict(int)
    agent_blocked_count: dict[str, int] = defaultdict(int)
    agent_handoff_count: dict[str, int] = defaultdict(int)
    task_durations: list[dict] = []

    # Task start tracking: agent -> most recent task_start timestamp
    task_starts: dict[str, str] = {}

    for ev in events:
        agent = ev.get("agent_id", "")
        etype = ev.get("event_type", "")

        if etype == "idle":
            agent_idle_count[agent] += 1
        elif etype == "blocked":
            agent_blocked_count[agent] += 1
        elif etype == "handoff":
            agent_handoff_count[agent] += 1
        elif etype == "task_start":
            task_starts[agent] = ev.get("timestamp", "")
        elif etype == "task_complete":
            start_ts = task_starts.pop(agent, None)
            if start_ts:
                try:
                    s = datetime.fromisoformat(start_ts)
                    e = datetime.fromisoformat(ev.get("timestamp", ""))
                    duration_sec = (e - s).total_seconds()
                    task_durations.append({
                        "agent_id": agent,
                        "duration_sec": duration_sec,
                        "event_id": ev.get("event_id", ""),
                    })
                except (ValueError, TypeError):
                    pass

    # Identify idle agents (more than 2 idle events)
    idle_agents = [
        {"agent_id": a, "idle_events": c}
        for a, c in agent_idle_count.items() if c > 2
    ]

    # Identify slow tasks (> 2x average duration)
    slow_tasks: list[dict] = []
    if task_durations:
        avg = sum(t["duration_sec"] for t in task_durations) / len(task_durations)
        slow_tasks = [t for t in task_durations if t["duration_sec"] > avg * 2]

    # Frequent handoffs
    handoff_heavy = [
        {"agent_id": a, "handoff_count": c}
        for a, c in agent_handoff_count.items() if c > 2
    ]

    return {
        "idle_agents": idle_agents,
        "slow_tasks": slow_tasks,
        "handoff_heavy": handoff_heavy,
    }


def get_heatmap(swarm_id: str) -> dict:
    """Activity heatmap by agent and hour bucket.

    Returns ``{"agents": [...], "hours": [0..23], "grid": {agent: {hour: count}}}``.
    """
    events = _load_events(swarm_id)
    agents: set[str] = set()
    grid: dict[str, dict[int, int]] = defaultdict(lambda: defaultdict(int))

    for ev in events:
        agent = ev.get("agent_id", "")
        agents.add(agent)
        try:
            ts = datetime.fromisoformat(ev.get("timestamp", ""))
            hour = ts.hour
            grid[agent][hour] += 1
        except (ValueError, TypeError):
            continue

    # Convert defaultdicts to plain dicts for JSON serialisation
    plain_grid = {a: dict(hours) for a, hours in grid.items()}

    return {
        "agents": sorted(agents),
        "hours": list(range(24)),
        "grid": plain_grid,
    }


def get_throughput(swarm_id: str, bucket_minutes: int = 60) -> list[dict]:
    """Tasks completed per time bucket.

    Returns a list of ``{"bucket_start": iso, "completed": int}`` dicts.
    """
    events = _load_events(swarm_id)
    completions: list[datetime] = []
    for ev in events:
        if ev.get("event_type") == "task_complete":
            try:
                completions.append(datetime.fromisoformat(ev["timestamp"]))
            except (ValueError, KeyError):
                continue

    if not completions:
        return []

    completions.sort()
    bucket_sec = bucket_minutes * 60
    buckets: list[dict] = []

    first = completions[0]
    # Truncate to bucket boundary
    bucket_start = first.replace(minute=0, second=0, microsecond=0)

    idx = 0
    while idx < len(completions):
        bucket_end_ts = bucket_start.timestamp() + bucket_sec
        count = 0
        while idx < len(completions) and completions[idx].timestamp() < bucket_end_ts:
            count += 1
            idx += 1
        buckets.append({
            "bucket_start": bucket_start.isoformat(),
            "completed": count,
        })
        bucket_start = datetime.fromtimestamp(bucket_end_ts, tz=bucket_start.tzinfo)

    return buckets


def get_efficiency_score(swarm_id: str) -> dict:
    """0-100 efficiency score based on idle time, completion rate, handoffs.

    Components:
    - idle_ratio (lower is better, 40% weight)
    - completion_rate (higher is better, 40% weight)
    - handoff_penalty (fewer is better, 20% weight)
    """
    events = _load_events(swarm_id)
    if not events:
        return {"score": 0, "breakdown": {}, "note": "No telemetry data"}

    total = len(events)
    idle_count = sum(1 for e in events if e.get("event_type") == "idle")
    complete_count = sum(1 for e in events if e.get("event_type") == "task_complete")
    fail_count = sum(1 for e in events if e.get("event_type") == "task_fail")
    start_count = sum(1 for e in events if e.get("event_type") == "task_start")
    handoff_count = sum(1 for e in events if e.get("event_type") == "handoff")

    # Idle ratio: 0 = no idle, 1 = all idle  ->  score = 100 * (1 - ratio)
    idle_ratio = idle_count / total if total else 0
    idle_score = 100 * (1 - idle_ratio)

    # Completion rate: completed / started
    completion_rate = (complete_count / start_count) if start_count else 0
    completion_score = min(100, completion_rate * 100)

    # Handoff penalty: more handoffs = lower score
    handoff_ratio = handoff_count / total if total else 0
    handoff_score = max(0, 100 - handoff_ratio * 200)

    weighted = (idle_score * 0.4) + (completion_score * 0.4) + (handoff_score * 0.2)
    score = round(max(0, min(100, weighted)))

    return {
        "score": score,
        "breakdown": {
            "idle_score": round(idle_score, 1),
            "completion_score": round(completion_score, 1),
            "handoff_score": round(handoff_score, 1),
            "idle_ratio": round(idle_ratio, 3),
            "completion_rate": round(completion_rate, 3),
            "handoff_ratio": round(handoff_ratio, 3),
        },
        "event_counts": {
            "total": total,
            "idle": idle_count,
            "task_start": start_count,
            "task_complete": complete_count,
            "task_fail": fail_count,
            "handoff": handoff_count,
        },
    }


def compare_swarms(swarm_ids: list[str]) -> list[dict]:
    """Side-by-side telemetry comparison for multiple swarms.

    Returns a list of per-swarm summaries including efficiency score,
    event counts, and throughput.
    """
    results = []
    for sid in swarm_ids:
        eff = get_efficiency_score(sid)
        timeline = get_swarm_timeline(sid)
        throughput = get_throughput(sid)
        total_completed = sum(b["completed"] for b in throughput) if throughput else 0
        results.append({
            "swarm_id": sid,
            "total_events": len(timeline),
            "efficiency_score": eff["score"],
            "total_completed": total_completed,
            "event_counts": eff.get("event_counts", {}),
        })
    return results
