import { mergeChainResults } from "./mergeChains";
import { parseProgress } from "./parseProgress";
import type {
  ChainResult,
  ChainStatus,
  MultiChainResult,
  SamplingOpts,
} from "./types";

// Runs chains remotely on a wasm-exec service: one POST /v1/execute job per
// chain, module = the model compiled as a pure-WASI command module (see
// wasi-prototype/). The driver's argv: data_json seed chain_id num_warmup
// num_samples init_radius refresh. Progress arrives as SSE stderr events,
// draws as CSV on SSE stdout events.

type SseEvent = { event: string; data: any };

async function* sseEvents(body: ReadableStream<Uint8Array>): AsyncGenerator<SseEvent> {
  const reader = body.getReader();
  const decoder = new TextDecoder();
  let buf = "";
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    buf += decoder.decode(value, { stream: true });
    let sep;
    while ((sep = buf.indexOf("\n\n")) !== -1) {
      const block = buf.slice(0, sep);
      buf = buf.slice(sep + 2);
      let event = "message";
      let data = "";
      for (const line of block.split("\n")) {
        if (line.startsWith("event:")) event = line.slice(6).trim();
        else if (line.startsWith("data:")) data += line.slice(5).trim();
      }
      if (data) yield { event, data: JSON.parse(data) };
    }
  }
}

const runOneRemoteChain = async (
  serverUrl: string,
  clientKey: string,
  moduleBytes: Uint8Array,
  args: string[],
  chainId: number,
  signal: AbortSignal,
  onStatus: (chainId: number, status: ChainStatus) => void,
): Promise<ChainResult> => {
  const form = new FormData();
  // POC: send the module bytes with every job. The server caches by sha256,
  // so this could be optimized to module_sha256 with a 404 fallback.
  form.append("module", new Blob([moduleBytes as BlobPart]), "main-wasi.wasm");
  form.append("args", JSON.stringify(args));
  form.append("timeout_ms", "300000");

  const resp = await fetch(`${serverUrl}/v1/execute`, {
    method: "POST",
    headers: { Authorization: `Bearer ${clientKey}` },
    body: form,
    signal,
  });
  if (!resp.ok || !resp.body) {
    let message = resp.statusText;
    try {
      message = (await resp.json()).message ?? message;
    } catch {
      // keep statusText
    }
    throw new Error(`chain ${chainId}: submit failed: ${message}`);
  }

  let stdoutText = "";
  let stderrText = "";
  let stderrLineBuf = "";
  let exitCode: number | undefined;

  for await (const { event, data } of sseEvents(resp.body)) {
    if (event === "stdout") {
      stdoutText += atob(data.b64);
    } else if (event === "stderr") {
      const text = atob(data.b64);
      stderrText += text;
      stderrLineBuf += text;
      let nl;
      while ((nl = stderrLineBuf.indexOf("\n")) !== -1) {
        const line = stderrLineBuf.slice(0, nl).trim();
        stderrLineBuf = stderrLineBuf.slice(nl + 1);
        const progress = parseProgress(line);
        if (progress) onStatus(chainId, progress);
      }
    } else if (event === "exit") {
      exitCode = data.exitCode;
    } else if (event === "error") {
      throw new Error(`chain ${chainId}: ${data.code}: ${data.message}`);
    }
  }

  if (exitCode === undefined)
    throw new Error(`chain ${chainId}: stream ended without exit event`);
  if (exitCode !== 0)
    throw new Error(
      `chain ${chainId}: module exited with code ${exitCode}\n${stderrText.slice(-1000)}`,
    );

  // stdout: header line of param names, then one CSV row per draw
  const lines = stdoutText.trim().split("\n");
  const paramNames = lines[0].split(",");
  const rows = lines.slice(1).map((l) => l.split(",").map(Number));
  const draws = paramNames.map((_, p) => rows.map((r) => r[p]));
  return { chainId, paramNames, draws, consoleText: stderrText };
};

const runRemoteChains = (
  serverUrl: string,
  clientKey: string,
  moduleBytes: Uint8Array,
  data: string,
  opts: SamplingOpts,
  onChainStatus: (chainId: number, status: ChainStatus) => void,
): { result: Promise<MultiChainResult>; cancel: () => void } => {
  const seed = opts.seed ?? Math.floor(Math.random() * Math.pow(2, 32));
  const refresh = Math.max(
    10,
    Math.floor((opts.num_warmup + opts.num_samples) / 100) * 10,
  );
  const startTimeSec = performance.now() / 1000;
  const abort = new AbortController();

  const chainPromises = Array.from({ length: opts.num_chains }, (_, i) => {
    const chainId = i + 1;
    const args = [
      data,
      String(seed),
      String(chainId),
      String(opts.num_warmup),
      String(opts.num_samples),
      String(opts.init_radius),
      String(refresh),
    ];
    return runOneRemoteChain(
      serverUrl.replace(/\/$/, ""),
      clientKey,
      moduleBytes,
      args,
      chainId,
      abort.signal,
      onChainStatus,
    ).then(
      (r) => {
        onChainStatus(chainId, "done");
        return r;
      },
      (e) => {
        onChainStatus(chainId, "error");
        throw e;
      },
    );
  });

  const result = Promise.all(chainPromises).then((chainResults) =>
    mergeChainResults(chainResults, performance.now() / 1000 - startTimeSec),
  );

  return { result, cancel: () => abort.abort() };
};

export default runRemoteChains;