import { jsonStringify } from '@aztec/foundation/json-rpc';
import { createLogger } from '@aztec/foundation/log';
import { promiseWithResolvers } from '@aztec/foundation/promise';
import { sleep } from '@aztec/foundation/sleep';
import { NodeConnector, TransportClient } from '@aztec/foundation/transport';
import { Worker } from 'worker_threads';
import { WorkerWalletSchema } from './worker_wallet_schema.js';
const log = createLogger('e2e:test-wallet:worker-wallet');
const WORKER_READY_TIMEOUT_MS = 120_000;
/**
 * Wallet implementation that offloads all work to a worker thread.
 * Implements the Wallet interface by proxying calls over a transport layer
 * using JSON serialization with Zod schema parsing on both ends.
 */ export class WorkerWallet {
    worker;
    client;
    constructor(worker, client){
        this.worker = worker;
        this.client = client;
    }
    /**
   * Creates a WorkerWallet by spawning a worker thread that creates a TestWallet internally.
   * @param nodeUrl - URL of the Aztec node to connect to.
   * @param pxeConfig - Optional PXE configuration overrides.
   * @returns A WorkerWallet ready to use.
   */ static async create(nodeUrl, pxeConfig) {
        // replace stc/ with dest/ so the wallet works in Jest tests
        const workerUrl = new URL('./wallet_worker_script.js', import.meta.url);
        workerUrl.pathname = workerUrl.pathname.replace('/src/', '/dest/');
        // remove JEST_WORKER_ID so the worker uses pino-pretty transport instead of Jest's raw output.
        const { JEST_WORKER_ID: _, ...parentEnv } = process.env;
        const worker = new Worker(workerUrl, {
            workerData: {
                nodeUrl,
                pxeConfig
            },
            env: {
                ...parentEnv,
                ...process.stderr.isTTY || process.env.FORCE_COLOR ? {
                    FORCE_COLOR: '1'
                } : {},
                LOG_LEVEL: process.env.WORKER_LOG_LEVEL ?? 'warn'
            }
        });
        const connector = new NodeConnector(worker);
        const client = new TransportClient(connector);
        await client.open();
        const wallet = new WorkerWallet(worker, client);
        const { promise: workerDied, reject: rejectWorkerDied } = promiseWithResolvers();
        // reject if the worker exits or errors before the warmup completes.
        const onError = (err)=>{
            worker.off('exit', onExit);
            rejectWorkerDied(new Error(`Worker wallet thread error: ${err.message}`));
        };
        const onExit = (code)=>{
            worker.off('error', onError);
            rejectWorkerDied(new Error(`Worker wallet thread exited with code ${code} before becoming ready`));
        };
        worker.once('error', onError);
        worker.once('exit', onExit);
        const timeout = sleep(WORKER_READY_TIMEOUT_MS).then(()=>{
            throw new Error(`Worker wallet creation timed out after ${WORKER_READY_TIMEOUT_MS / 1000}s`);
        });
        try {
            // wait for worker wallet to start
            await Promise.race([
                wallet.getChainInfo(),
                workerDied,
                timeout
            ]);
        } catch (err) {
            log.error('Worker wallet creation failed, cleaning up', {
                error: String(err)
            });
            client.close();
            await worker.terminate();
            throw err;
        } finally{
            worker.off('error', onError);
            worker.off('exit', onExit);
        }
        return wallet;
    }
    async callRaw(fn, ...args) {
        const argsJson = jsonStringify(args);
        return await this.client.request({
            fn,
            args: argsJson
        });
    }
    async call(fn, ...args) {
        const resultJson = await this.callRaw(fn, ...args);
        const methodSchema = WorkerWalletSchema[fn];
        return methodSchema.returnType().parseAsync(JSON.parse(resultJson));
    }
    getChainInfo() {
        return this.call('getChainInfo');
    }
    getContractMetadata(address) {
        return this.call('getContractMetadata', address);
    }
    getContractClassMetadata(id) {
        return this.call('getContractClassMetadata', id);
    }
    getPrivateEvents(eventMetadata, eventFilter) {
        return this.call('getPrivateEvents', eventMetadata, eventFilter);
    }
    registerSender(address, alias) {
        return this.call('registerSender', address, alias);
    }
    getAddressBook() {
        return this.call('getAddressBook');
    }
    getAccounts() {
        return this.call('getAccounts');
    }
    registerContract(instance, artifact, secretKey) {
        return this.call('registerContract', instance, artifact, secretKey);
    }
    simulateTx(exec, opts) {
        return this.call('simulateTx', exec, opts);
    }
    executeUtility(call, opts) {
        return this.call('executeUtility', call, opts);
    }
    profileTx(exec, opts) {
        return this.call('profileTx', exec, opts);
    }
    sendTx(exec, opts) {
        return this.call('sendTx', exec, opts);
    }
    proveTx(exec, opts) {
        return this.call('proveTx', exec, opts);
    }
    /** Registers an account inside the worker's TestWallet, populating its accounts map. */ registerAccount(secret, salt) {
        return this.call('registerAccount', secret, salt);
    }
    createAuthWit(from, messageHashOrIntent) {
        return this.call('createAuthWit', from, messageHashOrIntent);
    }
    requestCapabilities(manifest) {
        return this.call('requestCapabilities', manifest);
    }
    batch(methods) {
        return this.call('batch', methods);
    }
    /** Shuts down the worker thread and closes the transport. */ async stop() {
        this.client.close();
        await this.worker.terminate();
    }
}
