// A model of nisse-dev's orchestrator (nisse-dev/src/nisse_dev/control.py),
// written to look for counterexamples: interleavings of operator messages,
// Gitea comments, work outcomes and crashes that break a property the
// design relies on.
//
// Everything lives in one record, `s`, and each handler is a pure function
// named after the method it mirrors. The model keeps what matters for the
// properties below and drops the rest (texts, transcripts, models, tools).
//
// Durable state (the state store) survives a crash; volatile state (the
// worker queue, work in flight, pollers, the cancel flag) does not. Gitea is
// the outside world and survives everything.
//
// Run:  quint run specs/orchestrator.qnt --invariant=all_properties
//       quint run specs/orchestrator.qnt --invariant=<one property>
module orchestrator {

  // ---- bounds --------------------------------------------------------------

  pure val MAX_MS = 2        // milestones per task
  pure val MAX_TASKS = 2     // tasks per run ('use' / 'abandon' / 'complete')
  pure val MAX_QUEUE = 3     // jobs waiting on the worker
  pure val MAX_IDS = 8       // Gitea comments per run

  // ---- vocabulary ----------------------------------------------------------

  pure val WORKING = Set("PLANNING_ARCH", "MILESTONE_PLANNING", "EXECUTING")
  pure val SIGNOFF = Set("AWAIT_ARCH_SIGNOFF", "AWAIT_MILESTONE_SIGNOFF")
  // Work that sets _working, so a message then gets "busy" (control._handle).
  pure val LONG_WORK = Set("plan", "exec", "review")

  // What the operator had in front of them: the plan last put up for sign-off.
  type Plan = { g: str, rev: int }
  pure val NO_PLAN: Plan = { g: "", rev: 0 }

  type Gate = { on: bool, g: str, rev: int, approved: bool }
  pure val NO_GATE: Gate = { on: false, g: "", rev: 0, approved: false }

  type Pending = { kind: str, g: str, interrupted: str }
  pure val NO_PENDING: Pending = { kind: "", g: "", interrupted: "" }

  // A message or comment waiting on the worker. `saw` is what the operator
  // had seen when they sent it; the agent never knows it, the properties do.
  type Job = { src: str, kind: str, saw: Plan, cid: int }

  type PlanComment = { id: int, tid: int, kind: str, saw: Plan }
  type PrComment = { id: int, key: int }

  type Watch = { on: bool, floor: int }
  pure val OFF: Watch = { on: false, floor: 0 }

  type Sys = {
    // -- durable: the state store (control._data and the state) --
    st: str,
    tid: int,                 // task counter; a task's branches are its own
    task: bool,               // a project is in use
    gate: Gate,
    revs: str -> int,         // latest revision per gate
    lastRev: int,             // revisions run across the task (L3)
    ms: int,
    pending: Pending,
    prOn: bool, prKey: int, prLast: int,          // _data["pr"]
    reviewing: int,           // pr["reviewing"]: the comment being answered, 0 if none
    boardOn: bool, boardPosted: bool, boardLast: int,  // _data["board"]
    intents: Set[int],        // pending open_pr intents, by PR key
    // -- Gitea --
    planThread: List[PlanComment],
    prThread: List[PrComment],
    nextId: int,
    giteaPr: int -> str,      // PR key -> "open" / "closed"
    opened: int -> int,       // PR key -> how many PRs Gitea ever opened for it
    openIssues: Set[int],     // plan issues still open, by task
    // -- volatile: lost in a crash --
    queue: List[Job],
    work: str,                // none / plan / exec / publish / review
    workGate: str,
    workCid: int,
    pubStage: int,
    cancel: bool,
    planWatch: Watch,
    prWatch: Watch,
    // -- the operator's view, and history the properties read --
    shown: Plan,
    taken: Set[int],          // review comments the agent took up
    answered: Set[int],       // ... and answered (pushed, asked, or reported)
    badApproval: bool,
    staleFeedback: bool,
  }

  // ---- helpers ------------------------------------------------------------

  pure def msGate(ms: int): str = if (ms == 1) "ms1" else "ms2"
  pure def prKeyOf(s: Sys): int = s.tid * 10 + s.ms
  pure def prState(s: Sys, key: int): str =
    if (s.giteaPr.keys().contains(key)) s.giteaPr.get(key) else "none"
  pure def planIds(s: Sys): int =
    s.planThread.foldl(0, (acc, c) => if (c.tid == s.tid and c.id > acc) c.id else acc)
  pure def prIds(s: Sys, key: int): int =
    s.prThread.foldl(0, (acc, c) => if (c.key == key and c.id > acc) c.id else acc)
  pure def max(a: int, b: int): int = if (a > b) a else b
  // Gitea opens a PR for a branch.
  pure def openPr(s: Sys, key: int): Sys =
    { ...s, giteaPr: s.giteaPr.put(key, "open"),
      opened: s.opened.put(key, (if (s.opened.keys().contains(key)) s.opened.get(key) else 0) + 1) }

  pure def queued(s: Sys, src: str): bool = s.queue.foldl(false, (acc, j) => acc or j.src == src)

  // _await_input
  pure def awaitInput(s: Sys, p: Pending): Sys = { ...s, st: "AWAIT_INPUT", pending: p }

  // _settled: a job that leaves a working state behind with nothing running
  // moves to AWAIT_INPUT with the restart choice.
  pure def settle(s: Sys): Sys =
    if (s.work == "none" and WORKING.contains(s.st))
      awaitInput(s, { kind: "boot", g: "", interrupted: s.st })
    else s

  pure def endWork(s: Sys): Sys = { ...s, work: "none", workGate: "", workCid: 0, cancel: false }

  // _run_planning_turn, first half: the turn is now in flight.
  pure def startPlanning(s: Sys, g: str): Sys =
    { ...s, st: if (g == "arch") "PLANNING_ARCH" else "MILESTONE_PLANNING",
      work: "plan", workGate: g }

  // _begin_project
  pure def beginProject(s: Sys, giteaUp: bool): Sys =
    val t = s.tid + 1
    startPlanning(
      { ...s, tid: t, task: true, gate: NO_GATE,
        revs: Map("arch" -> 0, "ms1" -> 0, "ms2" -> 0), lastRev: 0, ms: 1, pending: NO_PENDING,
        prOn: false, prKey: 0, prLast: 0, reviewing: 0,
        boardOn: giteaUp, boardPosted: false, boardLast: 0,
        openIssues: if (giteaUp) s.openIssues.union(Set(t)) else s.openIssues,
        planWatch: OFF, prWatch: OFF },
      "arch")

  // _run_planning_turn, second half: the planner submitted a plan.
  // _post_plan_to_board then watches the issue from here (a fresh watch
  // skips what is already there).
  pure def planSubmitted(s: Sys, giteaUp: bool): Sys =
    val g = s.workGate
    val rev = s.lastRev + 1
    val posted = s.boardOn and giteaUp
    val t = endWork({ ...s,
      revs: s.revs.set(g, rev), lastRev: rev,
      gate: { on: true, g: g, rev: rev, approved: false },
      st: if (g == "arch") "AWAIT_ARCH_SIGNOFF" else "AWAIT_MILESTONE_SIGNOFF",
      shown: { g: g, rev: rev } })
    if (posted)
      { ...t, boardPosted: true, boardLast: planIds(t), planWatch: { on: true, floor: planIds(t) } }
    else t

  // _execute_milestone: start
  pure def startExec(s: Sys): Sys = { ...s, st: "EXECUTING", work: "exec" }

  // _apply_gate_reply and Gate.evaluate, for Matrix and Gitea alike.
  // A job's `saw` doubles as its arrival stamp: the revision up when it was
  // sent (Matrix) or the latest plan posted before it (Gitea). A reply that
  // names no revision applies only if that is the current one.
  pure def applyGate(s: Sys, j: Job): Sys =
    val current = { g: s.gate.g, rev: s.gate.rev }
    if (s.gate.approved)
      if (j.kind == "start") startExec(s) else s
    else if (j.kind == "approveRev" and j.saw.rev != s.gate.rev)
      s  // names another revision: ignored with a notice
    else if (j.kind != "approveRev" and j.saw.rev != s.gate.rev)
      s  // about another revision: not applied, the agent says so
    else if (j.kind == "approve" or j.kind == "approveRev")
      val s2 = { ...s, badApproval: s.badApproval or j.saw != current }
      if (s.gate.g == "arch")
        startPlanning({ ...s2, gate: NO_GATE }, msGate(s.ms))
      else
        { ...s2, gate: { ...s.gate, approved: true }, st: "AWAIT_MILESTONE_SIGNOFF" }
    else
      // Anything else is rejection feedback, and the gate is planned again.
      startPlanning({ ...s, staleFeedback: s.staleFeedback or j.saw != current, gate: NO_GATE },
                    s.gate.g)

  // _return_to_pr_open / _watch_pr
  pure def returnToPrOpen(s: Sys): Sys =
    { ...s, st: "PR_OPEN", prWatch: if (s.prOn) { on: true, floor: s.prLast } else OFF }

  // _abandon: IDLE first, then close the issue and the PR, abort intents.
  pure def abandon(s: Sys): Sys =
    { ...s, st: "IDLE", task: false, gate: NO_GATE, pending: NO_PENDING,
      prOn: false, reviewing: 0, boardOn: false, boardPosted: false, intents: Set(),
      openIssues: s.openIssues.exclude(Set(s.tid)),
      giteaPr: if (s.prOn) s.giteaPr.put(s.prKey, "closed") else s.giteaPr,
      planWatch: OFF, prWatch: OFF }

  // _handle_pr_open: 'project complete' closes the plan issue; the PR is
  // the operator's to merge.
  pure def complete(s: Sys): Sys =
    { ...s, st: "IDLE", task: false, gate: NO_GATE, pending: NO_PENDING,
      prOn: false, reviewing: 0, boardOn: false, boardPosted: false,
      openIssues: s.openIssues.exclude(Set(s.tid)), planWatch: OFF, prWatch: OFF }

  // _resume_after_restart
  pure def resume(s: Sys, p: Pending, j: Job): Sys =
    if (j.kind != "continue" and j.kind != "replan")
      { ...s, pending: p }  // asked again
    else if (p.interrupted == "PR_OPEN")
      returnToPrOpen(s)
    else if (j.kind == "continue" and p.interrupted == "EXECUTING" and s.gate.approved)
      { ...s, st: "AWAIT_MILESTONE_SIGNOFF" }
    else
      val g =
        if (p.interrupted == "PLANNING_ARCH") "arch"
        else if (p.interrupted == "MILESTONE_PLANNING" or p.interrupted == "EXECUTING") msGate(s.ms)
        else if (s.gate.g != "") s.gate.g
        else if (s.revs.get("arch") == 0) "arch"
        else msGate(s.ms)
      startPlanning({ ...s, gate: NO_GATE }, g)

  // _handle_answer
  pure def answer(s: Sys, j: Job, giteaUp: bool): Sys =
    val p = s.pending
    val s0 = { ...s, pending: NO_PENDING }
    if (p.kind == "")
      if (j.kind == "use") beginProject({ ...s, planWatch: OFF, prWatch: OFF }, giteaUp) else s
    else if (p.kind == "question" or p.kind == "planning_failed") startPlanning(s0, p.g)
    else if (p.kind == "review_stopped") returnToPrOpen(s0)
    else if (p.kind == "milestone_stopped") { ...s0, st: "AWAIT_MILESTONE_SIGNOFF" }
    else if (p.kind == "boot") resume(s0, p, j)
    else s0

  // _handle_now, for a Matrix message the worker has reached.
  pure def handleMatrix(s: Sys, j: Job, giteaUp: bool): Sys =
    if (j.kind == "abandon")
      if (s.st == "IDLE") s else abandon(s)
    else if (s.st == "IDLE")
      if (j.kind == "use") beginProject(s, giteaUp) else s
    else if (SIGNOFF.contains(s.st))
      if (not(s.gate.on))
        awaitInput(s, { kind: "boot", g: "",
          interrupted: if (s.st == "AWAIT_ARCH_SIGNOFF") "PLANNING_ARCH" else "MILESTONE_PLANNING" })
      else applyGate(s, j)
    else if (s.st == "PR_OPEN")
      if (j.kind == "next" and s.ms < MAX_MS)
        startPlanning({ ...s, ms: s.ms + 1, gate: NO_GATE, prOn: false, prWatch: OFF }, msGate(s.ms + 1))
      else if (j.kind == "complete") complete(s)
      else s
    else if (s.st == "AWAIT_INPUT") answer(s, j, giteaUp)
    else s  // busy

  // _handle_plan_comment_now
  pure def handlePlanComment(s: Sys, j: Job): Sys =
    if (not(SIGNOFF.contains(s.st))) s  // not taken
    else applyGate({ ...s, boardLast: max(s.boardLast, j.cid) }, j)

  // _handle_review_comment_now, first half: the executor starts on it,
  // with the comment marked (durably) until it returns.
  pure def handleReviewComment(s: Sys, j: Job): Sys =
    if (s.st != "PR_OPEN") s  // not taken
    else { ...s, prLast: max(s.prLast, j.cid), taken: s.taken.union(Set(j.cid)),
           reviewing: j.cid, work: "review", workCid: j.cid }

  pure def handleJob(s: Sys, j: Job, giteaUp: bool): Sys =
    if (j.src == "matrix") handleMatrix(s, j, giteaUp)
    else if (j.src == "plan") handlePlanComment(s, j)
    else handleReviewComment(s, j)

  // Orchestrator.start, after a crash: reconcile intents against Gitea,
  // then resume, re-ask, or ask the restart question.
  pure def boot(s: Sys, giteaUp: bool): Sys =
    val unresolved = s.intents.size() > 0 and not(giteaUp)
    val s0 = { ...s, intents: if (giteaUp) Set() else s.intents }
    if (unresolved)
      awaitInput(s0, { kind: "boot", g: "", interrupted: s.st })
    else if (s.st == "AWAIT_INPUT" and s.pending.kind != "")
      s0
    else if (s.st == "EXECUTING" and giteaUp and prState(s, prKeyOf(s)) == "open")
      // L8 boot step 3: the milestone's PR is open, so it was published and only
      // PR_OPEN went unrecorded. Adopt it; nothing watched it, so from 0.
      { ...s0, st: "PR_OPEN", prOn: true, prKey: prKeyOf(s), prLast: 0,
        prWatch: { on: true, floor: 0 }, planWatch: OFF }
    else if (s.st == "PR_OPEN")
      // A comment whose answer never finished is delivered again.
      { ...s0, prWatch: if (s.prOn)
          { on: true, floor: if (s.reviewing > 0) s.reviewing - 1 else s.prLast }
        else OFF }
    else if (s.st == "IDLE" or (SIGNOFF.contains(s.st) and s.gate.on))
      if (s.gate.on and s.boardPosted) { ...s0, planWatch: { on: true, floor: s.boardLast } } else s0
    else
      awaitInput(s0, { kind: "boot", g: "", interrupted: s.st })

  // ---- the system ---------------------------------------------------------

  var s: Sys

  action init = s' = {
    st: "IDLE", tid: 0, task: false, gate: NO_GATE,
    revs: Map("arch" -> 0, "ms1" -> 0, "ms2" -> 0), lastRev: 0, ms: 1, pending: NO_PENDING,
    prOn: false, prKey: 0, prLast: 0, reviewing: 0,
    boardOn: false, boardPosted: false, boardLast: 0,
    intents: Set(),
    planThread: List(), prThread: List(), nextId: 1,
    giteaPr: Map(), opened: Map(), openIssues: Set(),
    queue: List(), work: "none", workGate: "", workCid: 0, pubStage: 0, cancel: false,
    planWatch: OFF, prWatch: OFF,
    shown: NO_PLAN, taken: Set(), answered: Set(),
    badApproval: false, staleFeedback: false,
  }

  // Each action has a deterministic form taking its choices as parameters,
  // so scenarios (orchestrator_scenarios.qnt) can replay a trace exactly;
  // `step` picks the parameters at random.

  pure val MATRIX_KINDS = Set("use", "approve", "approveRev", "text", "start", "next",
                              "complete", "abandon", "continue", "replan", "stop")
  pure val COMMENT_KINDS = Set("approve", "approveRev", "text", "start")

  // The operator sends a Matrix message (control._handle, on the
  // transport's thread). While long work runs it only gets "busy"; `stop`
  // reaches the cancel flag whenever the worker is busy.
  action send(kind: str): bool = all {
    s.tid < MAX_TASKS or kind != "use",
    if (kind == "stop")
      s' = { ...s, cancel: s.cancel or s.work != "none" }
    else if (LONG_WORK.contains(s.work))
      s' = s
    else all {
      s.queue.length() < MAX_QUEUE,
      s' = { ...s, queue: s.queue.append({ src: "matrix", kind: kind, saw: s.shown, cid: 0 }) },
    },
  }

  // The operator comments on the plan issue.
  action commentPlan(kind: str): bool = all {
    s.task, s.openIssues.contains(s.tid), s.nextId <= MAX_IDS,
    s' = { ...s, nextId: s.nextId + 1,
           planThread: s.planThread.append({ id: s.nextId, tid: s.tid, kind: kind, saw: s.shown }) },
  }

  // The operator comments on the open PR.
  action commentPr: bool = all {
    s.prOn, s.nextId <= MAX_IDS, prState(s, s.prKey) == "open",
    s' = { ...s, nextId: s.nextId + 1, prThread: s.prThread.append({ id: s.nextId, key: s.prKey }) },
  }

  // A poller delivers the next new comment (CommentPoller.poll_once). It
  // waits on the worker, so one job per poller at a time.
  action pollPlan: bool = {
    val c = s.planThread.foldl({ id: 0, tid: 0, kind: "", saw: NO_PLAN },
      (acc, x) => if (x.tid == s.tid and x.id > s.planWatch.floor and (acc.id == 0 or x.id < acc.id)) x else acc)
    all {
      s.planWatch.on, c.id > 0, not(queued(s, "plan")), s.queue.length() < MAX_QUEUE,
      s' = { ...s, planWatch: { on: true, floor: c.id },
             queue: s.queue.append({ src: "plan", kind: c.kind, saw: c.saw, cid: c.id }) },
    }
  }

  action pollPr: bool = {
    val c = s.prThread.foldl({ id: 0, key: 0 },
      (acc, x) => if (x.key == s.prKey and x.id > s.prWatch.floor and (acc.id == 0 or x.id < acc.id)) x else acc)
    all {
      s.prWatch.on, c.id > 0, not(queued(s, "pr")), s.queue.length() < MAX_QUEUE,
      s' = { ...s, prWatch: { on: true, floor: c.id },
             queue: s.queue.append({ src: "pr", kind: "review", saw: NO_PLAN, cid: c.id }) },
    }
  }

  // The worker takes the next job, one at a time, and settles it.
  action runJob(giteaUp: bool): bool = all {
    s.work == "none", s.queue.length() > 0,
    s' = settle(handleJob({ ...s, queue: s.queue.tail() }, s.queue.head(), giteaUp)),
  }

  // A planning turn ends.
  action finishPlanning(outcome: str, giteaUp: bool): bool = all {
    s.work == "plan", outcome != "stopped" or s.cancel,
    s' =
      if (outcome == "submit") planSubmitted(s, giteaUp)
      else if (outcome == "ask" or outcome == "stopped")
        awaitInput(endWork(s), { kind: "question", g: s.workGate, interrupted: "" })
      else if (outcome == "error")
        awaitInput(endWork(s), { kind: "planning_failed", g: s.workGate, interrupted: "" })
      else settle(endWork(s)),
  }

  // The executor finishes a milestone.
  action finishExec(outcome: str): bool = all {
    s.work == "exec", outcome != "stopped" or s.cancel,
    s' =
      if (outcome == "completed") { ...s, work: "publish", pubStage: 0, cancel: false }
      else if (outcome == "exception") settle(endWork(s))
      else awaitInput(endWork(s), { kind: "milestone_stopped", g: "", interrupted: "" }),
  }

  // Publishing, one step at a time, so a crash can land between them
  // (MilestonePublisher.publish and control._publish_milestone):
  // 0 commit, begin_intent; 1 push, then adopt an open PR for the branch or
  // open one ("lost": Gitea opened it but the response never came back);
  // 2 commit_intent; 3 record the PR, watch it, PR_OPEN.
  action publishStep(outcome: str): bool = {
    val key = prKeyOf(s)
    all {
      s.work == "publish",
      s' =
        if (s.pubStage == 0)
          if (outcome == "nochanges" or outcome == "fail")
            awaitInput(endWork(s), { kind: "milestone_stopped", g: "", interrupted: "" })
          else if (outcome == "exception") settle(endWork(s))
          else { ...s, intents: s.intents.union(Set(key)), pubStage: 1 }
        else if (s.pubStage == 1)
          if (outcome == "exception") settle(endWork(s))
          else if (outcome == "fail" or outcome == "nochanges")
            awaitInput(endWork(s), { kind: "milestone_stopped", g: "", interrupted: "" })
          else if (prState(s, key) == "open")
            { ...s, pubStage: 2 }  // adopted, never a second
          else if (outcome == "lost")
            awaitInput(endWork(openPr(s, key)),
                       { kind: "milestone_stopped", g: "", interrupted: "" })
          else { ...openPr(s, key), pubStage: 2 }
        else if (s.pubStage == 2)
          { ...s, intents: s.intents.exclude(Set(key)), pubStage: 3 }
        else
          endWork({ ...s, prOn: true, prKey: key, prLast: prIds(s, key),
                    prWatch: { on: true, floor: prIds(s, key) }, planWatch: OFF,
                    st: "PR_OPEN" }),
    }
  }

  // The executor finishes responding to a review comment.
  action finishReview(outcome: str): bool = all {
    s.work == "review", outcome != "stopped" or s.cancel,
    s' =
      val t = endWork({ ...s, answered: s.answered.union(Set(s.workCid)), reviewing: 0 })
      if (outcome == "stopped" or outcome == "escalate")
        awaitInput(t, { kind: "review_stopped", g: "", interrupted: "" })
      else t,
  }

  // The process dies and starts again: volatile state is gone.
  action crash(giteaUp: bool): bool =
    s' = boot({ ...s, queue: List(), work: "none", workGate: "", workCid: 0, pubStage: 0,
                cancel: false, planWatch: OFF, prWatch: OFF }, giteaUp)

  action step = {
    nondet kind = MATRIX_KINDS.oneOf()
    nondet ckind = COMMENT_KINDS.oneOf()
    nondet up = Set(true, false).oneOf()
    nondet planOutcome = Set("submit", "ask", "error", "stopped", "exception").oneOf()
    nondet workOutcome = Set("completed", "stopped", "escalate", "exception").oneOf()
    nondet pubOutcome = Set("ok", "nochanges", "fail", "lost", "exception").oneOf()
    any {
      send(kind), commentPlan(ckind), commentPr, pollPlan, pollPr, runJob(up),
      finishPlanning(planOutcome, up), finishExec(workOutcome), publishStep(pubOutcome),
      finishReview(workOutcome), crash(up),
    }
  }

  // ---- properties ---------------------------------------------------------

  // #66: a working state only while its work is running.
  val no_stranding = WORKING.contains(s.st) implies s.work != "none"

  // Waiting on the operator always means a question to answer.
  val await_input_has_a_question = s.st == "AWAIT_INPUT" implies s.pending.kind != ""

  // A milestone runs only on an approved plan for that milestone.
  val execution_needs_approval =
    s.st == "EXECUTING" implies (s.gate.on and s.gate.approved and s.gate.g == msGate(s.ms))

  // An approval only ever approves the plan the operator was looking at.
  val approval_only_of_what_was_seen = not(s.badApproval)

  // Rejection feedback only ever lands on the plan it was written about.
  val feedback_only_on_what_was_seen = not(s.staleFeedback)

  // Never more than one PR opened for a milestone.
  val one_pr_per_milestone = s.opened.keys().forall(k => s.opened.get(k) <= 1)

  // A review comment the agent took up is answered, being worked on, or
  // marked to be delivered again: never silently dropped.
  val review_comments_answered =
    s.taken.forall(id => s.answered.contains(id) or (s.work == "review" and s.workCid == id)
                         or s.reviewing == id)

  // With no task, no plan issue is left open.
  val no_open_issue_without_a_task = not(s.task) implies s.openIssues.size() == 0

  val all_properties = and {
    no_stranding, await_input_has_a_question, execution_needs_approval,
    approval_only_of_what_was_seen, feedback_only_on_what_was_seen,
    one_pr_per_milestone, review_comments_answered, no_open_issue_without_a_task,
  }
}