import * as path from "path";
import { ts } from "../../ir/builders.js";
/**
 * Convert an absolute module path to one rooted at the codegen cwd
 * for readability in the generated registry-key string literals.
 * Falls back to the original string when `abs` is already relative
 * (legacy callers that supply a relative `moduleId` directly).
 *
 * Safety: `register*Init` and `await*Init` both flow through this
 * function before becoming string literals in the emitted JS, so
 * registry keys still match exactly within a single compilation
 * pass — the value of `process.cwd()` at runtime is irrelevant once
 * the literals are baked in.
 */
function displayModuleId(abs) {
    if (!path.isAbsolute(abs))
        return abs;
    return path.relative(process.cwd(), abs);
}
/**
 * Walk the program once and route each node into the right bucket:
 *
 *   - `static` assignments (and `with handler { static x = ... }`)
 *     contribute names + frozen init statements.
 *   - `global` assignments (and their `with handler` form) contribute
 *     `__ctx.globals.set(...)` calls.
 *   - Top-level declarations (per `isTopLevelDeclaration`) become
 *     module-scope generated statements.
 *   - Anything else (bare expressions, function calls, …) goes into
 *     `__initializeGlobals` so it can access `__ctx`.
 */
export function partitionProgram(program, deps) {
    const staticVarNames = new Set();
    const exportedStaticVarNames = new Set();
    // Tag each init statement with its var name (null = bare/anonymous)
    // so the partition can optionally reorder per the topsort plan
    // without re-walking the AST. Bare statements keep source order
    // because they have no name to key on.
    const staticInitTagged = [];
    const globalInitTagged = [];
    const topLevelStatements = [];
    const topLevelCallbackStatements = [];
    for (const node of program.nodes) {
        if (isTopLevelCallbackCall(node)) {
            topLevelCallbackStatements.push(deps.processNodeInGlobalInit(node));
            continue;
        }
        const staticAssign = unwrapStaticAssignment(node);
        if (staticAssign) {
            const { stmt, handlerName } = staticAssign;
            staticVarNames.add(stmt.variableName);
            if (stmt.exported)
                exportedStaticVarNames.add(stmt.variableName);
            const valueNode = deps.processNodeInGlobalInit(stmt.value);
            const frozenAssign = ts.assign(ts.id(stmt.variableName), ts.call(ts.id("__deepFreeze"), [valueNode]));
            staticInitTagged.push({
                varName: stmt.variableName,
                node: handlerName
                    ? ts.withHandler(deps.buildHandlerArrow(handlerName), frozenAssign)
                    : frozenAssign,
            });
            continue;
        }
        const globalAssign = unwrapGlobalAssignment(node);
        if (globalAssign) {
            const { stmt, handlerName } = globalAssign;
            const valueNode = deps.processNodeInGlobalInit(stmt.value);
            // These statements are emitted inside `__initializeGlobals(__ctx)`
            // (see buildInitializeGlobalsFn below). `__ctx` is the parameter
            // there — no ALS frame is installed by the caller — so pass the
            // lexical `__ctx` identifier as the receiver instead of letting
            // globalSet default to `getRuntimeContext().ctx`.
            const setNode = ts.globalSet(deps.moduleId, stmt.variableName, valueNode, ts.id("__ctx"));
            globalInitTagged.push({
                varName: stmt.variableName,
                node: handlerName
                    ? ts.withHandler(deps.buildHandlerArrow(handlerName), setNode)
                    : setNode,
            });
            continue;
        }
        // Bare top-level `<expr> with handler` (the form not caught by
        // `unwrapStaticAssignment` / `unwrapGlobalAssignment` because the
        // wrapped statement isn't an assignment). Without this branch, the
        // `withModifier` falls through to `processNodeInGlobalInit` →
        // `processNode` → `processWithModifier`, which calls
        // `stepPathTracker.currentId()` against an empty stack and throws
        // the cryptic internal invariant `StepPathTracker: currentId()
        // called with empty path` (issue #229). Emit the same lightweight
        // pushHandler/popHandler wrapper that the static/global assignment
        // cases above use — `withHandler` deliberately doesn't need a step
        // id, and `__initializeGlobals` already runs under
        // `runInBootstrapFrame(...)` so the handler stack is live. Resume
        // semantics match the rest of globalInit: this is a fresh-run
        // side effect, not re-executed on interrupt resume.
        if (node.type === "withModifier") {
            const innerStmt = deps.processNodeInGlobalInit(node.statement);
            globalInitTagged.push({
                varName: null,
                node: ts.withHandler(deps.buildHandlerArrow(node.handlerName), innerStmt),
            });
            continue;
        }
        // `static <expression-statement>` at module top level (PR 3). The
        // wrapper is purely a routing marker — unwrap to the inner
        // statement and route it into `staticInitTagged` so it runs once
        // per process in Phase A instead of once per agent run in Phase
        // B. Downstream codegen never sees the wrapper.
        if (node.type === "staticStatement") {
            staticInitTagged.push({
                varName: null,
                node: deps.processNodeInGlobalInit(node.statement),
            });
            continue;
        }
        if (deps.isTopLevelDeclaration(node)) {
            topLevelStatements.push(deps.processNode(node));
        }
        else {
            // Top-level statements (function calls, etc.) need __ctx access,
            // so they live inside __initializeGlobals.
            globalInitTagged.push({
                varName: null,
                node: deps.processNodeInGlobalInit(node),
            });
        }
    }
    return {
        staticVarNames,
        exportedStaticVarNames,
        staticInitStatements: reorderTagged(staticInitTagged, deps.staticOrder),
        globalInitStatements: reorderTagged(globalInitTagged, deps.globalOrder),
        topLevelStatements,
        topLevelCallbackStatements,
    };
}
/**
 * Apply the topsort plan's `localOrder` to tagged init statements.
 *
 * Bare (unnamed) statements are anchored to their source position —
 * each occupies exactly the slot where it originally appeared. Named
 * statements are placed per the plan, with the k-th name in
 * `localOrder` filling the k-th named slot in source order. This
 * preserves side-effect ordering for patterns like `foo(); const x =
 * ...; bar();` while still letting topsort sequence the named decls.
 *
 * **Invariant (post-Task-4 of the agent-init PR2 stragglers PR):** for
 * user-authored same-file named decls, the plan order is guaranteed by
 * `assertNoIntraFileUseBeforeDef` to match source order — so this loop
 * is a no-op reorder for them. The reorder branch still does real
 * work for re-export wrapper statics synthesized by `resolveReExports`
 * (`static const x = _reexport_x`): wrappers have no user-controlled
 * source position, so the section assembler may freely slot them in
 * dep-first order.
 *
 * No plan → source order, unchanged.
 */
function reorderTagged(tagged, order) {
    if (!order || order.length === 0) {
        return tagged.map((t) => t.node);
    }
    const byName = {};
    for (const { varName, node } of tagged) {
        if (varName)
            byName[varName] = node;
    }
    // Walk tagged once. Bare slots emit their own node immediately.
    // Named slots are placeholders to be filled from `order` in
    // top-to-bottom order — only names that have a matching tagged
    // node count, so cross-module-only names in `order` are skipped.
    const planNamesInTagged = order.filter((n) => byName[n]);
    const out = [];
    let namedSlotIdx = 0;
    for (const { varName, node } of tagged) {
        if (varName) {
            const fillName = planNamesInTagged[namedSlotIdx++];
            if (fillName)
                out.push(byName[fillName]);
        }
        else {
            out.push(node);
        }
    }
    return out;
}
/**
 * Detect a top-level statement that registers a callback via the stdlib
 * `callback(name, fn)` function. After the `liftCallbackBlocks`
 * preprocessor runs, block-form callbacks have already been rewritten
 * into named-fn calls of this shape, so this single shape covers both
 * source forms.
 *
 * Handler-wrapped top-level forms — `callback(...) with myHandler` and a
 * `handle { callback(...) } with (...)` block at module scope — are NOT
 * matched here. The `with`-modifier form is rejected at compile time by
 * `assertNoWrappedTopLevelCallbacks` in `liftCallbackBlocks` (it would
 * silently drop on resume). The `handle { ... }` form at module scope
 * is a pre-existing top-level-handle limitation that crashes the
 * typescriptBuilder via `StepPathTracker: currentId() called with empty
 * path`. Until top-level `handle` is supported, register top-level
 * callbacks at module scope without a wrapping handler, or move the
 * registration into a node body where the handle machinery works.
 */
function isTopLevelCallbackCall(node) {
    return node.type === "functionCall" && node.functionName === "callback";
}
/** If `node` is a `static x = ...` (optionally wrapped in `with handler`), return its parts. */
function unwrapStaticAssignment(node) {
    if (node.type === "assignment" && node.scope === "static") {
        return { stmt: node };
    }
    if (node.type === "withModifier" &&
        node.statement.type === "assignment" &&
        node.statement.scope === "static") {
        return { stmt: node.statement, handlerName: node.handlerName };
    }
    return null;
}
/** If `node` is a `global x = ...` (optionally wrapped in `with handler`), return its parts. */
function unwrapGlobalAssignment(node) {
    if (node.type === "assignment" && node.scope === "global") {
        return { stmt: node };
    }
    if (node.type === "withModifier" &&
        node.statement.type === "assignment" &&
        node.statement.scope === "global") {
        return { stmt: node.statement, handlerName: node.handlerName };
    }
    return null;
}
/**
 * Concatenate the per-section outputs into the final generated module,
 * in the canonical order:
 *
 *   preprocess
 *   importStatements
 *   generated imports (template)
 *   generated builtins (template)
 *   tool registrations
 *   type aliases
 *   static-var declarations + __initializeStatic + __getStaticVars (if any)
 *   __initializeGlobals
 *   generated statements
 *   postprocess
 *   __sourceMap export
 */
export function assembleSections(opts) {
    const sections = [];
    sections.push(...opts.preprocess);
    if (opts.importStatements.length > 0) {
        sections.push(ts.statements(opts.importStatements));
    }
    if (opts.generatedImports.trim() !== "") {
        sections.push(ts.raw(opts.generatedImports));
    }
    if (opts.generatedBuiltins.trim() !== "") {
        sections.push(ts.raw(opts.generatedBuiltins));
    }
    if (opts.toolRegistrations.length > 0) {
        sections.push(ts.statements(opts.toolRegistrations));
    }
    for (const alias of opts.typeAliases) {
        sections.push(alias);
    }
    // Registry id is the absolute path; matches the `sourceModuleId`
    // other modules' `awaitModules` carry. Fall back to `opts.moduleId`
    // (cwd-relative) when no plan provided — that path won't get
    // cross-module awaits, but local init still works.
    const registryId = opts.registryModuleId ?? opts.moduleId;
    // Emit `__initializeStatic` whenever this module has either named
    // static vars OR bare `static <expr>` statements (PR 3). Bare static
    // statements need the same Promise-based once-per-process guard, and
    // other modules whose globals depend on this module's statics still
    // need an `__awaitStaticInit(...)` entry to wait on — even if the
    // dependency is only a side effect, not a value.
    if (opts.staticVarNames.size > 0 ||
        opts.staticInitStatements.length > 0) {
        sections.push(...buildStaticVarSetup(opts));
        sections.push(ts.raw(`__registerStaticInit(${JSON.stringify(displayModuleId(registryId))}, __initializeStatic);`));
    }
    sections.push(buildInitializeGlobalsFn(opts));
    // Same idea for globals init.
    sections.push(ts.raw(`__registerGlobalsInit(${JSON.stringify(displayModuleId(registryId))}, __initializeGlobals);`));
    sections.push(buildRegisterTopLevelCallbacksFn(opts));
    sections.push(ts.statements(opts.generatedStatements));
    if (opts.postprocess.length > 0) {
        sections.push(ts.statements(opts.postprocess));
    }
    sections.push(ts.raw(`export const __sourceMap = ${opts.sourceMapJson};`));
    return ts.statements(sections);
}
/**
 * Build the leading banner-comment block for `__initializeStatic` /
 * `__initializeGlobals`. Surfaces, in two human-readable lines, the
 * plan-driven decisions baked into the function body: which other
 * modules' init must complete before this one (cross-module awaits),
 * and the order this module's own decls are assigned in (local
 * order). Both lists are shown cwd-relative for readability.
 *
 * Lets a reviewer skim a generated module and see exactly what the
 * topsort decided without re-deriving it from the body.
 */
function buildInitBanner(phase, localOrder, awaitModules) {
    if (localOrder.length === 0 && awaitModules.length === 0)
        return [];
    const awaitsLine = awaitModules.length > 0
        ? awaitModules.map(displayModuleId).join(", ")
        : "(none)";
    const localLine = localOrder.length > 0 ? localOrder.join(" → ") : "(none)";
    return [
        ts.comment(`Init plan (${phase} phase):`),
        ts.comment(`  awaits (cross-module): ${awaitsLine}`),
        ts.comment(`  local order:           ${localLine}`),
    ];
}
/**
 * Emit:
 *   let __staticInitPromise = null
 *   [export] let x        ←─── one per static var
 *   async function __initializeStatic(__ctx) {
 *     if (__staticInitPromise) return __staticInitPromise
 *     __staticInitPromise = (async () => { …staticInitStatements })()
 *     return __staticInitPromise
 *   }
 *   function __getStaticVars() { return { x, … } }
 *   __globalCtx.getStaticVars = __getStaticVars;
 */
function buildStaticVarSetup(opts) {
    const out = [];
    // Initialize each static `let` to the sentinel `__UNINIT_STATIC`
    // so that any read of the variable before its initializer has run
    // is caught by the `__readStatic` wrapper (emitted around static
    // reads by the pretty-printer). Without this initializer the
    // binding would be `undefined`, which collides with legitimate user
    // values and produced silent bugs like `fooStatic = barStatic + "!"`
    // evaluating to `"undefined!"` when bar's init had not yet run.
    const sentinel = ts.id("__UNINIT_STATIC");
    const staticLetDecls = [...opts.staticVarNames].map((name) => opts.exportedStaticVarNames.has(name)
        ? ts.export(ts.letDecl(name, sentinel))
        : ts.letDecl(name, sentinel));
    out.push(ts.statements([
        ts.letDecl("__staticInitPromise", ts.raw("null")),
        ...staticLetDecls,
    ]));
    // Promise-based guard: concurrent callers await the same init promise.
    // Body starts with awaits on any modules whose statics are
    // referenced by this module's static initializers — topsort guarantees
    // those dependencies have no cycles, so the await chain terminates.
    // Local assignments then run in topsort-order (sectionAssembler's
    // `reorderTagged` did that ordering during partition).
    const awaitPrelude = (opts.staticAwaitModules ?? []).map((m) => ts.raw(`await __awaitStaticInit(${JSON.stringify(displayModuleId(m.sourceModuleId))}, __ctx);`));
    out.push(ts.functionDecl("__initializeStatic", [{ name: "__ctx" }], ts.statements([
        ...buildInitBanner("static", opts.staticLocalOrder ?? [], (opts.staticAwaitModules ?? []).map((m) => m.sourceModuleId)),
        ts.if(ts.id("__staticInitPromise"), ts.return(ts.id("__staticInitPromise"))),
        ts.assign(ts.id("__staticInitPromise"), ts.iife({
            async: true,
            body: [...awaitPrelude, ...opts.staticInitStatements],
        })),
        ts.return(ts.id("__staticInitPromise")),
    ]), { async: true }));
    const staticVarObj = ts.obj([...opts.staticVarNames].map((n) => ts.set(n, ts.id(n))));
    out.push(ts.statements([
        ts.functionDecl("__getStaticVars", [], ts.return(staticVarObj)),
        ts.assign(ts.prop(ts.runtime.globalCtx, "getStaticVars"), ts.id("__getStaticVars")),
    ]));
    return out;
}
/**
 * Emit:
 *   async function __initializeGlobals(__ctx) {
 *     __ctx.globals.markInitialized(moduleId)
 *     [await __initializeStatic(__ctx)]  ← only if there are static vars
 *     [await __ctx.writeStaticStateToTrace(__globalCtx.getStaticVars())]
 *     …globalInitStatements
 *   }
 */
function buildInitializeGlobalsFn(opts) {
    // Inside this function body `__ctx` is the parameter (no ALS frame
    // installed by the caller), so every receiver must be `ts.id("__ctx")`
    // — NOT `ts.runtime.ctx` (which is now the `__ctx()` accessor and
    // would read from ALS instead of the parameter).
    const ctxParam = ts.id("__ctx");
    const body = [
        ...buildInitBanner("global", opts.globalLocalOrder ?? [], (opts.globalAwaitModules ?? []).map((m) => m.sourceModuleId)),
        // Idempotence guard: Phase B runs once per execCtx. The closure-
        // wide bootstrap loop in `runNode` iterates every registered
        // module, and other modules' `__initializeGlobals` prelude also
        // `await __awaitGlobalsInit(...)` for any module they depend on
        // — without this guard those overlapping calls would each re-run
        // the body, double-executing top-level globals and bare
        // statements. Same per-execCtx flag that the per-function lazy
        // guard already checks; safe and cheap.
        ts.if(ts.methodCall(ts.prop(ctxParam, "globals"), "isInitialized", [
            ts.str(opts.moduleId),
        ]), ts.return()),
        // Mark this module as initialized BEFORE running init statements.
        // This prevents infinite recursion when a global init expression
        // calls a function defined in the same module (which would trigger
        // __initializeGlobals again via the isInitialized check).
        ts.methodCall(ts.prop(ctxParam, "globals"), "markInitialized", [ts.str(opts.moduleId)]),
    ];
    // Always run this module's static init before its global init when
    // the module emitted an `__initializeStatic` at all (named static
    // vars OR bare `static <expr>` statements). The Promise-based guard
    // in `__initializeStatic` makes the call cheap on the second-plus
    // run; the first run is what bare static statements rely on for
    // their once-per-process semantics.
    if (opts.staticVarNames.size > 0 ||
        opts.staticInitStatements.length > 0) {
        body.push(ts.awaitCall(ts.id("__initializeStatic"), [ctxParam]));
    }
    if (opts.staticVarNames.size > 0) {
        body.push(ts.awaitMethodCall(ctxParam, "writeStaticStateToTrace", [
            ts.methodCall(ts.runtime.globalCtx, "getStaticVars"),
        ]));
    }
    // Await per-phase imported modules' globals init. Same logic as the
    // static prelude — topsort guarantees no cycles among the awaits.
    for (const m of opts.globalAwaitModules ?? []) {
        body.push(ts.raw(`await __awaitGlobalsInit(${JSON.stringify(displayModuleId(m.sourceModuleId))}, __ctx);`));
    }
    body.push(...opts.globalInitStatements);
    return ts.functionDecl("__initializeGlobals", [{ name: "__ctx" }], ts.statements(body), { async: true });
}
/**
 * Emit:
 *   async function __registerTopLevelCallbacks(__ctx) {
 *     __ctx.topLevelCallbacks = []
 *     …topLevelCallbackStatements
 *   }
 *
 * The body clears any previously-registered top-level callbacks before
 * re-registering so this is safe to call on every fresh run and every
 * resume without accumulating duplicate registrations across restores.
 * Always emitted (even when empty) so the calling runtime can call it
 * unconditionally — generated modules don't have to detect the no-op
 * case.
 */
function buildRegisterTopLevelCallbacksFn(opts) {
    // Same parameter-context rule as buildInitializeGlobalsFn — `__ctx`
    // here is the function parameter, not an ALS-installed value.
    const body = [
        ts.assign(ts.prop(ts.id("__ctx"), "topLevelCallbacks"), ts.arr([])),
        ...opts.topLevelCallbackStatements,
    ];
    return ts.functionDecl("__registerTopLevelCallbacks", [{ name: "__ctx" }], ts.statements(body), { async: true });
}
