"""
Phase 96: Full Agent Lifecycle Management — Birth to Retirement.

Manages every agent through a well-defined state machine: onboarding, active,
idle, suspended, and retired. Tracks all transitions with timestamps and
reasons, enabling full timeline reconstruction and fleet-wide analytics.

Storage:
  state/broker/lifecycle/
    agents/{agent_id}.json      — current agent state + metadata
    transitions.jsonl           — append-only log of all transitions
"""

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

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

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

LIFECYCLE_DIR = BROKER_DIR / "lifecycle"
AGENTS_DIR = LIFECYCLE_DIR / "agents"
TRANSITIONS_FILE = LIFECYCLE_DIR / "transitions.jsonl"

for _d in (LIFECYCLE_DIR, AGENTS_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# ── Valid states & transitions ─────────────────────────────────────────────

VALID_STATES = ("onboarding", "active", "idle", "suspended", "retired")

ALLOWED_TRANSITIONS: Dict[str, tuple] = {
    "onboarding": ("active", "suspended", "retired"),
    "active":     ("idle", "suspended", "retired"),
    "idle":       ("active", "suspended", "retired"),
    "suspended":  ("active", "retired"),
    "retired":    (),  # terminal state
}


# ── Helpers ────────────────────────────────────────────────────────────────

def _agent_path(agent_id: str) -> Path:
    return AGENTS_DIR / f"{agent_id}.json"


def _load_agent(agent_id: str) -> Optional[dict]:
    p = _agent_path(agent_id)
    if not p.exists():
        return None
    try:
        return json.loads(p.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_agent(agent: dict) -> None:
    _write_json_sync(_agent_path(agent["agent_id"]), agent)


def _record_transition(agent_id: str, from_state: str, to_state: str,
                       reason: str = "") -> dict:
    entry = {
        "transition_id": uuid.uuid4().hex[:12],
        "agent_id": agent_id,
        "from_state": from_state,
        "to_state": to_state,
        "reason": reason,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "epoch": time.time(),
    }
    _append_jsonl_sync(TRANSITIONS_FILE, entry)
    return entry


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

def register_agent(agent_id: str, name: str, platform: str,
                   capabilities: Optional[List[str]] = None) -> dict:
    """Register a new agent; initial state is 'onboarding'."""
    if _load_agent(agent_id) is not None:
        raise ValueError(f"Agent '{agent_id}' already registered")
    now = datetime.now(timezone.utc).isoformat()
    agent = {
        "agent_id": agent_id,
        "name": name,
        "platform": platform,
        "capabilities": capabilities or [],
        "state": "onboarding",
        "registered_at": now,
        "last_activity": now,
        "state_changed_at": now,
        "retirement_reason": None,
    }
    _save_agent(agent)
    _record_transition(agent_id, "", "onboarding", "initial registration")
    return agent


def get_lifecycle_state(agent_id: str) -> str:
    """Return current lifecycle state for an agent."""
    agent = _load_agent(agent_id)
    if agent is None:
        raise KeyError(f"Agent '{agent_id}' not found")
    return agent["state"]


def transition(agent_id: str, new_state: str, reason: str = "") -> dict:
    """Transition an agent to a new state with validation."""
    if new_state not in VALID_STATES:
        raise ValueError(f"Invalid state '{new_state}'. Must be one of {VALID_STATES}")
    agent = _load_agent(agent_id)
    if agent is None:
        raise KeyError(f"Agent '{agent_id}' not found")
    current = agent["state"]
    if new_state == current:
        raise ValueError(f"Agent already in state '{current}'")
    allowed = ALLOWED_TRANSITIONS.get(current, ())
    if new_state not in allowed:
        raise ValueError(
            f"Cannot transition from '{current}' to '{new_state}'. "
            f"Allowed: {allowed}"
        )
    old_state = current
    now = datetime.now(timezone.utc).isoformat()
    agent["state"] = new_state
    agent["state_changed_at"] = now
    agent["last_activity"] = now
    _save_agent(agent)
    _record_transition(agent_id, old_state, new_state, reason)
    return agent


def retire_agent(agent_id: str, reason: str = "") -> dict:
    """Gracefully retire an agent: archive data, record reason."""
    agent = _load_agent(agent_id)
    if agent is None:
        raise KeyError(f"Agent '{agent_id}' not found")
    if agent["state"] == "retired":
        raise ValueError(f"Agent '{agent_id}' is already retired")
    old_state = agent["state"]
    now = datetime.now(timezone.utc).isoformat()
    agent["state"] = "retired"
    agent["state_changed_at"] = now
    agent["last_activity"] = now
    agent["retirement_reason"] = reason or "graceful retirement"
    _save_agent(agent)
    _record_transition(agent_id, old_state, "retired", reason or "graceful retirement")
    return agent


def get_agent_timeline(agent_id: str) -> List[dict]:
    """Return full transition history for an agent, chronological."""
    all_transitions = _read_jsonl_sync(TRANSITIONS_FILE)
    return [t for t in all_transitions if t.get("agent_id") == agent_id]


def get_agents_by_state(state: str) -> List[dict]:
    """List all agents currently in a specific lifecycle state."""
    if state not in VALID_STATES:
        raise ValueError(f"Invalid state '{state}'. Must be one of {VALID_STATES}")
    results = []
    if not AGENTS_DIR.exists():
        return results
    for f in AGENTS_DIR.iterdir():
        if not f.name.endswith(".json"):
            continue
        try:
            agent = json.loads(f.read_text("utf-8"))
            if agent.get("state") == state:
                results.append(agent)
        except (json.JSONDecodeError, OSError):
            continue
    return results


def get_lifecycle_stats() -> dict:
    """Counts by state, avg time in each state, retirement rate."""
    counts: Dict[str, int] = {s: 0 for s in VALID_STATES}
    total = 0
    retired = 0
    if AGENTS_DIR.exists():
        for f in AGENTS_DIR.iterdir():
            if not f.name.endswith(".json"):
                continue
            try:
                agent = json.loads(f.read_text("utf-8"))
                st = agent.get("state", "unknown")
                if st in counts:
                    counts[st] += 1
                total += 1
                if st == "retired":
                    retired += 1
            except (json.JSONDecodeError, OSError):
                continue

    # Compute average time spent per state from transitions
    transitions = _read_jsonl_sync(TRANSITIONS_FILE)
    state_durations: Dict[str, List[float]] = {s: [] for s in VALID_STATES}
    # Group transitions by agent, ordered by epoch
    agent_trans: Dict[str, List[dict]] = {}
    for t in transitions:
        aid = t.get("agent_id", "")
        agent_trans.setdefault(aid, []).append(t)
    for aid, tlist in agent_trans.items():
        tlist.sort(key=lambda x: x.get("epoch", 0))
        for i in range(len(tlist) - 1):
            from_s = tlist[i].get("to_state", "")
            dur = tlist[i + 1].get("epoch", 0) - tlist[i].get("epoch", 0)
            if from_s in state_durations and dur >= 0:
                state_durations[from_s].append(dur)

    avg_duration = {}
    for s, durs in state_durations.items():
        avg_duration[s] = round(sum(durs) / len(durs), 2) if durs else 0

    return {
        "total_agents": total,
        "counts_by_state": counts,
        "retirement_rate": round(retired / total, 4) if total else 0,
        "avg_seconds_in_state": avg_duration,
    }


def auto_idle_detection(threshold_hours: int = 24) -> List[dict]:
    """Find active agents with no recent activity, transition to idle."""
    cutoff = (datetime.now(timezone.utc) - timedelta(hours=threshold_hours)).isoformat()
    idled = []
    if not AGENTS_DIR.exists():
        return idled
    for f in AGENTS_DIR.iterdir():
        if not f.name.endswith(".json"):
            continue
        try:
            agent = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if agent.get("state") == "active" and agent.get("last_activity", "") < cutoff:
            updated = transition(agent["agent_id"], "idle", "auto-idle detection")
            idled.append(updated)
    return idled
