/**
 * One running scene: the air grid, every object's grid, the generated (or
 * hand-edited) .m files compiled to GPU plans, the timestep, the microphone.
 *
 * Everything that is not rendering. The app and the tests both go through
 * this, so there is one place that decides how a SceneSpec becomes something
 * running on the GPU — and nothing about it is browser-specific beyond
 * needing a GPUDevice.
 *
 * The order of construction matters and is worth spelling out. The air grid
 * fixes h; with the speed of air that fixes a provisional timestep (the CFL
 * condition); the generator then sizes every object's grid at that dt and
 * instantiates the ops. The scene .m is evaluated next — normally it keeps
 * the medium at c0 and the provisional dt stands, but an edited scene that
 * puts a faster medium somewhere lowers dt, and the object grids are resized
 * once at the final value. Objects have no say in the timestep; each chooses
 * the finest node spacing that is stable at the air's dt for anything its
 * sliders can reach (see src/grids.ts), which is why moving a slider never
 * recompiles.
 */
import { makeAirGrid, stableDt, type AirGrid, type ObjectGrid } from '../grids.ts';
import { OpPlan } from '../mgpu/ops.ts';
import { GpuModel, type ModelParams } from '../mgpu/model.ts';
import { Scene } from '../scene/scene.ts';
import { Recorder } from '../audio/recorder.ts';
import { C_AIR, CFL, DOMAIN_X } from '../units.ts';
import { generate, suffixedParams, type GeneratedScene } from './codegen.ts';
import { templateOf } from './objects.ts';
import { BODY_TEMPLATES } from './bodies.ts';
import { defaultsOf, type Params, type SceneSpec } from './types.ts';

export interface StudioSessionOptions {
  device: GPUDevice;
  spec: SceneSpec;
  /** Model parameter values by suffixed key; defaults fill the gaps. */
  params?: Params;
  /** Scene parameter values (placement + body) by key; defaults fill gaps. */
  sceneParams?: Params;
  /** Override the generated sources — the Lab editor's working copies. */
  modelSource?: string;
  sceneSource?: string;
  /** Air grid points along x; y and z get half each. */
  nx: number;
  /** Domain length along x, metres. */
  Lx?: number;
  /** Grid fields one kernel may read, overriding what the device allows.
   *  Only for tests. */
  operandBudget?: number;
}

/** Model + scene parameter defaults for a spec, suffixed, spec values kept. */
export function specParams(spec: SceneSpec): { params: Params; sceneParams: Params } {
  const params: Params = {};
  const sceneParams: Params = {};
  for (const o of spec.objects) {
    const t = templateOf(o);
    Object.assign(params, suffixedParams(t.paramSpecs, o.params, o.id));
    Object.assign(sceneParams, suffixedParams(t.placeSpecs, o.place, o.id));
  }
  Object.assign(
    sceneParams,
    defaultsOf(BODY_TEMPLATES[spec.body.type].paramSpecs),
    spec.body.params,
  );
  return { params, sceneParams };
}

export class StudioSession {
  readonly device: GPUDevice;
  readonly spec: SceneSpec;
  readonly air: AirGrid;
  readonly gpu: GpuModel;
  readonly generated: GeneratedScene;
  /** The sources actually compiled (generated, or the editor's overrides). */
  readonly modelSource: string;
  readonly sceneSource: string;
  /** The microphone: the pressure at one point, every timestep. */
  readonly recorder: Recorder;

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

  #scene: Scene;
  #params: Params;
  #sceneParams: Params;
  #dt: number;

  private constructor(init: {
    device: GPUDevice;
    spec: SceneSpec;
    air: AirGrid;
    gpu: GpuModel;
    generated: GeneratedScene;
    modelSource: string;
    sceneSource: string;
    recorder: Recorder;
    scene: Scene;
    params: Params;
    sceneParams: Params;
    dt: number;
  }) {
    this.device = init.device;
    this.spec = init.spec;
    this.air = init.air;
    this.gpu = init.gpu;
    this.generated = init.generated;
    this.modelSource = init.modelSource;
    this.sceneSource = init.sceneSource;
    this.recorder = init.recorder;
    this.#scene = init.scene;
    this.#params = init.params;
    this.#sceneParams = init.sceneParams;
    this.#dt = init.dt;
  }

  static async create(opts: StudioSessionOptions): Promise<StudioSession> {
    const { device, spec } = opts;
    const air = makeAirGrid(opts.nx, opts.Lx ?? DOMAIN_X);

    const defaults = specParams(spec);
    const params = { ...defaults.params, ...opts.params };
    const sceneParams = { ...defaults.sceneParams, ...opts.sceneParams };

    // A provisional timestep from the speed of air; the scene is evaluated
    // against it, and only an edited scene with a faster medium changes it.
    let dt = stableDt(air.h, C_AIR, CFL);
    let generated = generate(spec, air, dt);
    const sceneSource = opts.sceneSource ?? generated.sceneSource;

    const buildScene = (): Scene =>
      Scene.create({
        air,
        source: sceneSource,
        outputs: generated.sceneOutputs,
        consts: generated.sceneConsts,
        paramNames: generated.sceneParamNames,
        params: sceneParams,
      });
    let scene = buildScene();
    const dtFinal = stableDt(air.h, Math.max(scene.cmax, C_AIR), CFL);
    if (dtFinal < dt) {
      // A faster medium tightens the CFL bound; resize the object grids once
      // at the final timestep. The scene's fields do not depend on dt.
      dt = dtFinal;
      generated = generate(spec, air, dt);
    }

    const modelSource = opts.modelSource ?? generated.modelSource;
    const ops = new OpPlan(device, generated.ops);

    // The scene's evaluated fields replace the generator's placeholders.
    const fields = { ...generated.modelFields, ...scene.fields };

    const gpu = await GpuModel.create({
      device,
      ops,
      consts: generated.modelConsts,
      fields,
      source: modelSource,
      paramNames: generated.modelParamNames,
      state: generated.state,
      operandBudget: opts.operandBudget,
    });
    gpu.setDt(dt);

    // 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('p')!,
      nx: air.nx,
      ny: air.ny,
      nz: air.nz,
    });

    const session = new StudioSession({
      device, spec, air, gpu, generated, modelSource, sceneSource,
      recorder, scene, params, sceneParams, dt,
    });
    gpu.setParams(params);
    return session;
  }

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

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

  get params(): Params {
    return this.#params;
  }

  get sceneParams(): Params {
    return this.#sceneParams;
  }

  /** Each object's grid, by id. */
  get grids(): Record<string, ObjectGrid> {
    return this.generated.grids;
  }

  /**
   * Re-evaluate the scene under a running model — a body edit, an object
   * moved. It is data, not code, so this needs no recompile. The timestep is
   * deliberately left alone: it was set from max(cmax, c_air) at build time,
   * and a scene edit that *raises* the fastest speed above that needs a
   * rebuild anyway (the caller compares `dtWanted` and rebuilds when they
   * disagree).
   */
  setScene(sceneParams: Params, source?: string): void {
    const scene = Scene.create({
      air: this.air,
      source: source ?? this.sceneSource,
      outputs: this.generated.sceneOutputs,
      consts: this.generated.sceneConsts,
      paramNames: this.generated.sceneParamNames,
      params: sceneParams,
    });
    this.#scene = scene;
    this.#sceneParams = sceneParams;
    this.gpu.uploadFields(scene.fields);
  }

  /** The timestep the current scene would ask for. Differs from `dt` only
   *  when a scene edit changed the fastest speed, which calls for a rebuild. */
  get dtWanted(): number {
    return stableDt(this.air.h, Math.max(this.#scene.cmax, C_AIR), CFL);
  }

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

  /** Run `init`: every object drawn into its excitation, silent air, t = 0. */
  excite(): void {
    this.gpu.init();
    this.recorder.clear();
    this.t = 0;
    this.steps = 0;
  }

  /** Put the microphone at the grid point nearest (x, y, z), metres. */
  setMic(x: number, y: number, z: number): void {
    const { Lx, Ly, Lz, h } = this.air;
    this.recorder.setProbe(
      (x + Lx / 2) / h - 0.5,
      (y + Ly / 2) / h - 0.5,
      (z + Lz / 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();
  }
}