"""The reconcile loop: desired agents -> placement -> worker jobs.

`reconcile(store, jobs, now)` is one pass and returns the actions it took; the runner
calls it every few seconds in a thread (Hub calls block). A worker row is a stable *slot*
(`<swarm>-<type>-wNNNN`); its job is replaced on restart, its agents never move.
"""

from __future__ import annotations

import asyncio
import logging
import os
import random
import time
import uuid
from collections import Counter
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta

from flock.controller import spawn as spawning
from flock.controller.jobs import JobsClient, JobState
from flock.controller.prices import cost_usd
from flock.controller.store import Row, Store
from flock.models import (
    CRASHLOOP_THRESHOLD,
    DEAD_AFTER,
    LAUNCHES_PER_WINDOW,
    RATE_WINDOW,
    UNKNOWN_AFTER,
    AgentState,
    AssignCommand,
    DrainCommand,
    Event,
    SwarmSpec,
    TypeHealth,
    TypeSpec,
    UnassignCommand,
    WorkerHeartbeat,
    WorkerState,
    agent_name,
    max_workers,
    parse_duration,
)

log = logging.getLogger("flock.reconcile")

BACKOFF_BASE = 30.0  # s, first crash backoff
BACKOFF_CAP = 900.0  # s, 15 min
INSPECT_EVERY = 15.0  # s, min gap between job-status checks for one non-fresh worker
INSPECT_GRACE = 60.0  # s after launch before a worker with no heartbeat is checked at all
LIST_THRESHOLD = 10  # more non-fresh workers than this in a swarm: one list_jobs, not N inspects
LAUNCH_TIMEOUT = 900.0  # s, a RUNNING job that never heartbeats in this long is a crash
RETIRE_GRACE = 120.0  # s, an emptied worker gets this long to drain before it is cancelled
STABLE_AFTER = 600.0  # s, a heartbeat this long after launch resets the crash streak
_UNKNOWN = UNKNOWN_AFTER.total_seconds()
_DEAD = DEAD_AFTER.total_seconds()


@dataclass
class Settings:
    public_url: str = "http://localhost:7860"
    token: str = ""
    secret_source: Mapping[str, str] = field(default_factory=lambda: os.environ)
    launches_per_window: int = LAUNCHES_PER_WINDOW
    rate_window: timedelta = RATE_WINDOW
    rng: random.Random = field(default_factory=random.Random)
    job_labels: dict[str, str] = field(default_factory=dict)  # deployment-wide extra job labels


@dataclass
class Action:
    kind: str  # launch, planned_restart, crash, crashloop, cancel, assign, unassign, retire, ...
    target: str
    detail: str = ""


def worker_id_for(swarm: str, type_name: str, idx: int) -> str:
    return f"{swarm}-{type_name}-w{idx:04d}"


def _iso(ts: float) -> str:
    return datetime.fromtimestamp(ts, UTC).isoformat()


# --------------------------------------------------------------------------- derived states


def worker_state(w: Row, now: float) -> WorkerState:
    if w["job_id"] is None:
        return WorkerState.dead
    hb = w["last_heartbeat"]
    if hb is None:
        return WorkerState.launching
    age = now - hb
    if age > _DEAD:
        return WorkerState.dead
    if age > _UNKNOWN:
        return WorkerState.unknown
    return WorkerState.draining if w["draining"] or w["retire_at"] else WorkerState.running


def agent_state(a: Row, w: Row | None, now: float) -> AgentState:
    hb = a["last_heartbeat"]
    if w is None or w["job_id"] is None:
        return AgentState.dead if hb is not None else AgentState.unknown
    if worker_state(w, now) == WorkerState.dead:
        return AgentState.dead
    if hb is None or (w["launched_at"] and hb < w["launched_at"]):
        return AgentState.unknown  # not yet reported by the current job
    age = now - hb
    if age > _DEAD:
        return AgentState.dead
    if age > _UNKNOWN or not a["hb_state"]:
        return AgentState.unknown
    return AgentState(a["hb_state"])


# --------------------------------------------------------------------------- one pass


@dataclass
class _Pass:
    store: Store
    jobs: JobsClient
    now: float
    settings: Settings
    budget: int
    actions: list[Action] = field(default_factory=list)
    listed: dict[str, JobState] | None = None  # job stages from one list_jobs, if used

    def act(self, kind: str, target: str, detail: str = "") -> None:
        self.actions.append(Action(kind, target, detail))

    def job_state(self, job_id: str, namespace: str) -> JobState | None:
        if self.listed is not None:
            return self.listed.get(job_id)
        try:
            return self.jobs.inspect(job_id, namespace=namespace)
        except Exception as exc:  # noqa: BLE001 - Hub hiccup: try again next pass
            log.warning("inspect %s failed: %s", job_id, type(exc).__name__)
            return None

    def cancel(self, job_id: str, namespace: str) -> None:
        try:
            self.jobs.cancel(job_id, namespace=namespace)
        except Exception as exc:  # noqa: BLE001 - the job may already be gone
            log.warning("cancel %s failed: %s", job_id, type(exc).__name__)


def _needs_check(w: Row, now: float) -> bool:
    """Fresh, non-draining workers are trusted from heartbeats; others get a job-status check."""
    if w["job_id"] is None:
        return False
    hb = w["last_heartbeat"]
    if hb is None and now - (w["launched_at"] or now) < INSPECT_GRACE:
        return False
    if hb is not None and now - hb <= _UNKNOWN and not w["draining"] and not w["retire_at"]:
        return False
    return w["last_inspect"] is None or now - w["last_inspect"] >= INSPECT_EVERY


def reconcile(
    store: Store, jobs: JobsClient, now: datetime, settings: Settings | None = None
) -> list[Action]:
    settings = settings or Settings()
    if settings.token and store.agent_token_key is None:
        spawning.set_token_key(store, settings.token)
    t = now.timestamp()
    used = store.launches_since(t - settings.rate_window.total_seconds())
    p = _Pass(store, jobs, t, settings, budget=settings.launches_per_window - used)
    for sw in store.list_swarms():
        spec: SwarmSpec = sw["spec"]
        p.listed = None
        if sum(_needs_check(w, t) for w in store.list_workers(spec.swarm)) > LIST_THRESHOLD:
            try:
                listing = jobs.list({"swarm": spec.swarm}, namespace=spec.namespace)
                p.listed = {j.job_id: JobState(j.stage, j.message) for j in listing}
            except Exception as exc:  # noqa: BLE001
                log.warning("list_jobs for %s failed: %s", spec.swarm, type(exc).__name__)
                p.listed = {}
        for type_name in sorted(set(spec.types) | set(store.type_names(spec.swarm))):
            ts = None if sw["removing"] else spec.types.get(type_name)
            _reconcile_type(p, spec, type_name, ts)
        gone = not store.list_workers(spec.swarm) and not store.list_agents(spec.swarm)
        if sw["removing"] and gone:
            store.delete_swarm(spec.swarm)
            p.act("remove_swarm", spec.swarm)
    return p.actions


def _reconcile_type(p: _Pass, spec: SwarmSpec, type_name: str, ts: TypeSpec | None) -> None:
    store = p.store
    for w in store.list_workers(spec.swarm, type_name):
        _refresh_worker(p, spec, w, ts)

    # Desired agents; scale down removes the highest slots first. Spawned agents (parent set)
    # are not replicas: they go with their parent, their type or the swarm (docs/spawn.md).
    replicas = ts.replicas if ts else 0
    existing = {a["name"]: a for a in store.list_agents(spec.swarm, type_name)}
    reason = "type_removed" if type_name not in spec.types else "swarm_removed"
    for a in sorted(existing.values(), key=lambda a: a["slot"], reverse=True):
        if ts is None or a["removing"] or (a["parent"] is None and a["slot"] >= replicas):
            for r in spawning.remove_subtree(store, a["name"], reason, p.now):
                p.act("unassign", r["name"], r["worker"] or "")
    for slot in range(replicas):
        name = agent_name(type_name, slot)
        if name not in existing:
            store.insert_agent(name, spec.swarm, type_name, slot, p.now)

    # Placement: fill live slots in order, never move a placed agent, open new slots as needed.
    if ts is not None:
        _place(p, spec, type_name, ts)

    # Empty workers drain and are cancelled after a grace period.
    counts = Counter(a["worker"] for a in store.list_agents(spec.swarm, type_name))
    for w in store.list_workers(spec.swarm, type_name):
        if counts[w["id"]]:
            continue
        if w["job_id"] is None:
            store.delete_worker(w["id"])
            p.act("remove_worker", w["id"])
        elif w["retire_at"] is None:
            store.enqueue_command(w["id"], DrainCommand().model_dump_json())
            store.update_worker(w["id"], retire_at=p.now + RETIRE_GRACE)
            p.act("retire", w["id"])
        elif p.now >= w["retire_at"]:
            p.cancel(w["job_id"], spec.namespace)
            store.delete_worker(w["id"])
            p.act("cancel", w["id"], w["job_id"])

    # Launch slots that have agents and no job.
    if ts is None or store.type_health(spec.swarm, type_name) == TypeHealth.crashlooping:
        return
    for w in store.list_workers(spec.swarm, type_name):
        if w["job_id"] is not None or not counts[w["id"]]:
            continue
        if w["next_launch_at"] is not None and w["next_launch_at"] > p.now:
            continue
        if p.budget <= 0:
            p.act("budget_exhausted", w["id"])
            return
        _launch(p, spec, type_name, ts, w)


def _place(p: _Pass, spec: SwarmSpec, type_name: str, ts: TypeSpec) -> None:
    store = p.store
    workers = [w for w in store.list_workers(spec.swarm, type_name) if w["retire_at"] is None]
    by_id = {w["id"]: w for w in workers}
    agents = store.list_agents(spec.swarm, type_name)
    counts = Counter(a["worker"] for a in agents if a["worker"] in by_id)
    unplaced = [a for a in agents if a["worker"] not in by_id]
    if not unplaced:
        return
    cap = max(1, ts.agents_per_worker)
    queue = iter(unplaced)
    pending = len(unplaced)

    def fill(w: Row) -> None:
        nonlocal pending
        while counts[w["id"]] < cap and pending:
            a = next(queue)
            pending -= 1
            counts[w["id"]] += 1
            store.update_agent(a["name"], worker=w["id"])
            if w["job_id"]:
                _send_assign(store, w["id"], a, type_name, ts)
            p.act("assign", a["name"], w["id"])

    for w in workers:
        fill(w)
    # A type that declares `workers` never has more than initial + max_extra worker slots
    # (retiring ones included); agents that do not fit wait unplaced. Others: v1, unbounded.
    slots = max_workers(ts) if ts.workers is not None else None
    while pending:
        if slots is not None and len(store.list_workers(spec.swarm, type_name)) >= slots:
            p.act("workers_full", f"{spec.swarm}/{type_name}", f"{pending} unplaced")
            return
        idx = store.next_worker_index(spec.swarm, type_name)
        wid = worker_id_for(spec.swarm, type_name, idx)
        store.insert_worker(wid, spec.swarm, type_name, idx, p.now)
        fill(store.get_worker(wid))


def _send_assign(store: Store, worker_id: str, a: Row, type_name: str, ts: TypeSpec) -> None:
    if a["removing"]:
        return  # mid-removal (docs/spawn.md): never re-assign it or re-issue its token
    resume = ts.resume_prompt if a["resume_pending"] else None
    cmd = AssignCommand(
        agent=a["name"], type=type_name, spec=ts, restore=True, resume_prompt=resume,
        parent=a["parent"], agent_token=spawning.agent_token(store, a["name"], a["swarm"], time.time()),
    )
    store.enqueue_command(worker_id, cmd.model_dump_json())
    if a["resume_pending"]:
        store.update_agent(a["name"], resume_pending=0)


def _refresh_worker(p: _Pass, spec: SwarmSpec, w: Row, ts: TypeSpec | None) -> None:
    if not _needs_check(w, p.now):
        if w["job_id"] is not None:
            p.store.update_worker(w["id"], state=worker_state(w, p.now).value)
        return
    p.store.update_worker(w["id"], last_inspect=p.now)
    st = p.job_state(w["job_id"], spec.namespace)
    if st is None:
        return
    if st.terminal:
        _job_ended(p, spec, w, ts, st)
        return
    hb = w["last_heartbeat"]
    silent = hb is not None and p.now - hb > _DEAD
    never = hb is None and st.stage != "SCHEDULING" and p.now - w["launched_at"] > LAUNCH_TIMEOUT
    if silent or never:
        p.cancel(w["job_id"], spec.namespace)
        p.act("cancel", w["id"], w["job_id"])
        _job_ended(p, spec, w, ts, JobState("CANCELED", "unresponsive"))
        return
    p.store.update_worker(w["id"], state=worker_state(w, p.now).value)


def _job_ended(p: _Pass, spec: SwarmSpec, w: Row, ts: TypeSpec | None, st: JobState) -> None:
    store, wid = p.store, w["id"]
    store.clear_commands(wid)
    if w["retire_at"] is not None:
        store.delete_worker(wid)
        p.act("remove_worker", wid, st.stage)
        return
    drain = parse_duration(ts.drain_window).total_seconds() if ts else 0.0
    near_deadline = w["deadline"] is not None and p.now >= w["deadline"] - drain
    planned = (st.stage == "ERROR" and st.message == "Job timeout") or (
        st.stage == "COMPLETED" and (bool(w["draining"]) or near_deadline)
    )
    common = {
        "job_id": None,
        "incarnation": None,
        "state": WorkerState.dead.value,
        "last_heartbeat": None,
        "draining": 0,
    }
    if planned:
        store.update_worker(wid, **common, consecutive_crashes=0, next_launch_at=p.now)
        p.act("planned_restart", wid, f"{st.stage}: {st.message}")
        return
    streak = w["consecutive_crashes"] + 1
    backoff = min(BACKOFF_BASE * 2 ** (streak - 1), BACKOFF_CAP)
    store.update_worker(
        wid,
        **common,
        crashes=w["crashes"] + 1,
        consecutive_crashes=streak,
        next_launch_at=p.now + backoff,
    )
    p.act("crash", wid, f"{st.stage}: {st.message}; backoff {backoff:.0f}s")
    for a in store.list_agents(worker=wid):
        if a["turn_id"]:
            ev = Event(
                type="turn_interrupted",
                ts=datetime.fromtimestamp(p.now, UTC),
                agent=a["name"],
                turn_id=a["turn_id"],
                data={"worker": wid, "job_id": w["job_id"], "reason": f"{st.stage}: {st.message}"},
            )
            store.append_events(a["swarm"], a["name"], [ev])
            store.update_agent(a["name"], resume_pending=1, turn_id=None)
    if streak >= CRASHLOOP_THRESHOLD:
        store.set_type_health(w["swarm"], w["type"], TypeHealth.crashlooping.value)
        p.act("crashloop", f"{w['swarm']}/{w['type']}")


def _launch(p: _Pass, spec: SwarmSpec, type_name: str, ts: TypeSpec, w: Row) -> None:
    store, s, wid = p.store, p.settings, w["id"]
    lifetime = parse_duration(ts.worker_lifetime).total_seconds()
    jitter = parse_duration(ts.lifetime_jitter).total_seconds()
    timeout = lifetime + s.rng.uniform(0, jitter)
    deadline = p.now + timeout
    # Binds polls and heartbeats to this job: a parked poll or a late heartbeat from the job
    # this one replaces must not take the new job's commands (e2e run 1, see the API).
    incarnation = uuid.uuid4().hex
    env = {
        **ts.env,
        "FLOCK_URL": s.public_url,
        "WORKER_ID": wid,
        "WORKER_INCARNATION": incarnation,
        "SWARM": spec.swarm,
        "TYPE": type_name,
        "DEADLINE": _iso(deadline),
        "DRAIN_WINDOW": ts.drain_window,  # known before the first assign arrives
        "BUCKET": spec.bucket,
    }
    secrets = {"FLOCK_TOKEN": s.token}
    for name in [*ts.secrets, *([ts.msgboard_invite] if ts.msgboard_invite else [])]:
        if name in s.secret_source:
            secrets[name] = s.secret_source[name]
        else:
            log.warning(
                "secret %s for %s/%s is not set on the controller", name, spec.swarm, type_name
            )
    labels = {**s.job_labels, "swarm": spec.swarm, "type": type_name, "worker": wid}
    p.budget -= 1
    store.record_launch(p.now)
    try:
        job_id = p.jobs.launch(wid, ts, env, secrets, labels, timeout, namespace=spec.namespace)
    except Exception as exc:  # noqa: BLE001 - any Hub failure: retry later, keep the loop alive
        store.update_worker(wid, next_launch_at=p.now + BACKOFF_BASE)
        p.act("launch_failed", wid, type(exc).__name__)
        log.warning("launch of %s failed: %s", wid, type(exc).__name__)
        return
    store.clear_commands(wid)
    store.update_worker(
        wid,
        job_id=job_id,
        incarnation=incarnation,
        launched_at=p.now,
        deadline=deadline,
        last_heartbeat=None,
        last_inspect=None,
        draining=0,
        state=WorkerState.launching.value,
        next_launch_at=None,
    )
    for a in store.list_agents(worker=wid):
        _send_assign(store, wid, a, type_name, ts)
    p.act("launch", wid, job_id)


# --------------------------------------------------------------------------- heartbeats


def record_heartbeat(store: Store, worker_id: str, hb: WorkerHeartbeat, now: datetime) -> bool:
    """Apply a worker heartbeat. Returns False if the worker slot has no live job."""
    t = now.timestamp()
    w = store.get_worker(worker_id)
    if w is None or w["job_id"] is None:
        return False
    fields = {
        "last_heartbeat": t,
        "deadline": hb.deadline.timestamp(),
        "draining": int(hb.draining),
        "state": (WorkerState.draining if hb.draining else WorkerState.running).value,
    }
    if w["consecutive_crashes"] and t - (w["launched_at"] or t) >= STABLE_AFTER:
        fields["consecutive_crashes"] = 0
    store.update_worker(worker_id, **fields)
    queued = {c["body"] for c in store.commands_for(worker_id)}
    for ah in hb.agents:
        a = store.get_agent(ah.name)
        if (
            a is not None
            and a["worker"] is None
            and (a["swarm"], a["type"])
            == (
                w["swarm"],
                w["type"],
            )
        ):
            store.update_agent(ah.name, worker=worker_id)  # adopt (controller lost its state)
        elif a is None or a["worker"] != worker_id:
            body = UnassignCommand(agent=ah.name).model_dump_json()
            if body not in queued:
                store.enqueue_command(worker_id, body)
            continue
        store.record_agent_heartbeat(
            ah.name,
            {
                "hb_state": ah.state,
                "turn_id": ah.turn_id,
                "idle_since": ah.idle_since.timestamp() if ah.idle_since else None,
                "session_id": ah.session_id,
                "resumed": None if ah.resumed is None else int(ah.resumed),
                "harness_estimate_usd": ah.harness_estimate_usd,
                "turns": ah.turns,
                "last_sync": ah.last_sync.timestamp() if ah.last_sync else None,
                "msgboard_principal": ah.msgboard_principal,
                "last_heartbeat": t,
            },
            ah.tokens,
            cost_usd,
        )
    if not hb.draining and not w["retire_at"]:
        _reassign_missing(store, w, {ah.name for ah in hb.agents}, queued)
    return True


def _reassign_missing(store: Store, w: Row, reported: set[str], queued: set[str]) -> None:
    """Agents placed on this worker that its heartbeat does not list, with no assign queued,
    lost their assign on the way (a dropped poll response). Send it again; a worker treats a
    repeated assign for an agent it already hosts as a spec update, so this is idempotent."""
    missing = [a for a in store.list_agents(worker=w["id"])
               if a["name"] not in reported and not a["removing"]]
    if not missing:
        return
    sw = store.get_swarm(w["swarm"])
    ts = sw["spec"].types.get(w["type"]) if sw else None
    if ts is None:
        return
    pending = {AssignCommand.model_validate_json(b).agent for b in queued if '"assign"' in b}
    for a in missing:
        if a["name"] not in pending:
            log.warning("re-sending assign for %s: missing from %s heartbeat", a["name"], w["id"])
            _send_assign(store, w["id"], a, w["type"], ts)


# --------------------------------------------------------------------------- startup and runner


def rebuild_from_jobs(store: Store, jobs: JobsClient, now: datetime) -> list[Action]:
    """After a controller restart, re-adopt live worker jobs by label, then wait for heartbeats."""
    t = now.timestamp()
    actions: list[Action] = []
    for sw in store.list_swarms():
        spec: SwarmSpec = sw["spec"]
        live: dict[str, list] = {}
        for j in jobs.list({"swarm": spec.swarm}, namespace=spec.namespace):
            wid = j.labels.get("worker") or j.env.get("WORKER_ID")
            type_name = j.labels.get("type") or j.env.get("TYPE")
            if wid and type_name and j.stage not in ("COMPLETED", "CANCELED", "ERROR", "DELETED"):
                live.setdefault(wid, []).append((j, type_name))
        for wid, found in live.items():
            w = store.get_worker(wid)
            keep = next((x for x in found if w and x[0].job_id == w["job_id"]), found[-1])
            for j, _ in found:
                if j is not keep[0]:
                    jobs.cancel(j.job_id, namespace=spec.namespace)
                    actions.append(Action("cancel", wid, j.job_id))
            j, type_name = keep
            if w is None:
                idx = int(wid.rsplit("-w", 1)[-1]) if wid.rsplit("-w", 1)[-1].isdigit() else 0
                store.insert_worker(wid, spec.swarm, type_name, idx, t)
                store.bump_worker_index(spec.swarm, type_name, idx + 1)
                w = store.get_worker(wid)
            deadline = None
            if j.env.get("DEADLINE"):
                deadline = datetime.fromisoformat(j.env["DEADLINE"]).timestamp()
            store.update_worker(
                wid,
                job_id=j.job_id,
                incarnation=j.env.get("WORKER_INCARNATION"),
                deadline=deadline or w["deadline"],
                launched_at=w["launched_at"] or t,
                # one heartbeat interval of grace before the age rules apply
                last_heartbeat=t if w["last_heartbeat"] is not None else None,
            )
            actions.append(Action("adopt", wid, j.job_id))
    return actions


async def run_reconciler(
    store: Store, jobs: JobsClient, settings: Settings, interval: float = 5.0, wake=None
) -> None:
    """Background task: rebuild once, then reconcile every `interval` s (or when woken)."""
    now = lambda: datetime.now(UTC)
    try:
        await asyncio.to_thread(rebuild_from_jobs, store, jobs, now())
    except Exception:
        log.exception("rebuild from jobs failed")
    while True:
        try:
            for a in await asyncio.to_thread(reconcile, store, jobs, now(), settings):
                if a.kind not in ("assign", "budget_exhausted"):
                    log.info("%s %s %s", a.kind, a.target, a.detail)
        except Exception:
            log.exception("reconcile pass failed")
        if wake is not None:
            await wake(interval)
        else:
            await asyncio.sleep(interval)