// A conversation between two machines: an agent and its counterpart, which
// plays the other side (for example a simulated user with a hidden scenario).
// The counterpart speaks first, from the case. They alternate until the
// agent's turn meets the counterpart's `until` or the agent has taken
// `maxTurns` turns. Every turn of each side is an ordinary stored run whose
// history is the conversation so far from that side's point of view, so the
// agent's last turn holds the whole conversation and is the run that is
// graded, exported, and compared like any other.

import type { Machine } from "./machine.ts";
import { newRunId } from "./store.ts";
import type { ConversationTurn, Message, Run } from "./types.ts";
import type { RunOptions, Workspace } from "./workspace.ts";

export const DEFAULT_MAX_TURNS = 8;

export type ConverseOptions = Omit<RunOptions, "history"> & {
  /** Grade the counterpart's last turn too (against the counterpart's own rubric). */
  evaluateCounterpart?: boolean;
  /** Called as each side speaks, for live transcripts. */
  onTurn?: (side: "agent" | "counterpart", run: Run) => void;
};

export type Conversation = {
  id: string;
  /** The agent's last turn: the whole conversation, graded when `evaluate` is set. */
  run: Run;
  /** Every turn of the agent, in order (the last is `run`). */
  agentRuns: Run[];
  /** Every turn of the counterpart, in order. */
  counterpartRuns: Run[];
  ended: "done" | "max_turns";
};

/** Run one conversation between `agent` and its counterpart, starting from `caseInput`. */
export async function converse(workspace: Workspace, agent: Machine, caseInput: string, options: ConverseOptions = {}): Promise<Conversation> {
  const counterpart = agent.counterpart;
  if (!counterpart) throw new Error(`Machine "${agent.name}" has no counterpart to converse with.`);
  const { evaluate, evaluateCounterpart, onTurn, model, ...runOptions } = options;
  const other = counterpart.machine;
  const maxTurns = counterpart.maxTurns ?? DEFAULT_MAX_TURNS;
  const id = newRunId();
  const agentHistory: Message[] = [];
  const otherHistory: Message[] = [];
  const agentRuns: Run[] = [];
  const otherRuns: Run[] = [];
  let otherInput = counterpart.opening ? counterpart.opening(caseInput) : caseInput;
  let ended: Conversation["ended"] = "max_turns";

  for (let turn = 1; turn <= maxTurns; turn++) {
    // The counterpart runs on its own model: overriding the agent's model never changes the other side.
    const said = await workspace.run(other, otherInput, { ...runOptions, model: undefined, revision: undefined, history: [...otherHistory] });
    otherRuns.push(said);
    otherHistory.push(...ownTurn(said));
    onTurn?.("counterpart", said);

    const reply = await workspace.run(agent, said.output || "(no message)", { ...runOptions, model, history: [...agentHistory] });
    agentRuns.push(reply);
    agentHistory.push(...ownTurn(reply));
    onTurn?.("agent", reply);
    if (counterpart.until?.(reply)) {
      ended = "done";
      break;
    }
    otherInput = reply.output || "(no reply)";
  }

  const tag = (runs: Run[], withName: string) =>
    runs.map((run, i): Run => {
      const final = i === runs.length - 1;
      const conversation: ConversationTurn = { id, case: caseInput, with: withName, turn: i + 1, final, ...(final ? { ended } : {}) };
      return { ...run, conversation };
    });
  const taggedAgent = tag(agentRuns, other.name);
  const taggedOther = tag(otherRuns, agent.name);
  await Promise.all([...taggedAgent, ...taggedOther].map((run) => workspace.store.putRun(run)));

  let last = taggedAgent.at(-1)!;
  if (evaluate) {
    last = await workspace.evaluate(agent, last, { signal: options.signal });
    taggedAgent[taggedAgent.length - 1] = last;
  }
  if (evaluateCounterpart) {
    taggedOther[taggedOther.length - 1] = await workspace.evaluate(other, taggedOther.at(-1)!, { signal: options.signal });
  }
  return { id, run: last, agentRuns: taggedAgent, counterpartRuns: taggedOther, ended };
}

/**
 * Run one case of a machine: a single run, or for a machine with a
 * counterpart, a whole conversation, returning the agent's last turn.
 */
export async function runCase(workspace: Workspace, machine: Machine, input: string, options: RunOptions = {}): Promise<Run> {
  if (!machine.counterpart) return workspace.run(machine, input, options);
  const { history: _history, ...rest } = options;
  return (await converse(workspace, machine, input, rest)).run;
}

/** A run's own messages as history for its next turn: its input, then what it did (tool calls included). */
function ownTurn(run: Run): Message[] {
  const turn = run.trajectory[0]?.role === "user" ? run.trajectory : [{ role: "user" as const, content: run.input }, ...run.trajectory];
  // Reasoning stays on the run; history carries what was said and done.
  return turn.map(({ reasoning: _reasoning, ...m }) => m);
}

/** The case a run was made for: its conversation's starting case, or its own input. */
export function caseOf(run: Run): string {
  return run.conversation?.case ?? run.input;
}