/**
 * Execute a ScriptPlan on a GPUDevice.
 *
 * GPU ops stream into a command encoder and are submitted in batches; the
 * host ops the script asked for (tic/toc/disp/fprintf/echo) are the only
 * synchronization points. `tic` and `toc` flush pending work and await
 * `onSubmittedWorkDone`, so `toc` reports wall-clock time for work that has
 * actually finished — the same thing MATLAB's synchronous tic/toc measures,
 * which is what makes the number comparable when the script is pasted there.
 *
 * A `loop` op re-encodes its body per iteration, switching the loop-variable
 * uniform's dynamic offset; everything still lands in one submit.
 */
import { WORKGROUP_SIZE } from './wgsl.ts';
import type { EmitPart, Op, ScriptPlan, Slot, ValueRef } from './plan.ts';

void WORKGROUP_SIZE;

export interface TimingSegment {
  /** 1-based tic..toc pair index. */
  seq: number;
  seconds: number;
}

export interface RunResult {
  /** Everything the script printed, in order. */
  output: string;
  segments: TimingSegment[];
  /** Wall time of the whole execution (excluding compilation). */
  totalSeconds: number;
  error?: string;
}

const LV_STRIDE = 256;

export async function executePlan(
  device: GPUDevice,
  plan: ScriptPlan,
  onOutput?: (text: string) => void,
): Promise<RunResult> {
  let output = '';
  const segments: TimingSegment[] = [];
  const print = (text: string): void => {
    output += text;
    onOutput?.(text);
  };

  let encoder: GPUCommandEncoder | null = null;
  let pass: GPUComputePassEncoder | null = null;
  const inEncoder = (): GPUCommandEncoder => {
    if (!encoder) encoder = device.createCommandEncoder();
    return encoder;
  };
  const inPass = (): GPUComputePassEncoder => {
    if (!pass) pass = inEncoder().beginComputePass();
    return pass;
  };
  const endPass = (): void => {
    if (pass) {
      pass.end();
      pass = null;
    }
  };
  const flush = (): void => {
    endPass();
    if (encoder) {
      device.queue.submit([encoder.finish()]);
      encoder = null;
    }
  };
  const sync = async (): Promise<void> => {
    flush();
    await device.queue.onSubmittedWorkDone();
  };

  /** Values of tic/toc-produced variables, in seconds. */
  const hostVals = new Map<string, number>();
  let ticStartMs: number | null = null;

  const readBuffer = async (slot: Slot, count: number): Promise<Float32Array> => {
    flush();
    const bytes = Math.max(4, 4 * count);
    const staging = device.createBuffer({
      size: bytes,
      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
    const e = device.createCommandEncoder();
    e.copyBufferToBuffer(slot.buffer, 0, staging, 0, bytes);
    device.queue.submit([e.finish()]);
    await staging.mapAsync(GPUMapMode.READ);
    const data = new Float32Array(staging.getMappedRange().slice(0));
    staging.unmap();
    staging.destroy();
    return data.subarray(0, count);
  };

  const refValue = async (ref: ValueRef): Promise<number> => {
    switch (ref.kind) {
      case 'literal':
        return ref.value;
      case 'host':
        return hostVals.get(ref.cName) ?? NaN;
      case 'buffer':
        return (await readBuffer(ref.slot, 1))[0];
    }
  };

  const writeScalar = (slot: Slot, value: number): void => {
    device.queue.writeBuffer(slot.buffer, 0, new Float32Array([value]) as Float32Array<ArrayBuffer>);
  };

  async function execOps(ops: Op[], offsets: Map<string, number>): Promise<void> {
    for (const op of ops) {
      switch (op.kind) {
        case 'kernel': {
          const p = inPass();
          p.setPipeline(op.pipeline);
          if (op.loops.length) {
            p.setBindGroup(0, op.bindGroup, op.loops.map((cv) => offsets.get(cv) ?? 0));
          } else {
            p.setBindGroup(0, op.bindGroup);
          }
          p.dispatchWorkgroups(op.dispatch[0], op.dispatch[1]);
          if (op.copyBack) {
            endPass();
            inEncoder().copyBufferToBuffer(
              op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
            );
          }
          break;
        }
        case 'copy':
          endPass();
          inEncoder().copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
          break;
        case 'write':
          // Queue writes execute before any later submit; flush pending
          // encodes first so ordering matches program order.
          flush();
          device.queue.writeBuffer(op.slot.buffer, 0, op.data as Float32Array<ArrayBuffer>);
          break;
        case 'loop': {
          for (let it = 0; it < op.trips; it++) {
            offsets.set(op.cVar, it * LV_STRIDE);
            await execOps(op.body, offsets);
          }
          offsets.delete(op.cVar);
          break;
        }
        case 'tic': {
          await sync();
          ticStartMs = performance.now();
          if (op.assignTo) {
            const seconds = ticStartMs / 1000;
            hostVals.set(op.assignTo.cName, seconds);
            writeScalar(op.assignTo.slot, seconds);
          }
          break;
        }
        case 'toc': {
          await sync();
          const now = performance.now();
          const baseMs =
            op.sinceCName !== undefined
              ? (hostVals.get(op.sinceCName) ?? 0) * 1000
              : ticStartMs;
          if (baseMs === null) {
            print(`Error: toc without a preceding tic\n`);
            break;
          }
          const seconds = (now - baseMs) / 1000;
          segments.push({ seq: op.seq, seconds });
          if (op.print) {
            print(`Elapsed time is ${seconds.toFixed(6)} seconds.\n`);
          }
          if (op.assignTo) {
            hostVals.set(op.assignTo.cName, seconds);
            writeScalar(op.assignTo.slot, seconds);
          }
          break;
        }
        case 'emit': {
          const parts: string[] = [];
          for (const part of op.parts) {
            parts.push(await formatPart(part));
          }
          print(parts.join(''));
          break;
        }
        case 'display': {
          print(await formatDisplay(op));
          break;
        }
      }
    }
  }

  async function formatPart(part: EmitPart): Promise<string> {
    if (part.kind === 'text') return part.text;
    return formatSpec(part.spec, await refValue(part.ref));
  }

  async function formatDisplay(op: Op & { kind: 'display' }): Promise<string> {
    const head = op.label !== null ? `${op.label} =\n\n` : '';
    const count = op.shape.reduce((a, b) => a * b, 1);
    if (count === 1 || op.ref.kind !== 'buffer') {
      const v = await refValue(op.ref);
      return `${head}   ${formatShort(v)}\n\n`;
    }
    const [m, n] = op.shape.length === 2 ? op.shape : [count, 1];
    if (count > 400) {
      // MATLAB would print all of it; that is unreadable in a sandbox pane.
      const data = await readBuffer(op.ref.slot, Math.min(count, 4));
      const preview = Array.from(data).map(formatShort).join('  ');
      return `${head}   [${m}x${n}] ${preview} ...  (display truncated)\n\n`;
    }
    const data = await readBuffer(op.ref.slot, count);
    const lines: string[] = [];
    for (let r = 0; r < m; r++) {
      const cells: string[] = [];
      for (let c = 0; c < n; c++) {
        cells.push(formatShort(data[r + c * m]).padStart(12));
      }
      lines.push('  ' + cells.join(''));
    }
    return `${head}${lines.join('\n')}\n\n`;
  }

  const start = performance.now();
  let error: string | undefined;
  device.pushErrorScope('out-of-memory');
  device.pushErrorScope('validation');
  try {
    await execOps(plan.ops, new Map());
    await sync();
  } catch (e) {
    error = e instanceof Error ? e.message : String(e);
  }
  const validation = await device.popErrorScope();
  const oom = await device.popErrorScope();
  if (!error && validation) error = `GPU validation error: ${validation.message}`;
  if (!error && oom) error = `GPU out of memory: ${oom.message}`;
  const totalSeconds = (performance.now() - start) / 1000;

  return { output, segments, totalSeconds, error };
}

/** MATLAB `format short`-flavored scalar rendering. */
export function formatShort(v: number): string {
  if (!Number.isFinite(v)) return v > 0 ? 'Inf' : v < 0 ? '-Inf' : 'NaN';
  if (v === 0) return '0';
  if (Number.isInteger(v) && Math.abs(v) < 1e10) return String(v);
  const a = Math.abs(v);
  if (a >= 1e5 || a < 1e-3) return v.toExponential(4);
  return v.toFixed(4);
}

/** One printf-style conversion. */
function formatSpec(spec: string, v: number): string {
  const m = /^%([-+ 0#]*)(\d*)(?:\.(\d+))?([diufeEgGs])$/.exec(spec);
  if (!m) return String(v);
  const [, flags, widthS, precS, conv] = m;
  const width = widthS ? parseInt(widthS, 10) : 0;
  const prec = precS !== undefined ? parseInt(precS, 10) : undefined;
  let s: string;
  switch (conv) {
    case 'd':
    case 'i':
    case 'u':
      s = Number.isInteger(v) ? String(v) : v.toExponential(prec ?? 6);
      break;
    case 'f':
      s = v.toFixed(prec ?? 6);
      break;
    case 'e':
    case 'E': {
      s = v.toExponential(prec ?? 6);
      if (conv === 'E') s = s.toUpperCase();
      break;
    }
    case 'g':
    case 'G': {
      const p = prec === undefined || prec === 0 ? 6 : prec;
      const a = Math.abs(v);
      s = a !== 0 && (a < 1e-5 || a >= 10 ** p)
        ? v.toExponential(Math.max(0, p - 1)).replace(/\.?0+e/, 'e')
        : String(Number(v.toPrecision(p)));
      if (conv === 'G') s = s.toUpperCase();
      break;
    }
    case 's':
      s = String(v);
      break;
    default:
      s = String(v);
  }
  if (flags.includes('+') && v >= 0 && 'dfeg'.includes(conv.toLowerCase())) s = '+' + s;
  if (width > s.length) {
    s = flags.includes('-')
      ? s.padEnd(width)
      : flags.includes('0') && !flags.includes('-')
        ? (s.startsWith('-') ? '-' + s.slice(1).padStart(width - 1, '0') : s.padStart(width, '0'))
        : s.padStart(width);
  }
  return s;
}