// Loading of FMI-LS-WASM FMUs: ZIP extraction and modelDescription.xml parsing.
// Instantiating the component — transpiling it with jco and reaching its core
// exports — is `fmu-core.js`; what the page shows about an FMU is parsed here.

import { loadComponent } from './fmu-core.js';

async function inflateRaw(bytes) {
  const s = new Blob([bytes]).stream().pipeThrough(new DecompressionStream('deflate-raw'));
  return new Uint8Array(await new Response(s).arrayBuffer());
}

// Minimal ZIP reader over the central directory. Stored and deflated entries only.
export async function readZip(buf) {
  const dv = new DataView(buf), u8 = new Uint8Array(buf), dec = new TextDecoder();
  let eocd = -1;
  for (let i = buf.byteLength - 22; i >= 0 && i > buf.byteLength - 22 - 65536; i--) {
    if (dv.getUint32(i, true) === 0x06054b50) { eocd = i; break; }
  }
  if (eocd < 0) throw new Error('not a ZIP archive: no end-of-central-directory record');
  const count = dv.getUint16(eocd + 10, true);
  let p = dv.getUint32(eocd + 16, true);
  const files = new Map();
  for (let i = 0; i < count; i++) {
    if (dv.getUint32(p, true) !== 0x02014b50) throw new Error('corrupt ZIP central directory');
    const method = dv.getUint16(p + 10, true);
    const csize = dv.getUint32(p + 20, true);
    const nameLen = dv.getUint16(p + 28, true);
    const extraLen = dv.getUint16(p + 30, true);
    const cmtLen = dv.getUint16(p + 32, true);
    const lho = dv.getUint32(p + 42, true);
    const name = dec.decode(u8.subarray(p + 46, p + 46 + nameLen));
    if (csize === 0xffffffff || lho === 0xffffffff) throw new Error(`ZIP64 entry not supported: ${name}`);
    const start = lho + 30 + dv.getUint16(lho + 26, true) + dv.getUint16(lho + 28, true);
    const raw = u8.subarray(start, start + csize);
    if (!name.endsWith('/')) {
      if (method !== 0 && method !== 8) throw new Error(`unsupported ZIP compression method ${method}: ${name}`);
      files.set(name, method === 8 ? await inflateRaw(raw) : raw.slice());
    }
    p += 46 + nameLen + extraLen + cmtLen;
  }
  return files;
}

async function deflateRaw(bytes) {
  const s = new Blob([bytes]).stream().pipeThrough(new CompressionStream('deflate-raw'));
  return new Uint8Array(await new Response(s).arrayBuffer());
}

const CRC = Int32Array.from({ length: 256 }, (_, n) => {
  let c = n;
  for (let k = 0; k < 8; k++) c = c & 1 ? 0xedb88320 ^ (c >>> 1) : c >>> 1;
  return c;
});
function crc32(bytes) {
  let c = -1;
  for (const b of bytes) c = CRC[(c ^ b) & 0xff] ^ (c >>> 8);
  return (c ^ -1) >>> 0;
}

// Write a Map of name -> bytes back out as a deflated ZIP, so an FMU can be handed
// back with entries added. Counterpart of [`readZip`]; no ZIP64.
export async function writeZip(files) {
  const enc = new TextEncoder();
  const local = [];
  const central = [];
  let offset = 0;
  for (const [name, bytes] of files) {
    const nameBytes = enc.encode(name);
    const deflated = await deflateRaw(bytes);
    const stored = deflated.length >= bytes.length;
    const data = stored ? bytes : deflated;
    const head = new DataView(new ArrayBuffer(30));
    head.setUint32(0, 0x04034b50, true);
    head.setUint16(4, 20, true);
    head.setUint16(8, stored ? 0 : 8, true);
    head.setUint16(12, 33, true); // 1980-01-01, the epoch a ZIP date field can hold
    head.setUint32(14, crc32(bytes), true);
    head.setUint32(18, data.length, true);
    head.setUint32(22, bytes.length, true);
    head.setUint16(26, nameBytes.length, true);
    local.push(new Uint8Array(head.buffer), nameBytes, data);

    const dir = new DataView(new ArrayBuffer(46));
    dir.setUint32(0, 0x02014b50, true);
    dir.setUint16(4, 20, true);
    dir.setUint16(6, 20, true);
    dir.setUint16(10, stored ? 0 : 8, true);
    dir.setUint16(14, 33, true);
    dir.setUint32(16, crc32(bytes), true);
    dir.setUint32(20, data.length, true);
    dir.setUint32(24, bytes.length, true);
    dir.setUint16(28, nameBytes.length, true);
    dir.setUint32(42, offset, true);
    central.push(new Uint8Array(dir.buffer), nameBytes);
    offset += 30 + nameBytes.length + data.length;
  }
  const centralSize = central.reduce((n, p) => n + p.length, 0);
  const end = new DataView(new ArrayBuffer(22));
  end.setUint32(0, 0x06054b50, true);
  end.setUint16(8, files.size, true);
  end.setUint16(10, files.size, true);
  end.setUint32(12, centralSize, true);
  end.setUint32(16, offset, true);
  return new Blob([...local, ...central, new Uint8Array(end.buffer)]);
}

// FMI 3.0 variable element name -> the get/set suffix used by the WIT methods.
// Enumeration is an Int64 in FMI 3.0.
const TYPES = {
  Float32: 'Float32', Float64: 'Float64',
  Int8: 'Int8', Int16: 'Int16', Int32: 'Int32', Int64: 'Int64',
  UInt8: 'UInt8', UInt16: 'UInt16', UInt32: 'UInt32', UInt64: 'UInt64',
  Boolean: 'Boolean', String: 'String', Binary: 'Binary',
  Enumeration: 'Int64', Clock: 'Clock',
};
const NUMERIC = new Set(['Float32', 'Float64', 'Int8', 'Int16', 'Int32', 'Int64',
  'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Boolean', 'Enumeration']);

const attr = (e, n) => (e && e.hasAttribute(n) ? e.getAttribute(n) : null);
const numAttr = (e, n) => { const v = attr(e, n); return v == null || v === '' ? null : parseFloat(v); };
const boolAttr = (e, n) => { const v = attr(e, n); return v == null ? null : v === 'true' || v === '1'; };

function iface(e) {
  if (!e) return null;
  return {
    modelIdentifier: attr(e, 'modelIdentifier') || '',
    needsExecutionTool: boolAttr(e, 'needsExecutionTool') === true,
    hasEventMode: boolAttr(e, 'hasEventMode') === true,
    canHandleVariableCommunicationStepSize: boolAttr(e, 'canHandleVariableCommunicationStepSize') !== false,
    fixedInternalStepSize: numAttr(e, 'fixedInternalStepSize'),
    providesDirectionalDerivatives: boolAttr(e, 'providesDirectionalDerivatives') === true,
  };
}

// The OpenModelica <Figures> vendor annotation. Absent / unknown-version /
// malformed degrades to []: such an FMU is plotted the ordinary way, never rejected.
function parseFigures(root) {
  const viz = root.querySelector(':scope > Annotations > Tool[name="OpenModelica"] > Figures');
  if (!viz) return [];
  const version = attr(viz, 'version');
  if (version != null && version !== '1') return [];   // unknown schema: ignore, don't guess
  const axis = (plot, role) => {
    const a = plot.querySelector(`:scope > Axis[role="${role}"]`);
    if (!a) return null;
    return {
      label: attr(a, 'label') || '', unit: attr(a, 'unit') || '',
      min: numAttr(a, 'min'), max: numAttr(a, 'max'),
      log: (attr(a, 'scale') || 'Linear') === 'Log',
    };
  };
  const figures = [];
  for (const f of viz.querySelectorAll(':scope > Figure')) {
    const plots = [];
    for (const p of f.querySelectorAll(':scope > Plot')) {
      const curves = [];
      for (const c of p.querySelectorAll(':scope > Curve')) {
        const y = attr(c, 'y'); if (!y) continue;
        curves.push({ x: attr(c, 'x') || '', y, legend: attr(c, 'legend') || '' });
      }
      const tr = p.querySelector(':scope > TerminalRef');
      if (!curves.length) continue;   // TerminalRef-only plot: skip, we render explicit curves
      plots.push({
        title: attr(p, 'title') || '', preferred: boolAttr(p, 'preferred') === true,
        terminal: tr ? attr(tr, 'terminal') : null, curves,
        x: axis(p, 'x'), y: axis(p, 'y'), y2: axis(p, 'y2'),
      });
    }
    if (!plots.length) continue;
    const cap = f.querySelector(':scope > Caption');
    figures.push({
      title: attr(f, 'title') || '', group: attr(f, 'group') || '',
      preferred: boolAttr(f, 'preferred') === true,
      caption: cap ? cap.textContent : '', plots,
    });
  }
  return figures;
}

// The OpenModelica <Visualization> annotation naming the _visual.xml resource, or null.
function parseVisualization(root) {
  const v = root.querySelector(':scope > Annotations > Tool[name="OpenModelica"] > Visualization');
  if (!v) return null;
  const version = attr(v, 'version');
  if (version != null && version !== '1') return null;
  const file = attr(v, 'file');
  return file ? { file } : null;
}

export function parseModelDescription(xml) {
  const doc = new DOMParser().parseFromString(xml, 'application/xml');
  const perr = doc.querySelector('parsererror');
  if (perr) throw new Error('modelDescription.xml is not well-formed: ' + perr.textContent.trim());
  const root = doc.documentElement;
  if (root.tagName !== 'fmiModelDescription') throw new Error('modelDescription.xml: root element is not fmiModelDescription');
  const fmiVersion = attr(root, 'fmiVersion') || '';
  if (!fmiVersion.startsWith('3.')) throw new Error(`FMI version ${fmiVersion || '?'} is not supported; this simulator requires FMI 3.0`);

  const variables = [];
  // FMI 3.0 <Alias> name -> primary name; many visualization crefs are aliases.
  const aliases = new Map();
  const mv = root.querySelector('ModelVariables');
  for (const e of mv ? Array.from(mv.children) : []) {
    const type = TYPES[e.tagName];
    if (!type) continue;
    // String and Binary carry their start in a <Start value="…"/> child.
    const startEl = e.querySelector(':scope > Start');
    const startAttr = attr(e, 'start');
    const start = startEl ? attr(startEl, 'value')
      : startAttr == null ? null : startAttr.trim().split(/\s+/)[0];
    const name = attr(e, 'name') || '';
    variables.push({
      name,
      vr: parseInt(attr(e, 'valueReference'), 10),
      tag: e.tagName,
      type,
      numeric: NUMERIC.has(e.tagName),
      causality: attr(e, 'causality') || 'local',
      variability: attr(e, 'variability') || (e.tagName.startsWith('Float') ? 'continuous' : 'discrete'),
      start,
      initial: attr(e, 'initial'),
      unit: attr(e, 'unit') || attr(e, 'displayUnit') || '',
      description: attr(e, 'description') || '',
      derivative: attr(e, 'derivative'),
    });
    for (const al of e.querySelectorAll(':scope > Alias')) {
      const an = attr(al, 'name');
      if (an) aliases.set(an, name);
    }
  }

  const de = root.querySelector('DefaultExperiment');
  const ms = root.querySelector('ModelStructure');
  return {
    modelName: attr(root, 'modelName') || '',
    fmiVersion,
    instantiationToken: attr(root, 'instantiationToken') || '',
    description: attr(root, 'description') || '',
    generationTool: attr(root, 'generationTool') || '',
    version: attr(root, 'version') || '',
    me: iface(root.querySelector('ModelExchange')),
    cs: iface(root.querySelector('CoSimulation')),
    se: iface(root.querySelector('ScheduledExecution')),
    defaultExperiment: {
      startTime: de ? numAttr(de, 'startTime') : null,
      stopTime: de ? numAttr(de, 'stopTime') : null,
      stepSize: de ? numAttr(de, 'stepSize') : null,
      tolerance: de ? numAttr(de, 'tolerance') : null,
    },
    variables,
    aliases,
    nStates: ms ? ms.querySelectorAll(':scope > ContinuousStateDerivative').length : 0,
    nEventIndicators: ms ? ms.querySelectorAll(':scope > EventIndicator').length : 0,
    figures: parseFigures(root),
    visualization: parseVisualization(root),
  };
}

// The text of an FMU resource (path relative to resources/), or null.
export function readResource(files, name) {
  const bytes = files.get('resources/' + name);
  return bytes ? new TextDecoder().decode(bytes) : null;
}

// The wasm component lives at binaries/wasm32-wasip2/<modelIdentifier>.wasm; fall
// back to the single .wasm in that directory when the identifier does not match.
export function findComponent(files, md) {
  const dir = 'binaries/wasm32-wasip2/';
  const ids = [md.cs, md.me, md.se].filter(Boolean).map((i) => i.modelIdentifier);
  for (const id of ids) {
    const bytes = files.get(`${dir}${id}.wasm`);
    if (bytes) return bytes;
  }
  const candidates = [...files.keys()].filter((n) => n.startsWith(dir) && n.endsWith('.wasm'));
  if (candidates.length === 1) return files.get(candidates[0]);
  if (!candidates.length) throw new Error(`no wasm component found in ${dir} (is this an FMI-LS-WASM FMU?)`);
  throw new Error(`no component matching modelIdentifier ${ids.join('/')} in ${dir}`);
}

// Read an FMU. `withComponent` transpiles and instantiates the wasm component
// here as well, which the page does not need: it shows what the model
// description says and leaves running the FMU to the worker.
export async function loadFmu(buf, { onLog, withComponent = false } = {}) {
  const files = await readZip(buf);
  const mdFile = files.get('modelDescription.xml');
  if (!mdFile) throw new Error('modelDescription.xml is missing from the archive');
  const md = parseModelDescription(new TextDecoder().decode(mdFile));
  if (!withComponent) return { md, files, component: null, instance: null };
  const component = await loadComponent(findComponent(files, md), { files, onLog });
  const instance = async () => ({ md, ...(await component.probe()) });
  return { md, files, component, instance };
}
