/**
 * One running simulation: grid, medium, compiled .m, timestep.
 *
 * Everything that is not rendering. The app and the tests both go through
 * this, so there is one place that decides how a pair of .m files becomes
 * something running on the GPU — and nothing about it is browser-specific
 * beyond needing a GPUDevice.
 */
import { makeGrid, stableDt, type Grid } from '../grid.ts';
import { StencilPlan } from './stencil.ts';
import { GpuModel, type ModelParams } from './model.ts';
import { Scene } from '../scene/scene.ts';
import { Recorder } from '../audio/recorder.ts';
import type { MModel, Params } from './registry.ts';
import type { MScene } from '../scene/registry.ts';

export interface ModelSessionOptions {
  device: GPUDevice;
  model: MModel;
  params: Params;
  /** Override the model source — the editor's working copy. */
  source?: string;
  scene: MScene;
  sceneParams: Params;
  /** Override the scene source — the editor's working copy. */
  sceneSource?: string;
  /** Grid points per side. */
  n: number;
  /** Side length of the square domain, in metres. Defaults to a small toy
   *  domain (2 m); the app passes its real DOMAIN constant explicitly. */
  L?: number;
  /** Fraction of the stability limit to take as the timestep. */
  cfl?: number;
  /** Grid fields one kernel may read, overriding what the device allows.
   *  Only for tests, which use it to exercise the planner's kernel splitting
   *  on a device that would never need it. */
  operandBudget?: number;
}

export class ModelSession {
  readonly device: GPUDevice;
  readonly model: MModel;
  readonly grid: Grid;
  readonly gpu: GpuModel;
  /** The microphone: the pressure at one point, every timestep. */
  readonly recorder: Recorder;

  /** Model time and step count since the last reset. */
  t = 0;
  steps = 0;

  #scene: Scene;
  #sceneModel: MScene;
  #params: Params;
  #cfl: number;

  private constructor(init: {
    device: GPUDevice;
    model: MModel;
    grid: Grid;
    gpu: GpuModel;
    recorder: Recorder;
    scene: Scene;
    sceneModel: MScene;
    params: Params;
    cfl: number;
  }) {
    this.device = init.device;
    this.model = init.model;
    this.grid = init.grid;
    this.gpu = init.gpu;
    this.recorder = init.recorder;
    this.#scene = init.scene;
    this.#sceneModel = init.sceneModel;
    this.#params = init.params;
    this.#cfl = init.cfl;
  }

  static async create(opts: ModelSessionOptions): Promise<ModelSession> {
    const { device, model, params, scene, sceneParams } = opts;
    const grid = makeGrid(opts.n, opts.L ?? 2);
    const cfl = opts.cfl ?? 0.5;

    // The scene first: the model's timestep depends on the fastest speed in
    // it, and the medium is an argument the compiled step reads.
    const built = Scene.create({
      grid,
      source: opts.sceneSource ?? scene.source,
      paramNames: scene.params.map((p) => p.key),
      params: sceneParams,
    });

    const stencil = StencilPlan.create(device, { nx: grid.nx, ny: grid.ny, h: grid.h });
    const gpu = await GpuModel.create({
      device,
      stencil,
      grid,
      medium: built,
      source: opts.source ?? model.source,
      paramNames: model.params.map((p) => p.key),
      state: model.state,
      operandBudget: opts.operandBudget,
    });

    // The microphone listens to the host-owned pressure buffer, which is
    // where both `init` and `step` leave their result.
    const recorder = await Recorder.create({
      device,
      field: gpu.stateBuffer(model.state[0])!,
      nx: grid.nx,
      ny: grid.ny,
    });

    const session = new ModelSession({
      device, model, grid, gpu, recorder, scene: built, sceneModel: scene, params, cfl,
    });
    session.#applyDt();
    gpu.setParams(params);
    return session;
  }

  get scene(): Scene {
    return this.#scene;
  }

  get sceneModel(): MScene {
    return this.#sceneModel;
  }

  /** The pressure field's name — the first state field the model advances. */
  get pressureName(): string {
    return this.model.state[0];
  }

  get dt(): number {
    return this.gpu.dt;
  }

  /** Fraction of the stability limit the timestep is taken at. */
  get cfl(): number {
    return this.#cfl;
  }

  /**
   * Change the timestep, as a fraction of what the scheme is stable at.
   *
   * A fraction rather than a number of seconds, because the stable step
   * follows from the grid spacing and the fastest speed in the scene: refine
   * the grid or drop in a faster scatterer and a dt that was fine becomes a
   * dt that diverges. This way the meaning of the setting survives both.
   *
   * Applied to the running simulation as it stands. Leapfrog carries two
   * fields a timestep apart, so changing dt between steps is inconsistent by
   * the size of the change; a small drag is a small transient, and a large
   * jump is worth a Restart.
   */
  setCfl(cfl: number): void {
    this.#cfl = cfl;
    this.#applyDt();
    this.gpu.setParams(this.#params);
    // The trace is one sample per timestep, so a recording made at one dt
    // cannot be spliced onto one made at another: it would be two different
    // sample rates in the same buffer.
    this.recorder.clear();
  }

  /** The largest timestep this scheme is stable at on this grid and medium. */
  get dtLimit(): number {
    return stableDt(this.grid.h, this.#scene.cmax, this.model.order, 1);
  }

  #applyDt(): void {
    this.gpu.setDt(stableDt(this.grid.h, this.#scene.cmax, this.model.order, this.#cfl));
  }

  /**
   * Swap the medium under a running model. It is data, not code, so this
   * needs no recompile — but the timestep follows from it, and changing dt
   * mid-run would leave the leapfrog's two histories half a step apart, so
   * the caller is expected to reset afterwards.
   */
  setScene(sceneModel: MScene, params: Params, source?: string): void {
    const built = Scene.create({
      grid: this.grid,
      source: source ?? sceneModel.source,
      paramNames: sceneModel.params.map((p) => p.key),
      params,
    });
    this.#scene = built;
    this.#sceneModel = sceneModel;
    this.gpu.uploadMedium(built);
    this.#applyDt();
    this.gpu.setParams(this.#params);
    this.recorder.clear();
  }

  setParams(params: ModelParams): void {
    this.#params = params;
    this.gpu.setParams(params);
  }

  /** Run `init`: a silent grid at t = 0. */
  reset(): void {
    this.gpu.init();
    this.recorder.clear();
    this.t = 0;
    this.steps = 0;
  }

  /** Put the microphone at the grid point nearest the given coordinates. */
  setMic(x: number, y: number): void {
    const { L, h } = this.grid;
    this.recorder.setProbe((x + L / 2) / h - 0.5, (y + L / 2) / h - 0.5);
  }

  /** Advance `n` steps. Synchronous: records and submits, nothing read back.
   *  The microphone samples inside the same submission, once per step. */
  step(n = 1): void {
    this.gpu.step(n, (enc) => this.recorder.encode(enc));
    this.t += n * this.dt;
    this.steps += n;
  }

  /** Wait for the submitted steps to finish, without reading anything back. */
  sync(): Promise<undefined> {
    return this.device.queue.onSubmittedWorkDone();
  }

  /** Read a named field back to the CPU. */
  read(name: string): Promise<Float32Array> {
    return this.gpu.read(name);
  }

  describe(): { init: string[]; step: string[] } {
    return this.gpu.describe();
  }

  destroy(): void {
    this.recorder.destroy();
    this.gpu.destroy();
  }
}