/**
 * MATLAB script source -> numbl's JIT IR, ready for the WGSL planner.
 *
 * Unlike turing-surface — which specializes named functions against
 * host-supplied argument types — the sandbox lowers a whole *script*: shapes
 * come from the script itself (`n = 2048; A = rand(n);`), pinned static by
 * numbl's exact-value propagation through the type lattice.
 *
 * Two numbl passes matter here:
 *  - `lowerProgram` lowers the top-level statements to IR, one statement per
 *    operation (ANF), with every node's type fixed.
 *  - `inlinePass` then folds single-use temps back into their consumer, so a
 *    source line like `y = a - u + u.*u.*v` becomes ONE statement whose RHS is
 *    an expression tree — i.e. one fused GPU kernel instead of four.
 */
import { parseMFile } from 'numbl-src/numbl-core/parser/index.ts';
import { Workspace, Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
import type { IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
import { applyBuiltinPatches } from './patches.ts';
import { fuseTemps } from './fuse.ts';
import { inScript, ScriptCompileError } from './errors.ts';

export interface CompiledScript {
  /** The lowered, inline-folded top-level statements. */
  stmts: IRStmt[];
  /**
   * Should the statement covering source offset `at` echo its result,
   * MATLAB-style? True exactly when the source statement has no trailing
   * semicolon. Compiler temps sit inside their source statement's span, so
   * the caller must additionally skip `_mtoc2_*` names.
   */
  isEchoed(at: number): boolean;
}

export function compileScript(source: string, fileName = 'script.m'): CompiledScript {
  applyBuiltinPatches();

  const ast = inScript(() => parseMFile(source, fileName));

  // Statements the parser marked unsuppressed (no `;`), by source range.
  const echoed: { start: number; end: number }[] = [];
  for (const s of ast.body) {
    if (s.suppressed === false) echoed.push({ start: s.span.start, end: s.span.end });
  }

  if (ast.body.some((s) => s.type === 'Function')) {
    // A script may syntactically end with local functions, but nothing here
    // compiles calls to them — say so up front rather than at the call site.
    throw new ScriptCompileError(
      'local functions are not supported in the sandbox yet — inline their bodies',
    );
  }

  const prog = inScript(() => {
    const ws = new Workspace(fileName, []);
    ws.addFile({ name: fileName, source, ast });
    ws.finalize();
    const lowered = new Lowerer(ws).lowerProgram(ast);
    inlinePass(lowered);
    // numbl's pass stops at what its C backend fuses; fold the rest of the
    // single-use temps the WGSL emitter can absorb (sin/exp, comparisons,
    // logicals, generators) so one source line is one kernel.
    fuseTemps(lowered.topLevelStmts);
    return lowered;
  });

  return {
    stmts: prog.topLevelStmts,
    isEchoed: (at) => echoed.some((r) => at >= r.start && at <= r.end),
  };
}