import type { ProblemSpec } from "../spec/types";
import type { MethodDefinition } from "../methods/types";
import type { GpuStages } from "../runner/types";
import { generateParamsBlock } from "./params";
import { LIBRARY_FILES, type SourceFile } from "./library";

/** A complete, self-contained set of `.m` files plus the one to run. */
export interface SourceBundle {
  files: SourceFile[];
  /** File executed as the script */
  entry: string;
  /** Workspace variable read back after the run */
  resultVariable: string;
}

export const MAIN_FILE = "main.m";

/**
 * The script a method runs: the params block followed by the solver call. With
 * `gpu`, for a method whose loop has a WebGPU port, the call is written out as
 * its three stages, since that is what the host then runs (see buildGpuStages).
 */
export function generateMain(
  spec: ProblemSpec,
  method: MethodDefinition,
  opts: { gpu?: boolean } = {},
): string {
  const g = opts.gpu ? method.gpu : undefined;
  const solve = g
    ? [
        `% Solve.  With the GPU option the host runs ${g.setupFunction} and ${g.outputFunction} here in numbl and`,
        `% ${g.loopFile} as its WebGPU port; without WebGPU, or when this script or ${g.loopFile} is edited,`,
        `% all three run in numbl.  The library files are in the other tabs.`,
        `S = ${g.setupFunction}(params);`,
        `R = ${g.loopFile.replace(/\.m$/, "")}(S);      % replaced by the WebGPU port when the loop runs on the GPU`,
        `result = ${g.outputFunction}(S, R);`,
      ]
    : [
        `% Solve.  ${method.solverFunction} and the functions it calls are the library files in the other tabs.`,
        `result = ${method.solverFunction}(params);`,
      ];
  return [generateParamsBlock(spec), "", ...solve, ""].join("\n");
}

/** Bundle for running one method on a spec. */
export function buildRunBundle(
  spec: ProblemSpec,
  method: MethodDefinition,
  opts: { gpu?: boolean } = {},
): SourceBundle {
  return {
    files: [{ name: MAIN_FILE, source: generateMain(spec, method, opts) }, ...LIBRARY_FILES],
    entry: MAIN_FILE,
    resultVariable: "result",
  };
}

/**
 * The two numbl scripts of a GPU run (see runner/executeGpu.ts): the setup,
 * whose struct S the host reads (sparse operators as find() triplets, since
 * a sparse matrix does not cross to the host), and the output, which takes the
 * loop's record as the workspace variable R. Null for a method without a port.
 */
export function buildGpuStages(spec: ProblemSpec, method: MethodDefinition): GpuStages | null {
  const g = method.gpu;
  if (!g) return null;
  const params = generateParamsBlock(spec);
  const setup = [
    params,
    "",
    `% Stage 1 of a GPU run: the setup, read by the WebGPU port of ${g.loopFile}.`,
    `S = ${g.setupFunction}(params);`,
    "if S.sparse",
    "    % the sparse operators cross to the host as triplets",
    "    [ki, kj, kv] = find(S.Km); S.Km_i = ki; S.Km_j = kj; S.Km_v = kv;",
    "    [li, lj, lv] = find(S.Lm); S.Lm_i = li; S.Lm_j = lj; S.Lm_v = lv;",
    "end",
    "",
  ].join("\n");
  const output = [
    params,
    "",
    `% Stage 3 of a GPU run: R is the record of the WebGPU loop, passed in by the host.`,
    `S = ${g.setupFunction}(params);`,
    `result = ${g.outputFunction}(S, R);`,
    "",
  ].join("\n");
  return { setup, output, loopFile: g.loopFile };
}

/** Bundle that evaluates a method's cost plan (currently only the FD method has one). */
export function buildPlanBundle(spec: ProblemSpec, method: MethodDefinition): SourceBundle | null {
  if (!method.planFunction) return null;
  const main = [generateParamsBlock(spec), "", `plan = ${method.planFunction}(params);`, ""].join(
    "\n",
  );
  return {
    files: [{ name: "plan.m", source: main }, ...LIBRARY_FILES],
    entry: "plan.m",
    resultVariable: "plan",
  };
}

/** Bundle comparing two results, which the runner passes in as workspace variables A and B. */
export function buildCompareBundle(): SourceBundle {
  const main = [
    "% A and B are two result structs for the same problem (A: modal, B: finite differences).",
    "cmp = plate_compare(A, B);",
    "",
  ].join("\n");
  return {
    files: [{ name: "compare.m", source: main }, ...LIBRARY_FILES],
    entry: "compare.m",
    resultVariable: "cmp",
  };
}

/**
 * Replace files in a bundle with edited versions, keeping everything else. Used
 * when the user edits source in the code view: the run then uses exactly the
 * text on screen.
 */
export function withEdits(bundle: SourceBundle, edits: Record<string, string>): SourceBundle {
  return {
    ...bundle,
    files: bundle.files.map((f) =>
      edits[f.name] !== undefined ? { ...f, source: edits[f.name] } : f,
    ),
  };
}