"""
Phase 71: Dynamic Role Reassignment

If an agent is stuck (no progress updates within a timeout), another agent
in the swarm can take over its role and task context automatically.

Storage:
  state/broker/dynamic_roles/{swarm_id}.json   — role assignments + progress
  state/broker/dynamic_roles/{swarm_id}.jsonl   — reassignment history
"""

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

ROLES_DIR = BROKER_DIR / "dynamic_roles"


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


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


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


def _load_state(swarm_id: str) -> dict:
    path = _state_path(swarm_id)
    if not path.exists():
        return {"swarm_id": swarm_id, "roles": {}, "progress": {}}
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return {"swarm_id": swarm_id, "roles": {}, "progress": {}}


def _save_state(state: dict) -> None:
    _ensure_dir()
    _write_json_sync(_state_path(state["swarm_id"]), state)


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


def assign_role(
    swarm_id: str,
    agent_id: str,
    role: str,
    task_description: str = "",
) -> dict:
    """Assign a role within a swarm to an agent."""
    _ensure_dir()
    state = _load_state(swarm_id)
    now = datetime.now(timezone.utc).isoformat()

    state["roles"][agent_id] = {
        "role": role,
        "task_description": task_description,
        "assigned_at": now,
        "status": "active",
    }
    # Initialise progress tracker
    state["progress"][agent_id] = {
        "last_update": now,
        "messages": [],
    }
    _save_state(state)
    return state["roles"][agent_id]


def get_role(swarm_id: str, agent_id: str) -> Optional[dict]:
    """Return the current role info for an agent, or None."""
    state = _load_state(swarm_id)
    return state["roles"].get(agent_id)


def record_progress(swarm_id: str, agent_id: str, message: str = "") -> dict:
    """Record a progress heartbeat for an agent."""
    state = _load_state(swarm_id)
    now = datetime.now(timezone.utc).isoformat()

    if agent_id not in state["progress"]:
        state["progress"][agent_id] = {"last_update": now, "messages": []}

    state["progress"][agent_id]["last_update"] = now
    if message:
        state["progress"][agent_id]["messages"].append(
            {"message": message, "timestamp": now}
        )
    _save_state(state)
    return {"agent_id": agent_id, "last_update": now}


def check_stuck(swarm_id: str, timeout_minutes: int = 20) -> list[dict]:
    """Find agents whose last progress update exceeds *timeout_minutes*."""
    state = _load_state(swarm_id)
    now = datetime.now(timezone.utc)
    stuck: list[dict] = []

    for agent_id, prog in state.get("progress", {}).items():
        last = datetime.fromisoformat(prog["last_update"])
        elapsed = (now - last).total_seconds() / 60.0
        if elapsed >= timeout_minutes:
            role_info = state["roles"].get(agent_id, {})
            stuck.append({
                "agent_id": agent_id,
                "role": role_info.get("role", "unknown"),
                "minutes_since_update": round(elapsed, 1),
                "task_description": role_info.get("task_description", ""),
            })
    return stuck


def get_available_agents(swarm_id: str) -> list[str]:
    """Return agents in the swarm that are not currently assigned critical roles.

    An agent is considered *available* if it has no role entry, or its role
    status is ``idle`` or ``completed``.
    """
    state = _load_state(swarm_id)
    available: list[str] = []

    # Collect all known agent ids (from roles + progress)
    all_agents = set(state.get("roles", {}).keys()) | set(
        state.get("progress", {}).keys()
    )

    for agent_id in sorted(all_agents):
        role_info = state["roles"].get(agent_id, {})
        status = role_info.get("status", "idle")
        if status in ("idle", "completed", "reassigned"):
            available.append(agent_id)
    return available


def reassign_role(
    swarm_id: str,
    from_agent: str,
    to_agent: str,
    reason: str = "",
) -> dict:
    """Transfer a role (and its task context) from one agent to another."""
    _ensure_dir()
    state = _load_state(swarm_id)
    now = datetime.now(timezone.utc).isoformat()

    old_role = state["roles"].get(from_agent)
    if old_role is None:
        return {"error": f"Agent {from_agent} has no role to transfer"}

    # Mark source agent as reassigned
    state["roles"][from_agent]["status"] = "reassigned"
    state["roles"][from_agent]["reassigned_at"] = now

    # Assign the role to the target agent
    state["roles"][to_agent] = {
        "role": old_role["role"],
        "task_description": old_role.get("task_description", ""),
        "assigned_at": now,
        "status": "active",
        "inherited_from": from_agent,
    }
    state["progress"][to_agent] = {"last_update": now, "messages": []}

    _save_state(state)

    # Append to reassignment history
    record = {
        "id": uuid.uuid4().hex[:12],
        "swarm_id": swarm_id,
        "from_agent": from_agent,
        "to_agent": to_agent,
        "role": old_role["role"],
        "task_description": old_role.get("task_description", ""),
        "reason": reason,
        "timestamp": now,
    }
    _append_jsonl_sync(_history_path(swarm_id), record)
    return record


def auto_reassign_stuck(
    swarm_id: str, timeout_minutes: int = 20
) -> list[dict]:
    """Find stuck agents and reassign their roles to available replacements.

    Returns a list of reassignment records (one per stuck agent that was
    successfully reassigned).
    """
    stuck = check_stuck(swarm_id, timeout_minutes)
    available = get_available_agents(swarm_id)
    reassignments: list[dict] = []

    for entry in stuck:
        if not available:
            break
        replacement = available.pop(0)
        if replacement == entry["agent_id"]:
            # skip self
            if not available:
                break
            replacement = available.pop(0)
        result = reassign_role(
            swarm_id,
            from_agent=entry["agent_id"],
            to_agent=replacement,
            reason=f"auto: stuck for {entry['minutes_since_update']} min",
        )
        if "error" not in result:
            reassignments.append(result)
    return reassignments


def get_reassignment_history(swarm_id: str) -> list[dict]:
    """Return the full reassignment log for a swarm."""
    _ensure_dir()
    return _read_jsonl_sync(_history_path(swarm_id))
