// Graded runs as fine-tuning data: one chat conversation per run, in the
// JSONL shape tinker-cookbook's FromConversationFileBuilder reads (OpenAI
// style messages, tool calls in `tool_calls`). Exporting a machine with its
// descendants gives one dataset for the model they share.

import { caseOf } from "./converse.ts";
import type { Tool } from "./harness.ts";
import type { Machine } from "./machine.ts";
import type { Store } from "./store.ts";
import type { Message, Run } from "./types.ts";

/** Message content: a string, or parts when the assistant's reasoning is included. */
export type ChatContent = string | Array<{ type: "thinking"; thinking: string } | { type: "text"; text: string }>;

export type ChatMessage = {
  role: string;
  content: ChatContent;
  tool_calls?: Array<{ id: string; type: "function"; function: { name: string; arguments: string } }>;
  tool_call_id?: string;
  name?: string;
  /** In a whole-run example, whether this assistant message is a training target (false for a failed step). */
  trainable?: boolean;
};

/** A tool as the model saw it, in the OpenAI function format trainers render into the system prompt. */
export type ToolSpec = { name: string; description: string; parameters: Record<string, unknown> };

export type ExportLine = {
  messages: ChatMessage[];
  tools?: ToolSpec[];
  machine: string;
  run: string;
  score: number;
  /** One example for a whole run: train on every assistant message marked trainable, not only the last. */
  whole_run?: boolean;
};

/**
 * Runs scoring at least `minScore`, as conversations. Runs without stored
 * instructions (recorded since 0.4.0) fall back to the machine's current ones.
 */
export async function exportRuns(
  store: Store,
  machines: Machine[],
  options: {
    minScore?: number;
    limit?: number;
    reasoning?: boolean;
    /**
     * Write the runs as if `as` had produced them: its instructions as the
     * system message and its transformReasoning on the reasoning. Used to turn
     * a parent's good runs into training data for a child that changes how
     * the model thinks (export a child "--from" its parent).
     */
    as?: Machine;
    /** Only runs started at or after this time (ISO 8601), e.g. one sampling batch. */
    since?: string;
    /** Only runs made for these cases (a conversation's case, else the run's input). */
    cases?: string[];
    /**
     * Rejection sampling: of each case's runs that clear minScore, keep the
     * best `bestOf` (ties go to the newest).
     */
    bestOf?: number;
    /**
     * One example per assistant message, each ending at it, and for a
     * conversation every turn of the agent, not only the last. For trainers
     * that learn only each example's last assistant message (such as
     * scripts/train_sft.py on Qwen3.5/3.6), this trains everything the agent
     * did: its questions, its tool calls, and its replies.
     */
    perTurn?: boolean;
  } = {},
): Promise<{ lines: ExportLine[]; counts: Record<string, { graded: number; kept: number }> }> {
  const minScore = options.minScore ?? 1;
  const cases = options.cases ? new Set(options.cases) : undefined;
  const lines: ExportLine[] = [];
  const counts: Record<string, { graded: number; kept: number }> = {};
  for (const machine of machines) {
    const all = await store.listRuns(machine.name, { limit: options.limit ?? 100_000 });
    const runs = all.filter(
      (r) =>
        r.scorecard &&
        r.harness !== "echo" &&
        (!r.conversation || r.conversation.final) &&
        (!options.since || r.startedAt >= options.since) &&
        (!cases || cases.has(caseOf(r))),
    );
    let kept = runs.filter((r) => r.scorecard!.total >= minScore);
    if (options.bestOf !== undefined) kept = bestPerCase(kept, options.bestOf);
    counts[machine.name] = { graded: runs.length, kept: kept.length };
    const target = options.as ?? machine;
    const reasoning = options.reasoning === false ? null : (target.transformReasoning ?? ((r: string) => r));
    // The tools the model was offered, so training renders the same prompt sampling did.
    const tools = target.tools.length ? target.tools.map((t) => ({ name: t.name, description: t.description, parameters: t.inputSchema })) : undefined;
    for (const run of kept.reverse()) {
      // A conversation's earlier turns are runs of their own, with the reasoning the last turn's history leaves out.
      const turns = options.perTurn && run.conversation ? all.filter((r) => r.conversation?.id === run.conversation!.id).sort((a, b) => a.conversation!.turn - b.conversation!.turn) : [run];
      for (const turn of turns) {
        const asRun = options.as ? { ...turn, instructions: target.instructions } : turn;
        const messages = conversation(asRun, target.instructions, reasoning);
        const examples = options.perTurn ? endingAtEachReply(messages, messages.length - ownMessages(asRun, reasoning), target.tools) : [messages];
        for (const m of examples) lines.push({ messages: m, ...(tools ? { tools } : {}), machine: target.name, run: turn.id, score: run.scorecard!.total });
      }
    }
  }
  return { lines: uniformContent(lines), counts };
}

/** Per case, the best `k` runs, best first (ties: newest first, as listed). */
function bestPerCase(runs: Run[], k: number): Run[] {
  const byCase = new Map<string, Run[]>();
  for (const run of runs) byCase.set(caseOf(run), [...(byCase.get(caseOf(run)) ?? []), run]);
  const kept = new Set([...byCase.values()].flatMap((rs) => [...rs].sort((a, b) => b.scorecard!.total - a.scorecard!.total).slice(0, k)));
  return runs.filter((r) => kept.has(r));
}

/** How many of a converted conversation's messages are the run's own turns (after its input). */
function ownMessages(run: Run, reasoning: ((r: string) => string) | null): number {
  const turns = run.trajectory[0]?.role === "user" && run.trajectory[0].content === run.input ? run.trajectory.slice(1) : run.trajectory;
  return toChat(turns, reasoning).length;
}

/**
 * The conversation cut after each assistant message from `start` on: one
 * example per reply, with the context the model saw at that step. Harnesses
 * do not send earlier steps' reasoning back to the model, so earlier
 * assistant messages lose their thinking parts. Replies whose tool calls all
 * failed (see stepFailed) are left out, so the model is not taught to repeat
 * them.
 */
function endingAtEachReply(messages: ChatMessage[], start: number, tools: Tool[] = []): ChatMessage[][] {
  const out: ChatMessage[][] = [];
  for (let i = start; i < messages.length; i++) {
    if (messages[i].role === "assistant" && !stepFailed(messages, i, tools)) out.push([...messages.slice(0, i).map(withoutThinking), messages[i]]);
  }
  return out;
}

/** An assistant message without its thinking parts (other messages unchanged). */
export function withoutThinking(m: ChatMessage): ChatMessage {
  return Array.isArray(m.content) && m.role === "assistant" ? { ...m, content: m.content.filter((p) => p.type !== "thinking") } : m;
}

/**
 * Whether the assistant message at `i` made tool calls and every one failed:
 * its result is a harness error ({"error": ...}, what a thrown error becomes)
 * or one its tool's `failed` recognizes (e.g. a rejected input).
 */
export function stepFailed(messages: ChatMessage[], i: number, tools: Tool[] = []): boolean {
  const calls = messages[i].tool_calls ?? [];
  if (calls.length === 0) return false;
  const byName = new Map(tools.map((t) => [t.name, t]));
  const results = messages.slice(i + 1).filter((m) => m.role === "tool" && calls.some((c) => c.id === m.tool_call_id));
  return (
    results.length > 0 &&
    results.every((m) => {
      const result = contentText(m.content);
      const name = m.name ?? calls.find((c) => c.id === m.tool_call_id)?.function.name;
      return isToolError(result) || Boolean(name && byName.get(name)?.failed?.(result));
    })
  );
}

function contentText(content: ChatContent): string {
  return typeof content === "string" ? content : content.map((p) => ("text" in p ? p.text : "")).join("");
}

/** A harness's error result: a JSON object with only an "error" key. */
function isToolError(text: string): boolean {
  try {
    const parsed = JSON.parse(text);
    return Boolean(parsed && typeof parsed === "object" && "error" in parsed && Object.keys(parsed).length === 1);
  } catch {
    return false;
  }
}

/**
 * Every message's content as parts, if any message in the file uses parts.
 * tinker-cookbook loads JSONL through Hugging Face datasets, whose Arrow
 * columns cannot mix strings with lists: one assistant message carrying a
 * thinking part makes every content in the file need to be a list.
 */
export function uniformContent(lines: ExportLine[]): ExportLine[] {
  if (!lines.some((l) => l.messages.some((m) => Array.isArray(m.content)))) return lines;
  return lines.map((l) => ({
    ...l,
    messages: l.messages.map((m) => (typeof m.content === "string" ? { ...m, content: [{ type: "text" as const, text: m.content }] } : m)),
  }));
}

/**
 * System instructions, prior turns, the input, then the run's own turns with
 * tool calls folded into their assistant message. Reasoning on the run's own
 * turns becomes a "thinking" part, passed through `reasoning` (a transform,
 * or null to leave reasoning out).
 */
export function conversation(run: Run, fallbackInstructions: string, reasoning: ((r: string) => string) | null = (r) => r): ChatMessage[] {
  const out: ChatMessage[] = [{ role: "system", content: run.instructions ?? fallbackInstructions }];
  for (const m of run.history) out.push(...toChat([m]));
  out.push({ role: "user", content: run.input });
  // The trajectory repeats the input as its first message.
  const turns = run.trajectory[0]?.role === "user" && run.trajectory[0].content === run.input ? run.trajectory.slice(1) : run.trajectory;
  out.push(...toChat(turns, reasoning));
  return out;
}

function toChat(messages: Message[], reasoning: ((r: string) => string) | null = null): ChatMessage[] {
  const out: ChatMessage[] = [];
  const thinking = (m: Message, text: string): ChatContent =>
    reasoning && m.reasoning ? [{ type: "thinking", thinking: reasoning(m.reasoning) }, ...(text ? [{ type: "text" as const, text }] : [])] : text;
  for (const m of messages) {
    if (m.role === "assistant" && m.toolCallId) {
      const call = { id: m.toolCallId, type: "function" as const, function: { name: m.name ?? "tool", arguments: JSON.stringify(m.toolInput ?? {}) } };
      const last = out.at(-1);
      // Text and the calls that follow it are one assistant turn.
      if (last?.role === "assistant" && !(m.reasoning && reasoning)) (last.tool_calls ??= []).push(call);
      else out.push({ role: "assistant", content: thinking(m, ""), tool_calls: [call] });
    } else if (m.role === "tool") {
      out.push({ role: "tool", content: m.content, tool_call_id: m.toolCallId, ...(m.name ? { name: m.name } : {}) });
    } else if (m.role === "assistant") {
      out.push({ role: "assistant", content: thinking(m, m.content) });
    } else {
      out.push({ role: m.role, content: m.content });
    }
  }
  return out;
}