// The page's side of the engine: a Web Worker (engine_worker.js) that runs GENCO (ONNX, onnxruntime-web) and the
// repair and fast correction (Pyodide). The contract is in ENGINE_CONTRACT.md; this file is its executable form.
//
// Knobs (JSON, the fields of whatif_app's whatif.scenario.Knobs dataclass):
//   domes         [[lat, lon, radius_km, amplitude_degC], ...]   0-3 heat domes (the UI offers 3)
//   weather       [[lat, lon, dT_degC], ...]   a real-weather anomaly field (scenario.field_anomaly), [] for none
//   heat_load_wind  false: r and wind are real data that already include the heat (real days); default true
//   r             load level, share of the summer peak            UI 0.5-1.1 (training 0.6-1.0)
//   wind_area     8 numbers, wind availability per area           0-1, the case's area order (grid.json meta.areas)
//   solar_area    8 numbers, solar availability per area          0-1
//   gas_m, coal_m fuel price multipliers                          UI 0.25-5 and 0.25-2
//   branches_out  branch indices (0-based, case order)            the network must stay connected (the page checks)
//   gens_out      unit indices (0-based, case order)
//   noise         null (always: the browser never draws the generator's per-bus load noise)
//   dark          [] (bus indices of a hurricane's damaged zone; the heat tab never sets it)
//
// Results (one per call; typed arrays are transferred, not copied):
//   preview(knobs) -> {cost, max_loading, n_viol, viol_detail, feasible, time_ms,
//                      line_loading: Float32Array(n_branch), bus_vm: Float32Array(n_bus), gen_p: Float32Array(n_gen),
//                      v_out?: Int32Array}                       (v_out optional: buses out of band at the gate)
//   check(knobs)   -> the same, plus loops; or {status: "cancelled", loops: 0} when a newer request stopped it
//   exact(knobs, {start: "check" | "genco" | "flat" | "mid", limit_s})
//                  -> {status: "optimal" | "not_certified" | "cancelled", start, cost, iterations, time_ms (solve only),
//                      build_ms, package_ms, total_ms, feasible, gap_vs_check, line_loading, bus_vm, gen_p}
//                      (the exact AC-OPF, interior point from the checked point (the check run first if needed;
//                      its time in check_ms), GENCO's prediction, or a standard start without GENCO; hard limit 20 s
//                      from the pipeline, 60 s from scratch; arrays may be NaN when not certified)
//   cancel()       -> stops a running exact solve, which then resolves with {status: "cancelled"}; a preview, check or
//                     sent while an exact is outstanding stops it the same way (Cancellation in the contract)
//   preview results may also carry forward_ms, repair_ms, raw_max_mismatch_mw, raw_cost, raw_line_loading (GENCO's
//   own prediction before the repair)
//   init(opts)     -> {ready, load_ms, threads, can_cancel, interrupt_buffer (SharedArrayBuffer of 2 Int32, or null)}
//   progress       <- {type: "progress", id, stage: "check" | "solve" | "setup"}: the request's timed work starts (check,
//                     exact's interior point) or its one-time setup does; passed to check()'s or exact()'s onProgress
//   every result also gets rt_ms here: its round trip as the page saw it
// Units: cost $/h; loading as a fraction of the (heat-adjusted) rating, NaN for a branch out of service or without a
// rating; |V| p.u.; Pg MW (0 for units out); time_ms the engine's own time for the call.

export const WORKER_URL = "engine_worker.js";
export const WORKER_TYPE = "module";               // the worker is an ES module (import onnxruntime-web / pyodide.mjs)
// a private Space signs the page (?__sign=): the worker's requests pass only if the worker's own URL carries the signature
const SIGN = typeof location === "undefined" ? null : new URLSearchParams(location.search).get("__sign");
const signed = (u) => (SIGN ? `${u}${String(u).includes("?") ? "&" : "?"}__sign=${encodeURIComponent(SIGN)}` : u);

export function normalizeKnobs(k) {
  const n = (x, d) => (Number.isFinite(+x) ? +x : d);
  return {
    domes: (k.domes || []).map((d) => d.map(Number)),
    r: n(k.r, 0.9),
    wind_area: (k.wind_area || Array(8).fill(0.5)).map(Number),
    solar_area: (k.solar_area || Array(8).fill(0.5)).map(Number),
    gas_m: n(k.gas_m, 1.0),
    coal_m: n(k.coal_m, 1.0),
    gens_out: [...new Set((k.gens_out || []).map(Number))].sort((a, b) => a - b),
    branches_out: [...new Set((k.branches_out || []).map(Number))].sort((a, b) => a - b),
    noise: Array.isArray(k.noise) && k.noise.length === 2 ? k.noise.map(Number) : null,   // a stored hurricane draw: [seed, i]
    dark: (k.dark || []).map(Number),
    weather: (k.weather || []).map((p) => p.map(Number)),
    heat_load_wind: k.heat_load_wind !== false,
  };
}

// Throws if a result breaks the contract (used by the page in development and by tools/ui_test.py through the mock).
export function validateResult(r, dims, withLoops = false) {
  const need = ["cost", "max_loading", "n_viol", "viol_detail", "feasible", "time_ms"];
  for (const f of need) if (!(f in r)) throw new Error(`engine result lacks ${f}`);
  const arr = [["line_loading", dims.nBranch], ["bus_vm", dims.nBus], ["gen_p", dims.nGen]];
  for (const [f, len] of arr) {
    if (!(r[f] instanceof Float32Array)) throw new Error(`${f} must be a Float32Array`);
    if (r[f].length !== len) throw new Error(`${f} has ${r[f].length} entries, expected ${len}`);
  }
  if (withLoops && !Number.isInteger(r.loops)) throw new Error("check result lacks integer loops");
  return r;
}

// exact() times (ENGINE_CONTRACT.md): time_ms is the solve only; build_ms the model build and package_ms the package load
// on this call (0 afterwards); total_ms the call's wall time. An engine that reports timing: {solve_s, build_s} instead
// is read the same way. time_ms is never adjusted by the build.
export function normalizeExact(r) {
  if (!r || typeof r !== "object") return r;
  const t = r.timing || {};
  if (!Number.isFinite(r.build_ms) && Number.isFinite(t.build_s)) r.build_ms = t.build_s * 1000;
  if (!Number.isFinite(r.time_ms)) {
    if (Number.isFinite(t.solve_s)) r.time_ms = t.solve_s * 1000;
    else if (Number.isFinite(r.total_ms)) r.time_ms = r.total_ms;
  }
  return r;
}

// An exact() result: a status always; when optimal, the cost and the three arrays as for a preview.
export function validateExact(r, dims) {
  if (!r || !["optimal", "not_certified", "cancelled"].includes(r.status)) throw new Error("exact result lacks a valid status");
  if (r.status === "optimal") {
    for (const f of ["cost", "time_ms"]) if (!Number.isFinite(r[f])) throw new Error(`exact result lacks ${f}`);
    for (const [f, len] of [["line_loading", dims.nBranch], ["bus_vm", dims.nBus], ["gen_p", dims.nGen]]) {
      if (!(r[f] instanceof Float32Array) || r[f].length !== len) throw new Error(`exact ${f} must be a Float32Array(${len})`);
    }
  }
  return r;
}

export class Engine {
  constructor(url = WORKER_URL, type = WORKER_TYPE) {
    this.worker = new Worker(signed(url), { type });
    this.n = 0;
    this.pending = new Map();
    this.statusListeners = [];
    this.ready = false;
    this.iv = null;                                 // the interrupt buffer: [cancel flag, cancellation counter]
    this.slowOut = 0;                               // exact and check calls sent and not yet answered
    this.worker.onmessage = (e) => this._onMessage(e.data);
    this.worker.onerror = (e) => {
      const err = new Error(`engine worker error: ${e.message || e}`);
      for (const [, p] of this.pending) p.reject(err);
      this.pending.clear();
      this.statusListeners.forEach((cb) => cb({ type: "status", stage: "error", detail: err.message }));
    };
  }

  onStatus(cb) { this.statusListeners.push(cb); }

  _onMessage(m) {
    if (m.type === "status") { this.statusListeners.forEach((cb) => cb(m)); return; }
    const p = this.pending.get(m.id);
    if (!p) return;
    if (m.type === "progress") { if (p.onProgress) p.onProgress(m); return; }   // {type: "progress", id, stage}
    this.pending.delete(m.id);
    // rt_ms: this request's round trip as the page saw it (queue, setup, the work, messaging); the result's own times
    // say which part was the work
    if (m.ok && m.result && typeof m.result === "object") m.result.rt_ms = performance.now() - p.t0;
    if (m.ok) p.resolve(m.result); else p.reject(new Error(m.error || "engine error"));
  }

  _call(type, payload = {}, onProgress = null) {
    const id = ++this.n;
    return new Promise((resolve, reject) => {
      this.pending.set(id, { resolve, reject, onProgress, t0: performance.now() });
      this.worker.postMessage({ id, type, ...payload });
    });
  }

  async init(opts = {}) {
    const r = await this._call("init", { opts });
    this.ready = !!r.ready;
    this.iv = r.interrupt_buffer ? new Int32Array(r.interrupt_buffer) : null;
    return r;
  }
  // A newer request stops outstanding slow work (an exact; a check too where the engine supports it): bump the counter
  // and raise the flag before posting it. A running solve polls the buffer at safe points; queued work whose gen is
  // behind the counter is skipped. Writing the buffer when nothing reads it is harmless (ENGINE_CONTRACT.md).
  _interrupt() { if (this.iv && this.slowOut > 0) { Atomics.add(this.iv, 1, 1); Atomics.store(this.iv, 0, 2); } }
  _slow(type, payload, onProgress) {
    const gen = this.iv ? Atomics.load(this.iv, 1) : 0;
    this.slowOut += 1;
    const done = () => { this.slowOut -= 1; };
    const p = this._call(type, { ...payload, gen }, onProgress);
    p.then(done, done);
    return p;
  }
  preview(knobs) { this._interrupt(); return this._call("preview", { knobs: normalizeKnobs(knobs) }); }
  // onProgress(m): called with {stage: "check" | "solve" | "setup"} when the work the result times starts
  check(knobs, { onProgress = null } = {}) { this._interrupt(); return this._slow("check", { knobs: normalizeKnobs(knobs) }, onProgress); }
  exact(knobs, opts = { start: "genco" }) {
    const start = ["check", "genco", "flat", "mid"].includes(opts.start) ? opts.start : "genco";
    const lim = Number.isFinite(opts.limit_s) ? { limit_s: opts.limit_s } : {};      // the solve's hard limit (engine default if absent)
    return this._slow("exact", { knobs: normalizeKnobs(knobs), opts: { start, ...lim } }, opts.onProgress || null).then(normalizeExact);
  }
  // Stop a running exact solve (in practice the from-scratch one): it then resolves with {status: "cancelled", start}.
  cancel() { this._interrupt(); return this._call("cancel", {}); }
}

// Transferable buffers of a result (the worker passes these to postMessage).
export function transferables(result) {
  const out = [];
  const walk = (o) => {
    if (!o || typeof o !== "object") return;
    if (ArrayBuffer.isView(o)) { out.push(o.buffer); return; }
    for (const v of Object.values(o)) walk(v);
  };
  walk(result);
  return [...new Set(out)];
}