// Fine-tuning a machine's model on good runs: pick the runs, write them as
// training data, and run supervised fine-tuning on Tinker through
// tinker-cookbook (scripts/train_sft.py). Whether the result is better is
// decided elsewhere, by running it against the live version on held-out
// cases (see sdk.ts, train()).

import { spawn } from "node:child_process";
import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { caseOf } from "./converse.ts";
import { conversation, stepFailed, uniformContent, withoutThinking, type ChatMessage, type ExportLine, type ToolSpec } from "./export.ts";
import type { Tool } from "./harness.ts";
import type { Machine } from "./machine.ts";
import { TINKER_TRAIN_PRICES } from "./pricing.ts";
import { rubricFingerprint } from "./rubric.ts";
import type { Run } from "./types.ts";

export type DataSelection = {
  /** Lowest scorecard total to train on. Default 1: only runs that met every criterion. */
  minScore?: number;
  /** Keep at most this many runs per case, best first (then newest). Default 2. */
  bestOf?: number;
  /** Cases never to train on (the held-out inputs). */
  exclude?: Set<string>;
  /** Only runs by these models, when given (e.g. the live model, for rejection sampling). */
  models?: Set<string>;
};

/** Graded runs worth training on: scored under the current rubric, at least minScore, the best few per case. */
export function selectRuns(machine: Machine, runs: Run[], options: DataSelection = {}): Run[] {
  const minScore = options.minScore ?? 1;
  const bestOf = options.bestOf ?? 2;
  const fp = rubricFingerprint(machine.rubric);
  const byCase = new Map<string, Run[]>();
  for (const run of runs) {
    if (run.machine !== machine.name || run.harness === "echo" || !run.scorecard) continue;
    if (run.conversation && !run.conversation.final) continue;
    if (run.scorecard.rubric !== fp || run.scorecard.total < minScore) continue;
    if (options.exclude?.has(caseOf(run)) || options.exclude?.has(run.input)) continue;
    if (options.models && !options.models.has(run.model)) continue;
    // Train only on runs made under the instructions the model will run with (not, say, an optimizer's rewrite).
    if (run.instructions !== undefined && run.instructions !== machine.instructions) continue;
    byCase.set(caseOf(run), [...(byCase.get(caseOf(run)) ?? []), run]);
  }
  const out: Run[] = [];
  for (const group of byCase.values()) {
    group.sort((a, b) => b.scorecard!.total - a.scorecard!.total || b.id.localeCompare(a.id));
    out.push(...group.slice(0, bestOf));
  }
  return out;
}

/**
 * One training example per assistant step of a conversation: the context the
 * model saw at that step, ending with what it produced. Earlier turns of a
 * conversation carry no reasoning (harnesses do not send it back across
 * turns). Within the run, `keepStepReasoning` keeps earlier steps' thinking
 * in the context: a chat-completions harness sends each step's message back
 * whole, reasoning included, and Tinker's endpoints render it (checked against
 * their prompt_tokens for Qwen3.6 and Inkling-Small). Without it, earlier
 * steps lose their thinking. Each example is trained on its last message only
 * (train_sft.py's default). Steps whose tool calls all failed (e.g. an edit
 * whose snippet was not found) are left out, so the model is not taught to
 * repeat them; they stay in later steps' context.
 */
export function stepExamples(messages: ChatMessage[], options: { skipFailedSteps?: boolean; tools?: Tool[]; keepStepReasoning?: boolean } = {}): ChatMessage[][] {
  const skipFailed = options.skipFailedSteps ?? true;
  const lastUser = messages.findLastIndex((m) => m.role === "user");
  const out: ChatMessage[][] = [];
  for (let i = lastUser + 1; i < messages.length; i++) {
    if (messages[i].role !== "assistant" || (skipFailed && stepFailed(messages, i, options.tools))) continue;
    out.push([...messages.slice(0, i).map((m, j) => (options.keepStepReasoning && j > lastUser ? m : withoutThinking(m))), messages[i]]);
  }
  return out;
}

/** Harnesses that send a step's whole reply back to the model, reasoning included (chat completions). */
const RESENDS_REASONING = new Set(["tinker", "openrouter", "openai", "chat-completions"]);

/**
 * One example for a whole run: its conversation, with every assistant
 * message a target except steps whose tool calls all failed. When the harness
 * sends each step back with its reasoning (RESENDS_REASONING), every step's
 * prompt is a prefix of this sequence, so it trains the same targets on the
 * same contexts as stepExamples at a fraction of the tokens: on an
 * Inkling-Small dev run, 7,699 tokens against 121,144 for its 29 step
 * prompts, which were token for token prefixes of it.
 */
export function wholeRunExample(messages: ChatMessage[], tools: Tool[] = []): ChatMessage[] {
  const lastUser = messages.findLastIndex((m) => m.role === "user");
  return messages.map((m, i) => (m.role === "assistant" && i > lastUser ? { ...m, trainable: !stepFailed(messages, i, tools) } : m));
}

/** Whether a base model's training data can be whole runs (see wholeRunExample): verified for Inkling's renderer. */
export function wholeRunsFor(baseModel: string): boolean {
  return baseModel.startsWith("thinkingmachines/Inkling");
}

/**
 * Training examples from runs, with the tools the model was offered: one per
 * assistant step (see stepExamples), or with `wholeRuns`, one per run whose
 * harness resends reasoning (see wholeRunExample).
 */
export function trainingLines(machine: Machine, runs: Run[], options: { reasoning?: boolean; wholeRuns?: boolean } = {}): ExportLine[] {
  const reasoning = options.reasoning === false ? null : (machine.transformReasoning ?? ((r: string) => r));
  const tools: ToolSpec[] = machine.tools.map((t) => ({ name: t.name, description: t.description, parameters: t.inputSchema }));
  const lines: ExportLine[] = [];
  for (const run of runs) {
    if (options.wholeRuns && RESENDS_REASONING.has(run.harness)) {
      const messages = wholeRunExample(conversation(run, machine.instructions, reasoning), machine.tools);
      lines.push({ messages, ...(tools.length ? { tools } : {}), machine: machine.name, run: run.id, score: run.scorecard!.total, whole_run: true });
      continue;
    }
    for (const messages of stepExamples(conversation(run, machine.instructions, reasoning), { tools: machine.tools, keepStepReasoning: RESENDS_REASONING.has(run.harness) })) {
      lines.push({ messages, ...(tools.length ? { tools } : {}), machine: machine.name, run: run.id, score: run.scorecard!.total });
    }
  }
  return uniformContent(lines);
}

/**
 * A rough upper estimate of the tokens a training file holds: 2.5 characters
 * per token over the messages and tool declarations as JSON. (Measured on dev
 * runs with Qwen3's renderer: about 4 characters per token.)
 */
export function estimateTokens(lines: ExportLine[]): number {
  return Math.ceil(lines.reduce((n, l) => n + JSON.stringify(l.messages).length + (l.tools ? JSON.stringify(l.tools).length : 0), 0) / 2.5);
}

/** Estimated training cost in USD: tokens × epochs at Tinker's training price for the base model. */
export function trainingCost(baseModel: string, tokens: number, epochs: number): number {
  const price = TINKER_TRAIN_PRICES[baseModel];
  if (price === undefined) throw new Error(`No Tinker training price for ${baseModel}; cannot bound the cost of training it.`);
  return (tokens * epochs * price) / 1_000_000;
}

export type TrainJob = {
  /** JSONL training data. */
  data: string;
  /** Directory for the run's checkpoints and metrics, under .machine/training/ to show on the dashboard. */
  logDir: string;
  baseModel: string;
  epochs: number;
  learningRate?: number;
  batchSize?: number;
  /** Thinking effort the runs sampled at, 0 to 0.99, for models whose prompt carries it (Inkling). Default 0.9, the endpoint's. */
  effort?: number;
  log?: (line: string) => void;
  signal?: AbortSignal;
};

export type TrainResult = {
  /** Sampler path of the final checkpoint: the model id to run it with. */
  checkpoint: string;
  steps: number;
  /** Tokens processed across all steps (every epoch), as logged; what training is billed on. */
  tokens?: number;
  /** Training loss of the first and last logged steps, when logged. */
  firstLoss?: number;
  lastLoss?: number;
};

/** Something that trains: tinkerSft by default; tests pass a fake. */
export type Trainer = (job: TrainJob) => Promise<TrainResult>;

/**
 * Supervised fine-tuning on Tinker with scripts/train_sft.py, run by
 * `python` (default $MACHINE_PYTHON, else python3), which must have
 * tinker-cookbook installed. Needs TINKER_API_KEY. The trainer's own output
 * goes to <logDir>/train.log; progress is reported from its metrics.
 */
export function tinkerSft(options: { python?: string } = {}): Trainer {
  return async (job) => {
    const python = options.python ?? process.env.MACHINE_PYTHON ?? "python3";
    const script = fileURLToPath(new URL("../scripts/train_sft.py", import.meta.url));
    const args = [script, job.data, job.logDir, "--model", job.baseModel, "--epochs", String(job.epochs)];
    if (job.learningRate !== undefined) args.push("--lr", String(job.learningRate));
    if (job.batchSize !== undefined) args.push("--batch-size", String(job.batchSize));
    if (job.effort !== undefined) args.push("--effort", String(job.effort));
    await mkdir(job.logDir, { recursive: true, mode: 0o700 });
    const logFile = path.join(job.logDir, "train.log");
    let reported = -1;
    const progress = setInterval(async () => {
      const steps = (await readFile(path.join(job.logDir, "metrics.jsonl"), "utf8").catch(() => "")).split("\n").filter((l) => l.trim());
      if (steps.length === 0) return;
      const last = JSON.parse(steps.at(-1)!) as { step?: number; train_mean_nll?: number; progress?: number };
      // About ten lines per run: one per tenth of the way through.
      const tenth = Math.floor((last.progress ?? 0) * 10);
      if (tenth > reported) {
        job.log?.(`  step ${(last.step ?? steps.length - 1) + 1}${last.progress ? ` (${Math.round(last.progress * 100)}%)` : ""}: loss ${last.train_mean_nll?.toFixed(3) ?? "?"}`);
        reported = tenth;
      }
    }, 5_000);
    try {
      await new Promise<void>((resolve, reject) => {
        const proc = spawn(python, args, { stdio: ["ignore", "pipe", "pipe"], signal: job.signal });
        let tail = "";
        const onData = (chunk: Buffer) => {
          tail = (tail + chunk.toString()).slice(-4000);
          void appendFile(logFile, chunk, { mode: 0o600 }).catch(() => {});
        };
        proc.stdout.on("data", onData);
        proc.stderr.on("data", onData);
        proc.once("error", (error) => reject(new Error(`Could not start ${python} (set MACHINE_PYTHON to a Python with tinker-cookbook): ${error.message}`)));
        proc.once("exit", (exitCode) => (exitCode === 0 ? resolve() : reject(new Error(`Training exited with code ${exitCode} (full output in ${logFile}): ${tail.trim().slice(-1500)}`))));
      });
    } finally {
      clearInterval(progress);
    }
    return readTrainingLog(job.logDir);
  };
}

/** The final checkpoint and losses of a tinker-cookbook training log directory. */
export async function readTrainingLog(logDir: string): Promise<TrainResult> {
  const jsonl = async (file: string) =>
    (await readFile(path.join(logDir, file), "utf8").catch(() => ""))
      .split("\n")
      .filter((l) => l.trim())
      .map((l) => JSON.parse(l) as Record<string, unknown>);
  const checkpoints = await jsonl("checkpoints.jsonl");
  const final = checkpoints.findLast((c) => typeof c.sampler_path === "string");
  if (!final) throw new Error(`Training finished without a sampler checkpoint in ${logDir}/checkpoints.jsonl.`);
  const metrics = (await jsonl("metrics.jsonl")).filter((m) => typeof m.train_mean_nll === "number");
  const tokens = metrics.reduce((n, m) => n + (typeof m.num_tokens === "number" ? m.num_tokens : 0), 0);
  return {
    checkpoint: final.sampler_path as string,
    steps: metrics.length,
    ...(tokens ? { tokens } : {}),
    ...(metrics.length ? { firstLoss: metrics[0].train_mean_nll as number, lastLoss: metrics.at(-1)!.train_mean_nll as number } : {}),
  };
}

export async function writeLines(file: string, lines: ExportLine[]): Promise<void> {
  await mkdir(path.dirname(file), { recursive: true, mode: 0o700 });
  await writeFile(file, lines.map((l) => JSON.stringify(l)).join("\n") + (lines.length ? "\n" : ""), { mode: 0o600 });
}

/** Named reasoning efforts as Tinker's OpenAI endpoint maps them to Inkling's thinking effort. */
export const EFFORT_PRESETS: Record<string, number> = { none: 0, minimal: 0.1, low: 0.2, medium: 0.7, high: 0.9, xhigh: 0.99 };

/** The thinking effort runs sampled at: a preset or a number; unset means the endpoint's default, high. */
export function effortScalar(effort: string | number | undefined): number {
  if (effort === undefined) return EFFORT_PRESETS.high;
  if (typeof effort === "number") return effort;
  const value = EFFORT_PRESETS[effort];
  if (value === undefined) throw new Error(`Unknown reasoning effort "${effort}".`);
  return value;
}