// Copyright (C) 2026 Doan Luu Duc Tai <doan@luuduct.ai>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

// @license magnet:?xt=urn:btih:0b31508aeb0634b347b8270c7bee4d411b5d4109&dn=agpl-3.0.txt AGPL-3.0

async function execute(bytes: ArrayBuffer, input: string): Promise<StackChat.Execution> {
    let memory: WebAssembly.Memory | undefined;
    let output = '';
    let errors = '';
    const stdoutDecoder = new TextDecoder('utf-8', { fatal: true });
    const stderrDecoder = new TextDecoder('utf-8', { fatal: true });
    const stdin = new TextEncoder().encode(input);
    let position = 0;
    let randomCalls = 0;
    let randomBytes = 0;
    const memoryBuffer = (): ArrayBuffer => {
        if (!memory) throw new Error('WASI memory is unavailable');
        return memory.buffer;
    };
    const data = (): DataView => new DataView(memoryBuffer());
    const put32 = (address: number, value: number) => data().setUint32(address, value, true);
    const noEntries = (count: number, size: number) => { put32(count, 0); put32(size, 0); return 0; };
    class Exit extends Error {
        readonly code: number;
        constructor(code: number) { super(`WASI exit ${code}`); this.code = code; }
    }
    const wasi: Record<string, WebAssembly.ImportValue> = {
        args_sizes_get: noEntries,
        environ_sizes_get: noEntries,
        args_get: () => 0,
        environ_get: () => 0,
        random_get(address: number, length: number) {
            const buffer = new Uint8Array(memoryBuffer(), address, length);
            for (let offset = 0; offset < length; offset += 65536) {
                crypto.getRandomValues(buffer.subarray(offset, Math.min(offset + 65536, length)));
            }
            randomCalls += 1;
            randomBytes += length;
            return 0;
        },
        clock_time_get(clock: number, _precision: bigint, address: number) {
            if (clock !== 0 && clock !== 1) return 52;
            const milliseconds = clock === 0 ? Date.now() : performance.now();
            data().setBigUint64(address, BigInt(Math.floor(milliseconds * 1000000)), true);
            return 0;
        },
        clock_res_get(_clock: number, address: number) { data().setBigUint64(address, 1000000n, true); return 0; },
        fd_write(fd: number, vectors: number, count: number, written: number) {
            if (fd !== 1 && fd !== 2) return 8;
            let total = 0;
            for (let index = 0; index < count; index += 1) {
                const address = data().getUint32(vectors + index * 8, true);
                const length = data().getUint32(vectors + index * 8 + 4, true);
                const text = (fd === 1 ? stdoutDecoder : stderrDecoder).decode(new Uint8Array(memoryBuffer(), address, length), { stream: true });
                if (fd === 1) output += text; else errors += text;
                total += length;
                if (output.length + errors.length > 1048576) throw new Error('Probe output exceeded its limit');
            }
            put32(written, total);
            return 0;
        },
        fd_read(fd: number, vectors: number, count: number, read: number) {
            if (fd !== 0) return 8;
            let total = 0;
            for (let index = 0; index < count && position < stdin.length; index += 1) {
                const address = data().getUint32(vectors + index * 8, true);
                const capacity = data().getUint32(vectors + index * 8 + 4, true);
                const length = Math.min(capacity, stdin.length - position);
                new Uint8Array(memoryBuffer(), address, length).set(stdin.subarray(position, position + length));
                position += length;
                total += length;
            }
            put32(read, total);
            return 0;
        },
        fd_fdstat_get(fd: number, address: number) {
            if (fd > 2) return 8;
            new Uint8Array(memoryBuffer(), address, 24).fill(0);
            data().setUint8(address, 2);
            return 0;
        },
        fd_fdstat_set_flags: (fd: number) => fd <= 2 ? 0 : 8,
        fd_prestat_get: () => 8,
        fd_prestat_dir_name: () => 8,
        fd_close: (fd: number) => fd <= 2 ? 0 : 8,
        fd_seek: () => 70,
        fd_tell: () => 70,
        fd_filestat_get: () => 8,
        path_open: () => 76,
        sched_yield: () => 0,
        proc_exit(code: number) { throw new Exit(code); },
    };
    for (const name of [
        'fd_readdir', 'fd_advise', 'fd_allocate', 'fd_datasync', 'fd_sync',
        'fd_filestat_set_size', 'fd_filestat_set_times', 'fd_fdstat_set_rights',
        'path_filestat_get', 'path_filestat_set_times', 'path_readlink', 'path_unlink_file',
        'path_rename', 'path_create_directory', 'path_remove_directory', 'path_link', 'path_symlink',
        'sock_accept', 'sock_recv', 'sock_send', 'sock_shutdown', 'poll_oneoff',
    ]) wasi[name] = () => 76;
    const module = await WebAssembly.compile(bytes);
    for (const entry of WebAssembly.Module.imports(module)) {
        if (entry.module !== 'wasi_snapshot_preview1' || typeof wasi[entry.name] !== 'function') {
            throw new Error(`Unsupported WASI import: ${entry.module}.${entry.name}`);
        }
    }
    const instance = await WebAssembly.instantiate(module, { wasi_snapshot_preview1: wasi });
    const exportedMemory = instance.exports['memory'];
    const start = instance.exports['_start'];
    if (!(exportedMemory instanceof WebAssembly.Memory) || typeof start !== 'function') throw new Error('Invalid WASI exports');
    memory = exportedMemory;
    try { start(); }
    catch (error) { if (!(error instanceof Exit) || error.code !== 0) throw new Error(`${error instanceof Error ? error.message : 'WASI execution failed'}: ${errors}`); }
    output += stdoutDecoder.decode();
    errors += stderrDecoder.decode();
    stdin.fill(0);
    return { output, randomCalls, randomBytes };
}

async function runProbe(bytes: ArrayBuffer) {
    const result = await execute(bytes, '');
    if (!result.output.startsWith('PASS: backend=OpenSSL ') || result.randomCalls === 0) {
        throw new Error('OpenSSL probe did not pass');
    }
    return result;
}

async function runCommand(bytes: ArrayBuffer, request: StackChat.CryptoRequest): Promise<unknown> {
    const input = JSON.stringify(request);
    if (input.length > 8388608) throw new Error('Crypto request exceeds its limit');
    const result = await execute(bytes, input);
    const response: unknown = JSON.parse(result.output);
    if (typeof response !== 'object' || response === null || Array.isArray(response)) throw new Error('Invalid crypto response');
    if ('error' in response) throw new Error(typeof response.error === 'string' ? response.error : 'Cryptographic operation failed');
    return response;
}

globalThis.stackChatWasi = Object.freeze({ runCommand, runProbe });

// @license-end