"""
Phase 69: Rigid State Machine for Swarm Lifecycle Coordination.

Manages swarm sessions through a strict state machine:
  PLANNING -> RECRUITING -> EXECUTING -> REVIEWING -> COMPLETED | FAILED | BLOCKED

Every transition is validated and logged. Agents can be assigned roles,
report progress, and sessions can be dissolved.

Storage:
  state/broker/swarm_sessions/{session_id}.json
  state/broker/swarm_sessions/history.jsonl
"""

import json
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

# ── Paths ──────────────────────────────────────────────────────────────────

SESSIONS_DIR = BROKER_DIR / "swarm_sessions"
SESSIONS_DIR.mkdir(parents=True, exist_ok=True)

HISTORY_PATH = SESSIONS_DIR / "history.jsonl"

# ── State machine definitions ──────────────────────────────────────────────

STATES = ("PLANNING", "RECRUITING", "EXECUTING", "REVIEWING", "COMPLETED", "FAILED", "BLOCKED")
TERMINAL_STATES = ("COMPLETED", "FAILED")
ROLES = ("lead", "worker", "reviewer")

# Valid transitions: from_state -> set of allowed to_states
VALID_TRANSITIONS: Dict[str, set] = {
    "PLANNING": {"RECRUITING", "FAILED"},
    "RECRUITING": {"EXECUTING", "FAILED"},
    "EXECUTING": {"REVIEWING", "BLOCKED", "FAILED"},
    "REVIEWING": {"COMPLETED", "EXECUTING", "FAILED"},
    "BLOCKED": {"EXECUTING", "FAILED"},
}

DEFAULT_BLOCK_TIMEOUT = 1800  # 30 minutes in seconds


# ── Internal helpers ───────────────────────────────────────────────────────


def _read_session(session_id: str) -> Optional[dict]:
    """Read a swarm session by ID."""
    path = SESSIONS_DIR / f"{session_id}.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _write_session(session: dict) -> None:
    """Write a swarm session to its JSON file."""
    path = SESSIONS_DIR / f"{session['session_id']}.json"
    _write_json_sync(path, session)


def _log_event(session_id: str, event_type: str, details: dict) -> None:
    """Append an event to the history log."""
    event = {
        "session_id": session_id,
        "event_type": event_type,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "epoch": time.time(),
        **details,
    }
    _append_jsonl_sync(HISTORY_PATH, event)


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


def create_swarm_session(
    name: str,
    task_description: str,
    required_skills: Optional[List[str]] = None,
    max_agents: int = 5,
) -> str:
    """Create a swarm session in PLANNING state. Returns session_id."""
    if not name.strip():
        raise ValueError("Name cannot be empty")
    if not task_description.strip():
        raise ValueError("Task description cannot be empty")
    if max_agents < 1:
        raise ValueError("max_agents must be at least 1")

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

    session = {
        "session_id": session_id,
        "name": name,
        "task_description": task_description,
        "required_skills": required_skills or [],
        "max_agents": max_agents,
        "state": "PLANNING",
        "agents": {},  # agent_id -> {"role": str, "joined_at": str}
        "progress": {},  # agent_id -> {"pct": float, "message": str, "updated_at": str}
        "created_at": now,
        "updated_at": now,
        "dissolved_at": None,
        "dissolve_reason": None,
        "block_timeout": DEFAULT_BLOCK_TIMEOUT,
    }

    _write_session(session)
    _log_event(session_id, "created", {
        "name": name,
        "state": "PLANNING",
    })

    return session_id


def get_state(session_id: str) -> str:
    """Returns the current state of a swarm session."""
    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")
    return session["state"]


def transition(session_id: str, new_state: str, reason: str = "") -> dict:
    """Move to a new state. Validates transitions. Returns updated session."""
    if new_state not in STATES:
        raise ValueError(f"Invalid state '{new_state}'. Must be one of: {STATES}")

    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    current = session["state"]
    if current in TERMINAL_STATES:
        raise ValueError(f"Session is in terminal state '{current}' and cannot transition")

    allowed = VALID_TRANSITIONS.get(current, set())
    if new_state not in allowed:
        raise ValueError(
            f"Invalid transition from '{current}' to '{new_state}'. "
            f"Allowed: {sorted(allowed)}"
        )

    old_state = current
    session["state"] = new_state
    session["updated_at"] = datetime.now(timezone.utc).isoformat()
    _write_session(session)

    _log_event(session_id, "transition", {
        "from_state": old_state,
        "to_state": new_state,
        "reason": reason,
    })

    return session


def assign_agent(session_id: str, agent_id: str, role: str = "worker") -> dict:
    """Add an agent to the swarm. Returns updated session."""
    if role not in ROLES:
        raise ValueError(f"Invalid role '{role}'. Must be one of: {ROLES}")

    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    if session["state"] in TERMINAL_STATES:
        raise ValueError(f"Cannot assign agents to a session in state '{session['state']}'")

    if agent_id in session["agents"]:
        raise ValueError(f"Agent '{agent_id}' is already assigned to this session")

    if len(session["agents"]) >= session["max_agents"]:
        raise ValueError(
            f"Session already has {len(session['agents'])}/{session['max_agents']} agents"
        )

    now = datetime.now(timezone.utc).isoformat()
    session["agents"][agent_id] = {
        "role": role,
        "joined_at": now,
    }
    session["updated_at"] = now
    _write_session(session)

    _log_event(session_id, "agent_assigned", {
        "agent_id": agent_id,
        "role": role,
    })

    return session


def remove_agent(session_id: str, agent_id: str, reason: str = "") -> dict:
    """Remove an agent from the swarm. Returns updated session."""
    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    if agent_id not in session["agents"]:
        raise ValueError(f"Agent '{agent_id}' is not in this session")

    del session["agents"][agent_id]
    # Also remove progress
    session["progress"].pop(agent_id, None)
    session["updated_at"] = datetime.now(timezone.utc).isoformat()
    _write_session(session)

    _log_event(session_id, "agent_removed", {
        "agent_id": agent_id,
        "reason": reason,
    })

    return session


def get_agents(session_id: str) -> List[dict]:
    """List agents with roles in a swarm session."""
    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    agents = []
    for agent_id, info in session["agents"].items():
        agent = {
            "agent_id": agent_id,
            "role": info["role"],
            "joined_at": info["joined_at"],
        }
        if agent_id in session.get("progress", {}):
            agent["progress"] = session["progress"][agent_id]
        agents.append(agent)

    return agents


def report_progress(
    session_id: str,
    agent_id: str,
    progress_pct: float,
    status_message: str = "",
) -> dict:
    """Agent reports progress. Returns updated session."""
    if not 0.0 <= progress_pct <= 100.0:
        raise ValueError("progress_pct must be between 0 and 100")

    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    if agent_id not in session["agents"]:
        raise ValueError(f"Agent '{agent_id}' is not in this session")

    if session["state"] in TERMINAL_STATES:
        raise ValueError(f"Cannot report progress for a session in state '{session['state']}'")

    now = datetime.now(timezone.utc).isoformat()
    session["progress"][agent_id] = {
        "pct": progress_pct,
        "message": status_message,
        "updated_at": now,
        "epoch": time.time(),
    }
    session["updated_at"] = now
    _write_session(session)

    _log_event(session_id, "progress_reported", {
        "agent_id": agent_id,
        "progress_pct": progress_pct,
        "status_message": status_message,
    })

    return session


def check_blocked(session_id: str) -> dict:
    """Detect if swarm is stuck.

    Returns {"blocked": bool, "stuck_agents": [...], "reason": str}.
    An agent is stuck if they haven't reported progress within the
    session's block_timeout (default 30 minutes).
    """
    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    timeout = session.get("block_timeout", DEFAULT_BLOCK_TIMEOUT)
    now_epoch = time.time()
    stuck_agents = []

    for agent_id in session["agents"]:
        prog = session.get("progress", {}).get(agent_id)
        if prog is None:
            # Agent never reported progress -- check if session has been
            # active long enough (use session updated_at as baseline)
            stuck_agents.append(agent_id)
        else:
            last_update = prog.get("epoch", 0)
            if (now_epoch - last_update) > timeout:
                stuck_agents.append(agent_id)

    blocked = len(stuck_agents) > 0 and session["state"] not in TERMINAL_STATES

    reason = ""
    if blocked:
        reason = f"{len(stuck_agents)} agent(s) have not reported progress within {timeout}s"

    return {
        "blocked": blocked,
        "stuck_agents": stuck_agents,
        "reason": reason,
        "timeout_seconds": timeout,
    }


def auto_reassign(session_id: str) -> dict:
    """If an agent is stuck, reassign their work to another agent.

    Returns {"reassigned": [...]} with details of reassignments.
    """
    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    if session["state"] in TERMINAL_STATES:
        return {"reassigned": [], "reason": "Session is in terminal state"}

    check = check_blocked(session_id)
    if not check["blocked"]:
        return {"reassigned": [], "reason": "No agents are stuck"}

    stuck = check["stuck_agents"]
    # Find non-stuck agents that can take over
    available = [
        aid for aid in session["agents"]
        if aid not in stuck and session["agents"][aid]["role"] == "worker"
    ]

    reassigned = []
    now = datetime.now(timezone.utc).isoformat()

    for stuck_agent in stuck:
        if available:
            target = available[0]  # assign to first available
            reassigned.append({
                "from_agent": stuck_agent,
                "to_agent": target,
                "reassigned_at": now,
            })
            _log_event(session_id, "auto_reassigned", {
                "from_agent": stuck_agent,
                "to_agent": target,
            })
        else:
            reassigned.append({
                "from_agent": stuck_agent,
                "to_agent": None,
                "reassigned_at": now,
                "note": "No available agent for reassignment",
            })
            _log_event(session_id, "reassign_failed", {
                "from_agent": stuck_agent,
                "reason": "No available agent",
            })

    session["updated_at"] = now
    _write_session(session)

    return {"reassigned": reassigned}


def get_session_history(session_id: str) -> List[dict]:
    """Full timeline of all events for a session."""
    all_events = _read_jsonl_sync(HISTORY_PATH)
    return [e for e in all_events if e.get("session_id") == session_id]


def list_sessions(state: Optional[str] = None, limit: int = 20) -> List[dict]:
    """List swarm sessions, optionally filtered by state."""
    if state is not None and state not in STATES:
        raise ValueError(f"Invalid state '{state}'. Must be one of: {STATES}")

    sessions = []
    for f in SESSIONS_DIR.glob("*.json"):
        try:
            session = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        if state and session.get("state") != state:
            continue

        sessions.append({
            "session_id": session["session_id"],
            "name": session["name"],
            "state": session["state"],
            "agent_count": len(session.get("agents", {})),
            "created_at": session["created_at"],
            "updated_at": session["updated_at"],
        })

    sessions.sort(key=lambda x: x["created_at"], reverse=True)
    return sessions[:limit]


def dissolve(session_id: str, reason: str = "") -> dict:
    """End a swarm session. Returns updated session."""
    session = _read_session(session_id)
    if session is None:
        raise FileNotFoundError(f"Session '{session_id}' not found")

    if session["state"] in TERMINAL_STATES:
        raise ValueError(f"Session is already in terminal state '{session['state']}'")

    now = datetime.now(timezone.utc).isoformat()
    old_state = session["state"]
    session["state"] = "FAILED"
    session["dissolved_at"] = now
    session["dissolve_reason"] = reason
    session["updated_at"] = now
    _write_session(session)

    _log_event(session_id, "dissolved", {
        "from_state": old_state,
        "reason": reason,
    })

    return session


def get_swarm_stats() -> dict:
    """Return swarm statistics: active/completed/failed sessions, averages."""
    by_state: Dict[str, int] = {}
    total = 0
    total_agents = 0
    completed_count = 0
    completed_times: List[float] = []

    for f in SESSIONS_DIR.glob("*.json"):
        try:
            session = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        total += 1
        state = session.get("state", "UNKNOWN")
        by_state[state] = by_state.get(state, 0) + 1
        total_agents += len(session.get("agents", {}))

        if state == "COMPLETED":
            completed_count += 1
            # Try to compute completion time
            created = session.get("created_at", "")
            updated = session.get("updated_at", "")
            if created and updated:
                try:
                    from datetime import datetime as dt
                    t_start = dt.fromisoformat(created)
                    t_end = dt.fromisoformat(updated)
                    diff = (t_end - t_start).total_seconds()
                    completed_times.append(diff)
                except (ValueError, TypeError):
                    pass

    active = sum(
        c for s, c in by_state.items() if s not in TERMINAL_STATES
    )

    avg_agents = (total_agents / total) if total > 0 else 0.0
    avg_completion_time = (
        sum(completed_times) / len(completed_times) if completed_times else 0.0
    )

    return {
        "total_sessions": total,
        "active_sessions": active,
        "by_state": by_state,
        "avg_agents_per_swarm": round(avg_agents, 1),
        "avg_completion_time_seconds": round(avg_completion_time, 1),
    }
