"""
Phase 72: Hierarchical Swarms

Lead agents can spawn and manage sub-swarms, creating a tree of coordinated
work groups.  Parent swarms track their children's progress and aggregate
statistics.

Storage:
  state/broker/swarm_hierarchy/{swarm_id}.json
"""

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

from .server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

HIERARCHY_DIR = BROKER_DIR / "swarm_hierarchy"


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


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


def _load(swarm_id: str) -> Optional[dict]:
    path = _swarm_path(swarm_id)
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save(swarm: dict) -> None:
    _ensure_dir()
    _write_json_sync(_swarm_path(swarm["swarm_id"]), swarm)


def _ensure_root(swarm_id: str) -> dict:
    """If the swarm record doesn't exist yet, create a root-level stub."""
    existing = _load(swarm_id)
    if existing is not None:
        return existing
    root = {
        "swarm_id": swarm_id,
        "name": swarm_id,
        "task": "",
        "lead_agent_id": "",
        "parent_swarm_id": None,
        "sub_swarms": [],
        "status": "active",
        "agents": [],
        "created_at": datetime.now(timezone.utc).isoformat(),
        "completed_at": None,
        "result_summary": "",
    }
    _save(root)
    return root


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


def create_sub_swarm(
    parent_swarm_id: str,
    name: str,
    task: str,
    lead_agent_id: str,
) -> dict:
    """Create a child swarm under a parent.  Returns the new sub-swarm record."""
    _ensure_dir()
    parent = _ensure_root(parent_swarm_id)

    sub_id = uuid.uuid4().hex[:12]
    now = datetime.now(timezone.utc).isoformat()

    sub = {
        "swarm_id": sub_id,
        "name": name,
        "task": task,
        "lead_agent_id": lead_agent_id,
        "parent_swarm_id": parent_swarm_id,
        "sub_swarms": [],
        "status": "active",
        "agents": [lead_agent_id],
        "created_at": now,
        "completed_at": None,
        "result_summary": "",
    }
    _save(sub)

    # Register in parent
    parent["sub_swarms"].append(sub_id)
    _save(parent)

    return sub


def get_sub_swarms(parent_swarm_id: str) -> list[dict]:
    """List all immediate child swarms of a parent."""
    parent = _load(parent_swarm_id)
    if parent is None:
        return []
    children: list[dict] = []
    for sid in parent.get("sub_swarms", []):
        child = _load(sid)
        if child:
            children.append(child)
    return children


def get_parent(swarm_id: str) -> Optional[dict]:
    """Return the parent swarm record, or None if this is a root swarm."""
    swarm = _load(swarm_id)
    if swarm is None:
        return None
    pid = swarm.get("parent_swarm_id")
    if pid is None:
        return None
    return _load(pid)


def get_hierarchy(root_swarm_id: str) -> dict:
    """Return the full tree of swarms starting from *root_swarm_id*.

    Each node has a ``children`` key containing nested sub-swarm trees.
    """
    swarm = _load(root_swarm_id)
    if swarm is None:
        return {}

    node = {**swarm, "children": []}
    for sid in swarm.get("sub_swarms", []):
        child_tree = get_hierarchy(sid)
        if child_tree:
            node["children"].append(child_tree)
    return node


def propagate_status(swarm_id: str) -> dict:
    """When a sub-swarm completes or fails, propagate to the parent.

    Returns ``{"propagated": True, "parent_swarm_id": ...}`` or an error.
    """
    swarm = _load(swarm_id)
    if swarm is None:
        return {"error": "Swarm not found"}

    pid = swarm.get("parent_swarm_id")
    if pid is None:
        return {"propagated": False, "reason": "No parent swarm"}

    parent = _load(pid)
    if parent is None:
        return {"error": "Parent swarm not found"}

    # Check if all children of parent are complete
    all_complete = True
    any_failed = False
    for sid in parent.get("sub_swarms", []):
        child = _load(sid)
        if child is None:
            continue
        if child["status"] == "failed":
            any_failed = True
        if child["status"] not in ("completed", "dissolved", "failed"):
            all_complete = False

    if all_complete and parent["sub_swarms"]:
        parent["status"] = "failed" if any_failed else "completed"
        parent["completed_at"] = datetime.now(timezone.utc).isoformat()
    _save(parent)

    return {"propagated": True, "parent_swarm_id": pid, "parent_status": parent["status"]}


def dissolve_sub_swarm(swarm_id: str, reason: str = "") -> dict:
    """Dissolve a sub-swarm and report results to its parent."""
    swarm = _load(swarm_id)
    if swarm is None:
        return {"error": "Swarm not found"}

    now = datetime.now(timezone.utc).isoformat()
    swarm["status"] = "dissolved"
    swarm["completed_at"] = now
    swarm["result_summary"] = reason or "dissolved"
    _save(swarm)

    # Propagate to parent
    if swarm.get("parent_swarm_id"):
        propagate_status(swarm_id)

    return {
        "swarm_id": swarm_id,
        "status": "dissolved",
        "reason": reason,
        "dissolved_at": now,
    }


def get_hierarchy_stats(root_swarm_id: str) -> dict:
    """Aggregate stats across the full hierarchy tree.

    Returns total agents, total sub-swarms, and completion percentage.
    """
    tree = get_hierarchy(root_swarm_id)
    if not tree:
        return {"error": "Swarm not found"}

    stats = {"total_agents": 0, "total_sub_swarms": 0, "completed": 0, "total": 0}

    def _walk(node: dict) -> None:
        stats["total"] += 1
        agents = set(node.get("agents", []))
        stats["total_agents"] += len(agents)
        if node.get("status") in ("completed", "dissolved"):
            stats["completed"] += 1
        for child in node.get("children", []):
            stats["total_sub_swarms"] += 1
            _walk(child)

    _walk(tree)

    stats["completion_pct"] = (
        round(stats["completed"] / stats["total"] * 100, 1) if stats["total"] else 0.0
    )
    return stats
