// The browser engine's orchestration: Pyodide (the app's scenario builder, repair and correction, copied verbatim)
// and onnxruntime-web (GENCO, the dense-aggregation ONNX export), in one thread of one Web Worker.
// Shared by whatif_static/engine_worker.js (browser) and engine/tools/node_harness.mjs (Node), which inject the
// loaders: createEngine({ base, loadPyodide, ort, fetchBytes, threads, onProgress, warmup }).
//
// Per what-if: Python prepare() -> GENCO in ORT -> Python preview() (Pipeline.preview: the "dsq" repair) and, for a
// check, Python check() (Pipeline.check: the fast correction, settings.SQP). Results: scalars plus typed arrays
// (line_loading, bus_vm, bus_va, gen_p, v_out, genco_vm, genco_pg), with time_ms the wall time of the call.

const now = () => performance.now();

// whatif_app's DEFAULT what-if (app.py), used for the warm-up preview
export const DEFAULT_KNOBS = { domes: [[32.8, -97.05, 300.0, 5.0]], r: 0.9, wind_area: Array(8).fill(0.4),
  solar_area: Array(8).fill(0.8), gas_m: 1.0, coal_m: 1.0, gens_out: [], branches_out: [], noise: null, dark: [] };

export async function createEngine(opt) {
  const { base, loadPyodide, ort, fetchBytes } = opt;
  const progress = opt.onProgress ?? (() => {});
  const stages = {};
  const t0 = now();
  let t = t0;
  const stage = (name) => { const n = now(); stages[name] = n - t; t = n; progress({ stage: name, ms: stages[name] }); };

  // --- Python: Pyodide, numpy/scipy/pandas, the engine package, the base case
  progress({ stage: "start" });
  const fetched = {};
  const timedFetch = async (rel) => { const a = now(); const b = await fetchBytes(base + rel); fetched[rel] = { bytes: b.byteLength, ms: now() - a }; return b; };
  const filesP = Promise.all(["engine_py.zip", "data/case.npz", "data/case.json", "model/model.json"].map(timedFetch));
  // GENCO's session is created while Pyodide loads (onnxruntime-web's WASM compiles in the background)
  ort.env.wasm.numThreads = opt.threads ?? 1;
  const sessP = timedFetch("model/genco_small_texas.onnx").then((onnx) => ort.InferenceSession.create(onnx,
    { executionProviders: ["wasm"], graphOptimizationLevel: "all" }));
  sessP.catch(() => {});
  const py = await loadPyodide({ indexURL: base + "pyodide/", lockFileURL: base + "pyodide/pyodide-lock.json",
    stdout: opt.stdout ?? (() => {}), stderr: opt.stderr ?? ((s) => console.warn(s)) });
  stage("pyodide");
  await py.loadPackage(["numpy", "scipy", "pandas"], { messageCallback: () => {}, errorCallback: (s) => console.warn(s) });
  stage("packages");
  const [zip, npz, caseJson, modelJson] = await filesP;
  stage("fetch_rest");
  py.unpackArchive(zip, "zip", { extractDir: "/home/pyodide/engine" });
  py.FS.mkdirTree("/home/pyodide/data");
  py.FS.writeFile("/home/pyodide/data/case.npz", npz);
  py.FS.writeFile("/home/pyodide/data/case.json", caseJson);
  py.FS.writeFile("/home/pyodide/data/model.json", modelJson);
  py.runPython(`
import sys
sys.path.insert(0, "/home/pyodide/engine")
from webengine.core import Engine
ENGINE = Engine("/home/pyodide/data/case.npz", "/home/pyodide/data/case.json", "/home/pyodide/data/model.json")
`);
  const E = py.globals.get("ENGINE");
  const info = JSON.parse(E.info());
  stage("engine");

  // --- Cancellation of a running exact: a SharedArrayBuffer the page also holds, as Int32Array(2): [0] a flag, [1] a
  // cancellation counter. When the page sends a newer request while an exact is outstanding it adds 1 to [1] and
  // writes 2 to [0]. The interior point polls the buffer before every Newton step (webengine.exact: should_stop) and
  // stops at once: the exact returns "cancelled" within one step (about 0.1 s). An exact message carries the counter's
  // value when it was sent (gen): if the counter moved since, the exact is superseded before it starts.
  // Pyodide's own interrupt buffer (setInterruptBuffer, an asynchronous KeyboardInterrupt) was tried for this and is
  // not used: an interrupt that lands inside casadi's SWIG wrappers surfaces as a SystemError instead of a
  // KeyboardInterrupt. Polling the same shared memory at a safe point gives the same latency without that failure.
  // Without SharedArrayBuffer (no cross-origin isolation) a running solve cannot be cancelled.
  const canCancel = typeof SharedArrayBuffer !== "undefined" && globalThis.crossOriginIsolated !== false;
  const interruptSab = canCancel ? new SharedArrayBuffer(8) : null;
  const interrupt = canCancel ? new Int32Array(interruptSab) : null;

  // --- GENCO: onnxruntime-web, WASM backend
  const sess = await sessP;
  const threads = ort.env.wasm.numThreads;
  const modelInfo = JSON.parse(new TextDecoder().decode(modelJson));
  const inType = Object.fromEntries(modelInfo.inputs.map(([n, ty]) => [n, ty]));
  stage("model");

  // --- Python <-> JS
  const unpack = (r) => {                                    // {"json", "arrays"} -> [object, {name: typed array}]
    const meta = JSON.parse(r.get("json"));
    const ap = r.get("arrays");
    const arrays = {};
    for (const name of ap.keys()) {
      const a = ap.get(name);
      const buf = a.getBuffer();
      arrays[name] = { data: buf.data.slice(), shape: Array.from(buf.shape) };
      buf.release();
      a.destroy();
    }
    ap.destroy();
    r.destroy();
    return [meta, arrays];
  };

  async function genco(arrays) {
    const feeds = {};
    for (const n of sess.inputNames) {
      const a = arrays[n];
      feeds[n] = new ort.Tensor(inType[n] === "bool" ? "bool" : inType[n] === "int64" ? "int64" : "float32", a.data, a.shape);
    }
    const a = now();
    const out = await sess.run(feeds);
    const ms = now() - a;
    return { bus: out.bus_out.data, gen: out.gen_out.data, ms };
  }

  const result = (meta, arrays, extra) => {
    const r = { ...meta, ...extra };
    for (const [k, v] of Object.entries(arrays)) r[k] = v.data;
    return r;
  };

  async function runPreview(knobsJson, noHeat) {
    const a = now();
    const cached = E.cached_preview(knobsJson, noHeat);
    if (cached !== undefined && cached !== null) {           // same what-if as a recent call: its result, with the
      const [m, arr] = unpack(cached);                       // original computation's time_ms (cached: true)
      return result(m, arr, { lookup_ms: now() - a });
    }
    const [prep, inputs] = unpack(E.prepare(knobsJson, noHeat));
    if (prep.island) {                                       // the page checks connectivity first; this is the guard
      const nan = (m) => new Float32Array(m).fill(NaN);
      return { island: true, converged: false, cost: null, max_loading: null, n_viol: null,
        viol_detail: "these line outages split the network into islands", feasible: false, time_ms: now() - a,
        line_loading: nan(info.n_branch), bus_vm: nan(info.n_bus), gen_p: nan(info.n_gen), v_out: new Int32Array(0) };
    }
    const g = await genco(inputs);
    const [m, arr] = unpack(E.preview(g.bus, g.gen, g.ms / 1000));
    const times = { time_ms: now() - a, genco_ms: g.ms, prepare_ms: (prep.build_s + prep.graph_s) * 1000 };
    const tp = py.toPy(times);
    E.record_times(tp);
    tp.destroy();
    return result(m, arr, times);
  }

  let casadiP = null;                                       // the exact solver's package, loaded on the first exact call

  // The cancellation state of one request (exact or check): gen, the counter's value when the page sent it (or, if
  // absent, its value now); superseded() is true once the page bumped the counter (or set the flag) after that.
  function cancellation(opts) {
    const gen = Number.isInteger(opts.gen) ? opts.gen : undefined;
    const sent = () => interrupt !== null && (gen !== undefined ? Atomics.load(interrupt, 1) !== gen
      : Atomics.load(interrupt, 0) === 2);
    if (gen !== undefined && sent()) return { early: true };
    if (interrupt) Atomics.store(interrupt, 0, 0);
    const gen0 = interrupt ? (gen !== undefined ? gen : Atomics.load(interrupt, 1)) : undefined;
    const stopNow = () => interrupt !== null && (Atomics.load(interrupt, 1) !== gen0 || Atomics.load(interrupt, 0) === 2);
    return { early: false, gen, stopNow, poll: interrupt ? stopNow : undefined };
  }

  const api = {
    info, threads,
    async preview(knobs) { return runPreview(JSON.stringify(knobs), false); },
    // opts.gen: the cancellation counter when the page sent the check. A newer preview, check, compare or cancel that
    // bumps the counter stops it between its steps (before every correction step, repair and Newton solve).
    async check(knobs, opts = {}) {
      const a = now();
      const kj = JSON.stringify(knobs);
      const cs = cancellation(opts);
      if (!cs.early && opts.notify) opts.notify("check");      // the instant this check's time_ms starts
      const nan = (m) => new Float32Array(m).fill(NaN);
      const cancelledCheck = () => ({ status: "cancelled", reason: "cancelled by a newer request", cached: false,
        converged: false, cost: null, max_loading: null, n_viol: null, viol_detail: "cancelled", feasible: false,
        loops: 0, time_ms: now() - a, line_loading: nan(info.n_branch), bus_vm: nan(info.n_bus), gen_p: nan(info.n_gen),
        v_out: new Int32Array(0) });
      if (cs.early) return cancelledCheck();
      let pre = null;
      if (!E.has_preview(kj)) {
        pre = await runPreview(kj, false);
        if (pre.island) return { ...pre, loops: 0, time_ms: now() - a };
      }
      if (cs.stopNow && cs.stopNow()) return cancelledCheck();
      const reused = pre ? null : JSON.parse(E.preview_times(kj));
      let r;
      try {
        r = E.check(kj, cs.poll);
      } finally {
        if (interrupt) Atomics.store(interrupt, 0, 0);
      }
      const [m, arr] = unpack(r);
      if (m.status === "cancelled") return result(m, arr, { time_ms: now() - a });
      // time_ms: this call's wall time (the preview included when this call ran it); preview_ms: that preview's time,
      // or the time of the earlier preview whose GENCO output the check reused (preview_reused: true)
      return result(m, arr, { time_ms: now() - a, preview_ms: pre ? pre.time_ms : (reused.time_ms ?? null),
        preview_reused: !pre });
    },
    // exact(knobs): the AC-OPF optimum (CasADi + MIPS in Pyodide), started from GENCO's prediction for the what-if
    // (its preview is run first if not cached), hard limit EXACT_LIMIT_S on the interior point; the first call also
    // loads casadi and builds the model (package_ms, build_ms).
    // opts: start "check" (default: the checked point; the check runs first if missing) | "genco" | "flat" | "mid";
    // limit_s (default 20 s from the check or GENCO, 60 s from scratch); debug (start and solution vectors in the
    // result). A newer request written to the interrupt buffer during the solve cancels it.
    async exact(knobs, opts = {}) {
      const a = now();
      const kj = JSON.stringify(knobs);
      const STARTS = ["check", "genco", "flat", "mid"];
      if (opts.start !== undefined && !STARTS.includes(opts.start)) {
        throw new Error(`exact: unknown start ${JSON.stringify(opts.start)} (one of ${STARTS.join(", ")})`);
      }
      const start = opts.start ?? "check";
      const cs = cancellation(opts);
      const gen = cs.gen;
      let checkMs = 0;
      const cancelled = (packageMs = 0) => ({ status: "cancelled", start, start_used: start,
        reason: "cancelled by a newer request", cached: false, time_ms: 0, build_ms: 0, check_ms: checkMs,
        package_ms: packageMs, total_ms: now() - a });
      if (cs.early) return cancelled();
      const stopNow = cs.stopNow;                              // polled by the solver before every Newton step
      let pre = null;
      const island = (r) => ({ ...r, status: "not_certified", start, start_used: start, iterations: 0, gap_vs_check: null,
        time_ms: 0, build_ms: 0, check_ms: checkMs, total_ms: now() - a });
      if (start === "check" && !E.has_check(kj)) {            // the progressive pipeline: GENCO, repair, check, exact
        const chk = await api.check(knobs, gen !== undefined ? { gen, notify: opts.notify } : { notify: opts.notify });
        checkMs = chk.time_ms;
        if (chk.island) return island(chk);
        if (chk.status === "cancelled") return cancelled();
      } else if (!E.has_preview(kj)) {
        pre = await runPreview(kj, false);
        if (pre.island) return island(pre);
      }
      let packageMs = 0;
      const notify = opts.notify ?? (() => {});
      if (!casadiP || !E.exact_model_ready()) notify("setup");   // casadi to load and/or the model to build
      if (!casadiP) {
        const b = now();
        casadiP = py.loadPackage(["casadi"], { messageCallback: () => {}, errorCallback: (s) => console.warn(s) });
        casadiP.catch(() => { casadiP = null; });
        await casadiP;
        packageMs = now() - b;
      } else await casadiP;
      // time_ms (from Python): the solve only, MIPS with its retry plus the scoring; build_ms: the CasADi model built
      // on this call (0 once built); package_ms: casadi loaded on this call; total_ms: this call's wall time
      E.exact_model(kj);                                       // the model build (first call): never interrupted
      if (stopNow && stopNow()) {                              // cancelled while the preview ran, casadi loaded or the model built
        Atomics.store(interrupt, 0, 0);
        return cancelled(packageMs);
      }
      let r;
      try {
        r = E.exact(kj, Number.isFinite(opts.limit_s) ? opts.limit_s : undefined, !!opts.debug, start,   // undefined: None
          cs.poll, undefined, () => notify("solve"));                // "solve": right before the interior point
      } finally {
        if (interrupt) Atomics.store(interrupt, 0, 0);
      }
      const [m, arr] = unpack(r);
      // check_ms: the check run first by this call (start "check"; 0 if it was cached); time_ms stays the solve only
      return result(m, arr, { package_ms: packageMs, preview_ms: pre ? pre.time_ms : 0, preview_reused: !pre,
        check_ms: checkMs, total_ms: now() - a });
    },
    async compare(knobs) {
      const a = now();
      const kj = JSON.stringify(knobs);
      const baseRes = await runPreview(kj, false);
      const hasDome = Array.isArray(knobs.domes) && knobs.domes.length > 0;
      // app.run_compare: PIPE.compare_no_heat(k) = preview(k without domes); with no dome that is k itself (cached)
      const noheat = await runPreview(kj, true);
      return { base: baseRes, noheat, no_dome: !hasDome, time_ms: now() - a };
    },
  };

  if (opt.warmup !== false) {
    await api.preview(DEFAULT_KNOBS);
    stage("warmup");
  }
  const load_ms = now() - t0;
  api.initInfo = { ready: true, load_ms, threads, cross_origin_isolated: !!globalThis.crossOriginIsolated,
    can_cancel: canCancel, interrupt_buffer: interruptSab, stages,
    fetched, info };
  return api;
}

// every ArrayBuffer behind a typed array in a result (for postMessage's transfer list)
export function transferables(obj, out = new Set()) {
  if (obj && typeof obj === "object") {
    if (ArrayBuffer.isView(obj)) out.add(obj.buffer);
    else for (const v of Object.values(obj)) transferables(v, out);
  }
  return [...out];
}