/**
 * Hand-encoded WASM fixture builder for the WASI-stdio variant of the
 * bridge. Produces a validator-clean module that imports the WASI
 * snapshot-preview1 surface the bridge expects (`fd_read`, `fd_write`,
 * `environ_get`, `environ_sizes_get`, `proc_exit`) and exports the
 * required `memory` + `_start` pair.
 *
 * Per-call protocol: the guest's `_start` is invoked once per envelope
 * exchange by the host. Each invocation:
 *   1. Calls `fd_read` to drain stdin into a scratch buffer at offset
 *      0 (the host queues exactly one frame before calling `_start`,
 *      so the read returns the full frame).
 *   2. Optionally executes `unreachable` (trap-on-call-N variants).
 *   3. Picks one of up to four pre-staged reply frames based on a
 *      mutable global call counter.
 *   4. Calls `fd_write(fd=1, ...)` with the chosen reply frame.
 *   5. Increments the call counter and returns.
 *
 * The four reply slots map to the protocol steps in order:
 *   slot 0 → routes reply
 *   slot 1 → start ok
 *   slot 2 → request reply (or trap, per `trapOnRequest`)
 *   slot 3 → dispose ok
 *
 * Memory layout (single 1-page memory, 64 KiB):
 *   [   0,   16) — fd_read iov[0] = { ptr=128, len=4096 }
 *                  fd_write iov[0] = { ptr=?, len=? } (overwritten)
 *   [  16,   20) — nread / nwritten scratch
 *   [ 128, 4224) — stdin scratch buffer
 *   [5000, 5004) — call counter (global stash)
 *   [8192, 8192 + reply0Bytes.length) — routes reply
 *   [12288, 12288 + reply1Bytes.length) — start ok
 *   [16384, 16384 + reply2Bytes.length) — request reply
 *   [20480, 20480 + reply3Bytes.length) — dispose ok
 *
 * Globals:
 *   global 0 (mut i32) — call counter (default 0).
 */

import {
    concat,
    encodeI32,
    encodeName,
    encodeU32,
    section,
    vec,
} from "../leb128.js";
import {
    BLOCK_TYPE_EMPTY,
    funcBody,
    OP_BLOCK,
    OP_CALL,
    OP_DROP,
    OP_END,
    OP_I32_ADD,
    OP_I32_CONST,
    OP_I32_EQZ,
    OP_I32_STORE,
    OP_I32_SUB,
    OP_UNREACHABLE,
    opBrIf,
    opGlobalGet,
    opGlobalSet,
    opI32Const,
    opLocalGet,
    opLocalSet,
    VT_I32,
} from "../wasm-opcodes.js";

const STDIN_PTR = 128;
const STDIN_LEN = 4096;
const NREAD_PTR = 16;

// Each reply slot has 4096 bytes of headroom.
const REPLY_SLOTS = [8192, 12288, 16384, 20480] as const;

const MAGIC = new Uint8Array([0x00, 0x61, 0x73, 0x6d]);
const VERSION = new Uint8Array([0x01, 0x00, 0x00, 0x00]);

const ENCODER = new TextEncoder();

/** Configurable behaviour for the WASI-stdio fixture. */
export interface BuildWasiStdioFixtureOptions {
    /** JSON-encoded routes reply payload. Default: 1 route. */
    readonly routesReply?: string;
    /** JSON-encoded ok-shaped reply for `{kind:"start"}`. Default: `{"kind":"ok"}`. */
    readonly startReply?: string;
    /** JSON-encoded request reply (PluginEnvelopeOut). Default: 200 ok body. */
    readonly requestReply?: string;
    /** JSON-encoded reply for `{kind:"dispose"}`. Default: `{"kind":"ok"}`. */
    readonly disposeReply?: string;
    /** If set, `_start` invokes `unreachable` on the given (0-indexed) call. */
    readonly trapOnCall?: 0 | 1 | 2 | 3;
    /** Drop one of the required imports — exercises the ABI validator. */
    readonly dropImports?: readonly (
        | "fd_read"
        | "fd_write"
        | "environ_get"
        | "environ_sizes_get"
        | "proc_exit"
    )[];
    /** Drop a required export — exercises the ABI validator. */
    readonly dropExports?: readonly ("memory" | "_start")[];
    /** Emit a malformed (non-JSON) reply on the request slot. */
    readonly requestReplyRaw?: Uint8Array;
    /**
     * Make `_start` return without writing any frame to fd 1 — used by
     * the runner's STDIO_TIMEOUT test.
     */
    readonly silentStart?: boolean;
    /**
     * Add an extra import from an unknown module — used by the
     * runner's instantiation_failed test.
     */
    readonly addExtraUnknownImport?: boolean;
    /**
     * Have `_start` invoke `environ_sizes_get` (via memoryRef) once
     * per call. Used to drive the runner's memoryRef thunk coverage.
     */
    readonly callEnvironSizesGet?: boolean;
}

const DEFAULT_ROUTES_REPLY = JSON.stringify({
    kind: "routes",
    routes: [
        { method: "GET", path: "/hello" },
        { method: "POST", path: "/echo" },
    ],
});
const DEFAULT_OK_REPLY = JSON.stringify({ kind: "ok" });
const DEFAULT_REQUEST_REPLY = JSON.stringify({
    kind: "response",
    response: {
        status: 200,
        headers: [["content-type", "text/plain"]],
        bodyBase64: btoa("ok")
            .replace(/\+/g, "-")
            .replace(/\//g, "_")
            .replace(/=+$/, ""),
    },
});

/** Build a length-prefixed frame `[u32 BE len][JSON bytes]`. */
function frame(json: string | Uint8Array): Uint8Array {
    const payload = typeof json === "string" ? ENCODER.encode(json) : json;
    const out = new Uint8Array(4 + payload.length);
    new DataView(out.buffer).setUint32(0, payload.length, false);
    out.set(payload, 4);
    return out;
}

/**
 * Build the WASI-stdio fixture WASM module. The returned bytes are
 * accepted by `WebAssembly.compile` and instantiate against an imports
 * object exposing `wasi_snapshot_preview1.{fd_read, fd_write,
 * environ_get, environ_sizes_get, proc_exit}`.
 */
export function buildWasiStdioFixtureWasm(
    options: BuildWasiStdioFixtureOptions = {}
): Uint8Array {
    const replies: ReadonlyArray<Uint8Array> = [
        frame(options.routesReply ?? DEFAULT_ROUTES_REPLY),
        frame(options.startReply ?? DEFAULT_OK_REPLY),
        options.requestReplyRaw !== undefined
            ? frame(options.requestReplyRaw)
            : frame(options.requestReply ?? DEFAULT_REQUEST_REPLY),
        frame(options.disposeReply ?? DEFAULT_OK_REPLY),
    ];
    for (let i = 0; i < replies.length; i++) {
        const r = replies[i] as Uint8Array;
        if (r.length > 4096) {
            throw new RangeError(`reply slot ${i} payload too large`);
        }
    }
    const dropImports = new Set(options.dropImports ?? []);
    const dropExports = new Set(options.dropExports ?? []);

    // Type section.
    // 0: (i32, i32, i32, i32) → i32 — fd_read / fd_write signature.
    // 1: (i32, i32) → i32 — environ_get / environ_sizes_get.
    // 2: (i32) → ()        — proc_exit.
    // 3: () → ()           — _start.
    const typeSection = section(
        0x01,
        vec([
            new Uint8Array([
                0x60,
                0x04,
                VT_I32,
                VT_I32,
                VT_I32,
                VT_I32,
                0x01,
                VT_I32,
            ]),
            new Uint8Array([0x60, 0x02, VT_I32, VT_I32, 0x01, VT_I32]),
            new Uint8Array([0x60, 0x01, VT_I32, 0x00]),
            new Uint8Array([0x60, 0x00, 0x00]),
        ])
    );

    // Import section.
    const importEntries: Uint8Array[] = [];
    const pushImport = (name: string, typeIdx: number): void => {
        if (dropImports.has(name as never)) return;
        importEntries.push(
            concat(
                encodeName("wasi_snapshot_preview1"),
                encodeName(name),
                new Uint8Array([0x00]), // kind = function
                encodeU32(typeIdx)
            )
        );
    };
    pushImport("fd_read", 0);
    pushImport("fd_write", 0);
    pushImport("environ_get", 1);
    pushImport("environ_sizes_get", 1);
    pushImport("proc_exit", 2);
    if (options.addExtraUnknownImport === true) {
        importEntries.push(
            concat(
                encodeName("unknown_module"),
                encodeName("missing_fn"),
                new Uint8Array([0x00]),
                encodeU32(2)
            )
        );
    }
    const importedCount = importEntries.length;
    const importSection = section(0x02, vec(importEntries));

    // Function section — one defined function: `_start` (type 3).
    const funcSection = section(0x03, vec([encodeU32(3)]));

    // Memory section — 1 page (64 KiB).
    const memorySection = section(0x05, vec([new Uint8Array([0x00, 0x01])]));

    // Global section — global 0 (mut i32) = 0.
    const globalSection = section(
        0x06,
        vec([
            concat(
                new Uint8Array([VT_I32, 0x01]),
                new Uint8Array([OP_I32_CONST]),
                encodeI32(0),
                new Uint8Array([OP_END])
            ),
        ])
    );

    // Export section.
    const exportEntries: Uint8Array[] = [];
    if (!dropExports.has("memory")) {
        exportEntries.push(
            concat(encodeName("memory"), new Uint8Array([0x02, 0x00]))
        );
    }
    if (!dropExports.has("_start")) {
        // _start is the only defined function — its index is `importedCount + 0`.
        exportEntries.push(
            concat(
                encodeName("_start"),
                new Uint8Array([0x00]),
                encodeU32(importedCount)
            )
        );
    }
    const exportSection = section(0x07, vec(exportEntries));

    // Code section — _start body.
    const startBody = buildStartBody({
        importedCount,
        replies,
        trapOnCall: options.trapOnCall,
        haveFdRead: !dropImports.has("fd_read"),
        haveFdWrite:
            !dropImports.has("fd_write") && options.silentStart !== true,
        callEnvironSizesGet: options.callEnvironSizesGet === true,
    });
    const codeSection = section(0x0a, vec([startBody]));

    // Data section — pre-stage each reply frame at its slot offset.
    const dataSegments: Uint8Array[] = [];
    for (let i = 0; i < replies.length; i++) {
        const slot = REPLY_SLOTS[i] as number;
        dataSegments.push(
            concat(
                encodeU32(0),
                new Uint8Array([OP_I32_CONST]),
                encodeI32(slot),
                new Uint8Array([OP_END]),
                encodeU32((replies[i] as Uint8Array).length),
                replies[i] as Uint8Array
            )
        );
    }
    const dataSection = section(0x0b, vec(dataSegments));

    return concat(
        MAGIC,
        VERSION,
        typeSection,
        importSection,
        funcSection,
        memorySection,
        globalSection,
        exportSection,
        codeSection,
        dataSection
    );
}

interface StartBodyOptions {
    readonly importedCount: number;
    readonly replies: ReadonlyArray<Uint8Array>;
    readonly trapOnCall: number | undefined;
    readonly haveFdRead: boolean;
    readonly haveFdWrite: boolean;
    readonly callEnvironSizesGet: boolean;
}

/**
 * Build the `_start` function body. Locals: `local 0` = call index
 * (cached from global). The function logic is:
 *
 *   local0 = global0
 *   if (trapOnCall === local0) unreachable
 *
 *   // 1. fd_read(fd=0, iovs=stdinIovs, iovsLen=1, nreadPtr=NREAD_PTR)
 *   // 2. select reply slot by local0 (4-arm switch)
 *   // 3. fd_write(fd=1, iovs=writeIovs, iovsLen=1, nwrittenPtr=NREAD_PTR)
 *   //
 *   // global0 = local0 + 1
 *   //
 * The fixture writes the WHOLE pre-staged frame on every call. We use
 * scratch memory cells to set up iovec records before each WASI call.
 */
function buildStartBody(opts: StartBodyOptions): Uint8Array {
    const { importedCount, replies, trapOnCall, haveFdRead, haveFdWrite } =
        opts;

    const _fdReadIdx = importedCount > 0 ? 0 : -1; // first imported = fd_read
    const fdWriteIdx = 1;
    void fdWriteIdx;

    // Iovec scratch at memory[0..16]:
    //   [ 0, 4) = ptr (stdinPtr=128 for fd_read; replySlot for fd_write)
    //   [ 4, 8) = len (STDIN_LEN for fd_read; reply.length for fd_write)
    //
    // NREAD_PTR / NWRITTEN_PTR scratch at memory[16..20].

    // 1) fd_read iovec setup
    const setupReadIov = concat(
        opI32Const(0),
        opI32Const(STDIN_PTR),
        new Uint8Array([OP_I32_STORE]),
        encodeU32(2),
        encodeU32(0),
        opI32Const(4),
        opI32Const(STDIN_LEN),
        new Uint8Array([OP_I32_STORE]),
        encodeU32(2),
        encodeU32(0)
    );

    // fd_read call:  result = fd_read(0, 0, 1, NREAD_PTR); drop;
    const callFdRead = concat(
        opI32Const(0), // fd = 0 (stdin)
        opI32Const(0), // iovs ptr = 0
        opI32Const(1), // iovs len
        opI32Const(NREAD_PTR), // nread ptr
        new Uint8Array([OP_CALL]),
        encodeU32(0), // imported fd_read at function index 0 (if not dropped)
        new Uint8Array([OP_DROP])
    );

    // 4-arm dispatch: local 0 holds the call index.
    // We use chained br_if blocks: if local0 === 0 → load slot 0 + jump out, etc.

    const dispatchBlocks: Uint8Array[] = [];
    for (let slot = 0; slot < 4; slot++) {
        const replyLen = replies[slot]?.length ?? 0;
        const replyPtr = REPLY_SLOTS[slot] as number;
        // Block: if local0 !== slot, br 0 (skip); else set iovec + br outer.
        // Encoded as a single block with a conditional skip.
        // We'll encode each arm as: (block (local.get 0) (i32.const slot) (i32.sub) (br_if 0) <setup+write> (br outer))
        // outer block depth: we'll wrap the whole switch in one outer block (depth 1).
        const arm = concat(
            new Uint8Array([OP_BLOCK, BLOCK_TYPE_EMPTY]),
            opLocalGet(0),
            opI32Const(slot),
            new Uint8Array([OP_I32_SUB]),
            opBrIf(0), // if local0 - slot !== 0 (truthy) → break inner block
            // Setup write iovec: mem[0..4] = replyPtr; mem[4..8] = replyLen
            opI32Const(0),
            opI32Const(replyPtr),
            new Uint8Array([OP_I32_STORE]),
            encodeU32(2),
            encodeU32(0),
            opI32Const(4),
            opI32Const(replyLen),
            new Uint8Array([OP_I32_STORE]),
            encodeU32(2),
            encodeU32(0),
            // Call fd_write(1, 0, 1, NREAD_PTR); drop.
            opI32Const(1),
            opI32Const(0),
            opI32Const(1),
            opI32Const(NREAD_PTR),
            new Uint8Array([OP_CALL]),
            encodeU32(1), // imported fd_write at function index 1
            new Uint8Array([OP_DROP]),
            // Break the OUTER switch block (depth 1: we're inside arm-block inside outer).
            new Uint8Array([0x0c, 0x01]), // OP_BR depth=1
            new Uint8Array([OP_END])
        );
        dispatchBlocks.push(arm);
    }

    const switchBody = concat(
        new Uint8Array([OP_BLOCK, BLOCK_TYPE_EMPTY]),
        ...dispatchBlocks,
        new Uint8Array([OP_END])
    );

    // Optional trap injection: if trapOnCall === local0, hit unreachable.
    let trapGuard: Uint8Array = new Uint8Array();
    if (trapOnCall !== undefined) {
        trapGuard = concat(
            new Uint8Array([OP_BLOCK, BLOCK_TYPE_EMPTY]),
            opLocalGet(0),
            opI32Const(trapOnCall),
            new Uint8Array([OP_I32_SUB]),
            new Uint8Array([OP_I32_EQZ]),
            new Uint8Array([OP_I32_EQZ]),
            opBrIf(0), // skip unreachable if !== trapOnCall
            new Uint8Array([OP_UNREACHABLE]),
            new Uint8Array([OP_END])
        );
    }

    // Increment the call counter EARLY so a trap on call N does not
    // leave global0 stuck at N (which would re-trap dispose).
    const incrementCounter = concat(
        opLocalGet(0),
        opI32Const(1),
        new Uint8Array([OP_I32_ADD]),
        opGlobalSet(0)
    );

    // Optional: invoke environ_sizes_get(20, 24) — pushes the
    // memoryRef thunk in the runner. environ_sizes_get is import #3
    // (fd_read=0, fd_write=1, environ_get=2, environ_sizes_get=3).
    const optionalEnvironCall = opts.callEnvironSizesGet
        ? concat(
              opI32Const(20),
              opI32Const(24),
              new Uint8Array([OP_CALL]),
              encodeU32(3),
              new Uint8Array([OP_DROP])
          )
        : new Uint8Array();

    // Compose: load counter; increment; trap guard; fd_read; switch.
    const loadCounter = concat(opGlobalGet(0), opLocalSet(0));

    const body = concat(
        loadCounter,
        incrementCounter,
        trapGuard,
        haveFdRead ? concat(setupReadIov, callFdRead) : new Uint8Array(),
        optionalEnvironCall,
        haveFdWrite ? switchBody : new Uint8Array()
    );

    return funcBody([{ count: 1, type: VT_I32 }], body);
}

// Re-export the slot constants for tests that need to introspect memory.
export const WASI_STDIO_FIXTURE_REPLY_SLOTS = REPLY_SLOTS;
export const WASI_STDIO_FIXTURE_STDIN_PTR = STDIN_PTR;