// A machine is a small, spec'd agent loop with an attached evaluator. The
// spec is plain code: what the machine is for, how it is instructed, which
// model and tools it uses, what good output looks like (the rubric), and how
// each part of that is graded (the evaluator). A child machine specializes a
// parent: it inherits the parent's rubric, graders, instructions, tools, and
// model, and adds its own criteria and instructions on top.

import type { Evaluator } from "./evaluator.ts";
import type { Tool } from "./harness.ts";
import type { Rubric } from "./rubric.ts";
import type { Run } from "./types.ts";
import { existsSync, readFileSync } from "node:fs";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { priceCheckpointAs } from "./pricing.ts";
import { nameSegments } from "./store.ts";

export type MachineSpec<R extends Rubric> = {
  /** Readable address within a workspace. Nested machines use "/", e.g. "chatbot/triage". */
  name: string;
  /** One or two sentences: the job this machine does and for whom. */
  intent: string;
  /** System instructions given to the model on every run. */
  instructions: string;
  /** Default model id; a run may override it. */
  model: string;
  rubric: R;
  evaluator: Evaluator<R>;
  tools?: Tool[];
  /** Upper bound on model calls per run, counting tool round trips. Defaults to the harness's. Children inherit it. */
  maxSteps?: number;
  /**
   * Rewrites the model's reasoning when runs are exported for training, e.g.
   * into another script, so fine-tuning teaches the model to think that way.
   * Children inherit it unless they set their own.
   */
  transformReasoning?: (reasoning: string) => string;
  /**
   * Another machine that plays the other side of a conversation, such as a
   * simulated user. With one, each case is the counterpart's starting point
   * (e.g. a hidden scenario), and running a case means a whole conversation:
   * see converse. Children inherit it unless they set their own.
   */
  counterpart?: Counterpart;
};

/** The other side of a machine's conversations. */
export type Counterpart = {
  machine: Machine;
  /** The conversation ends after the first turn of the machine for which this is true (e.g. it filed its ticket). */
  until?: (run: Run) => boolean;
  /** Most turns of the machine before the conversation stops. Default 8. */
  maxTurns?: number;
  /** The counterpart's first input for a case. Default: the case itself. */
  opening?: (caseInput: string) => string;
};

export type Machine<R extends Rubric = Rubric> = Readonly<
  Required<Omit<MachineSpec<R>, "maxSteps" | "transformReasoning" | "counterpart">> & {
    maxSteps?: number;
    transformReasoning?: (reasoning: string) => string;
    counterpart?: Counterpart;
    /** The machine's own instructions. For a child, `instructions` is its parent's followed by these. */
    ownInstructions: string;
    /** The machine this one specializes, if any. */
    parent?: Machine;
  }
>;

export function defineMachine<R extends Rubric>(spec: MachineSpec<R>): Machine<R> {
  validate(spec);
  return Object.freeze({ ...spec, tools: spec.tools ?? [], ownInstructions: spec.instructions });
}

/** What a child declares: only its own additions. Model defaults to the parent's. */
export type ChildSpec<R extends Rubric> = Omit<MachineSpec<R>, "model"> & { model?: string };

/**
 * A child machine specializes its parent within the parent's constraints. It
 * inherits every rubric criterion with its grader (it may add criteria but
 * not redefine or drop inherited ones), the parent's instructions (its own
 * follow them), the parent's tools (it may add more), and the parent's model
 * unless it names its own. Its name nests under the parent's:
 * `defineChild(julia, { name: "julia/parametric", ... })`.
 */
export function defineChild<P extends Rubric, R extends Rubric>(parent: Machine<P>, spec: ChildSpec<R>): Machine<P & R> {
  if (!spec.name.startsWith(`${parent.name}/`) || spec.name.slice(parent.name.length + 1).includes("/")) {
    throw new Error(`Child machine "${spec.name}" must be named "${parent.name}/<child>".`);
  }
  for (const id of Object.keys(spec.rubric)) {
    if (id in parent.rubric) throw new Error(`Child machine "${spec.name}" redefines criterion "${id}", which it inherits from "${parent.name}".`);
  }
  validate({ ...spec, model: spec.model ?? parent.model });
  const tools = [...parent.tools, ...(spec.tools ?? [])];
  const seen = new Set<string>();
  for (const tool of tools) {
    if (seen.has(tool.name)) throw new Error(`Child machine "${spec.name}" declares tool "${tool.name}", which it already inherits.`);
    seen.add(tool.name);
  }
  return Object.freeze({
    name: spec.name,
    intent: spec.intent,
    instructions: composeInstructions(parent, spec.instructions),
    ownInstructions: spec.instructions,
    model: spec.model ?? parent.model,
    rubric: Object.freeze({ ...parent.rubric, ...spec.rubric }) as P & R,
    evaluator: Object.freeze({ ...parent.evaluator, ...spec.evaluator }) as Evaluator<P & R>,
    tools,
    ...(spec.maxSteps ?? parent.maxSteps ? { maxSteps: spec.maxSteps ?? parent.maxSteps } : {}),
    ...(spec.transformReasoning ?? parent.transformReasoning ? { transformReasoning: spec.transformReasoning ?? parent.transformReasoning } : {}),
    ...(spec.counterpart ?? parent.counterpart ? { counterpart: spec.counterpart ?? parent.counterpart } : {}),
    parent,
  });
}

/** The same machine with different own instructions; a child keeps its parent's in front. */
export function withInstructions<R extends Rubric>(machine: Machine<R>, ownInstructions: string): Machine<R> {
  return Object.freeze({
    ...machine,
    ownInstructions,
    instructions: machine.parent ? composeInstructions(machine.parent, ownInstructions) : ownInstructions,
  });
}

/** The same machine on another model. */
export function withModel<R extends Rubric>(machine: Machine<R>, model: string): Machine<R> {
  if (!model.trim()) throw new Error(`Machine "${machine.name}" needs a non-empty model.`);
  return Object.freeze({ ...machine, model });
}

/** The same machine graded by another rubric and evaluator (e.g. today's, to compare against an old version). */
export function withEvaluation<R extends Rubric>(machine: Machine, rubric: R, evaluator: Evaluator<R>): Machine<R> {
  validate({ ...machine, rubric, evaluator });
  return Object.freeze({ ...machine, rubric, evaluator });
}

/** The machine and its ancestors, root first. */
export function lineage(machine: Machine): Machine[] {
  return machine.parent ? [...lineage(machine.parent), machine] : [machine];
}

function composeInstructions(parent: Machine, own: string): string {
  return `${parent.instructions}\n\n${own}`;
}

function validate<R extends Rubric>(spec: MachineSpec<R>) {
  nameSegments(spec.name, "machine name");
  if (spec.counterpart) {
    const { maxTurns } = spec.counterpart;
    if (maxTurns !== undefined && !(Number.isInteger(maxTurns) && maxTurns > 0)) throw new Error(`Machine "${spec.name}" needs a positive whole counterpart.maxTurns.`);
    if (spec.counterpart.machine.name === spec.name) throw new Error(`Machine "${spec.name}" cannot be its own counterpart.`);
  }
  if (spec.maxSteps !== undefined && !(Number.isInteger(spec.maxSteps) && spec.maxSteps > 0)) throw new Error(`Machine "${spec.name}" needs a positive whole maxSteps.`);
  for (const field of ["intent", "instructions", "model"] as const) {
    if (!spec[field]?.trim()) throw new Error(`Machine "${spec.name}" needs a non-empty ${field}.`);
  }
  for (const id of Object.keys(spec.rubric)) {
    if (!spec.evaluator[id]) throw new Error(`Machine "${spec.name}" has no grader for criterion "${id}".`);
  }
  for (const id of Object.keys(spec.evaluator)) {
    if (!(id in spec.rubric)) throw new Error(`Machine "${spec.name}" grades "${id}", which its rubric does not declare.`);
  }
  const names = new Set<string>();
  for (const tool of spec.tools ?? []) {
    if (names.has(tool.name)) throw new Error(`Machine "${spec.name}" declares tool "${tool.name}" twice.`);
    names.add(tool.name);
  }
}

/** A machine whose instructions live in this file, beside its machine.ts, declares them an optimization target. */
export const INSTRUCTIONS_FILE = "instructions.md";

/** A machine whose model comes from modelFrom keeps its committed model in this file, beside its machine.ts. */
export const MODEL_FILE = "model.json";

/** What model.json holds: the live model and, for a tinker:// checkpoint, the base model it was trained from. */
export type ModelFile = { model: string; baseModel?: string };

const MODEL_FILE_DIRS = new Set<string>();

/**
 * The machine's model: the one committed to model.json beside the calling
 * module (by `machine commit` or `m.update({ model })`), else `fallback`.
 * `model: modelFrom(import.meta.url, "Qwen/Qwen3.6-35B-A3B")`. A committed
 * checkpoint is registered with its base model for pricing.
 */
export function modelFrom(moduleUrl: string, fallback: string): string {
  const dir = path.dirname(fileURLToPath(moduleUrl));
  MODEL_FILE_DIRS.add(dir);
  const file = path.join(dir, MODEL_FILE);
  if (!existsSync(file)) return fallback;
  const parsed = JSON.parse(readFileSync(file, "utf8")) as Partial<ModelFile>;
  if (typeof parsed.model !== "string" || !parsed.model.trim()) throw new Error(`${file} needs a "model" string.`);
  if (parsed.baseModel) priceCheckpointAs(parsed.model, parsed.baseModel);
  return parsed.model;
}

/** Whether the machine defined in `dir` takes its model from model.json (it called modelFrom). */
export function usesModelFile(dir: string): boolean {
  return MODEL_FILE_DIRS.has(path.resolve(dir));
}

/** Read instructions.md from the directory of the calling module: `instructionsFrom(import.meta.url)`. */
export function instructionsFrom(moduleUrl: string): string {
  const file = fileURLToPath(new URL(INSTRUCTIONS_FILE, moduleUrl));
  const text = readFileSync(file, "utf8").trim();
  if (!text) throw new Error(`${file} is empty.`);
  return text;
}