import assert from 'node:assert/strict';
import {EventEmitter} from 'node:events';
import {randomUUID} from 'node:crypto';
import test from 'node:test';
import type {Socket} from 'socket.io';
import {RELAY_METHODS, WORKFLOW_ACTIONS} from '../types.ts';
import type {ForwardAck, JsonValue, NodeRelayRequest, RelayClaims, RelayLimits, RelayMethod, RelayRequest, WorkflowRunSnapshot} from '../types.ts';
import {create_session_store} from '../runtime/session_store.ts';
import {DEFAULT_LIMITS} from '../utils/relay_config.ts';
import {encode_token} from '../utils/token_utils.ts';
import {error_ack, success} from '../utils/relay_validation.ts';

class MockSocket extends EventEmitter {
    connected = true;
    conn = {writeBuffer: [] as unknown[]};
    acks = new Map<number, (value: unknown) => void>();
    rooms = new Set<string>();
    requests: NodeRelayRequest[] = [];
    respond: (request: NodeRelayRequest) => unknown = () => success(null);
    private next_ack = 0;
    join(room: string) {this.rooms.add(room);}
    leave(room: string) {this.rooms.delete(room);}
    disconnect() {this.connected = false; this.emit('disconnect');}
    timeout(_ms: number) {return this;}
    emitWithAck(event: string, request: NodeRelayRequest) {
        assert.equal(event, 'relay_request');
        this.requests.push(request);
        return new Promise<unknown>((resolve, reject) => {
            const id = this.next_ack++;
            this.acks.set(id, resolve);
            Promise.resolve().then(() => this.respond(request)).then(value => {
                if (this.acks.delete(id)) resolve(value);
            }, reject);
        });
    }
    get socket(): Socket {return this as unknown as Socket;}
}

async function fixture(t: {after: (fn: () => void) => unknown}, overrides: Partial<RelayLimits> = {},
    web_permissions: RelayMethod[] = [...RELAY_METHODS], node_permissions: RelayMethod[] = [...RELAY_METHODS]) {
    const key = 'mock-session-test-secret-only';
    const scope = {forward_id: 1, forward_session_id: randomUUID(), node_session_id: 'main-session', node_id: 2, bot_id: 3};
    const store = create_session_store({forward_id: 1, key, limits: {...DEFAULT_LIMITS, ...overrides}, on_status() {}});
    t.after(() => store.close());
    store.open({...scope, lease_expires_at: Date.now() + 60_000});
    const grant_id = randomUUID();
    const expires_at = Date.now() + 50_000;
    store.authorize({forward_session_id: scope.forward_session_id, grant_id, user_id: 4, permissions: web_permissions, expires_at});
    const join = async (socket: MockSocket, role: 'node' | 'web', permissions: RelayMethod[]) => {
        const now = Date.now();
        const claims: RelayClaims = {...scope, v: 1, role, permissions, issued_at: now, join_before: now + 10_000,
            expires_at, jti: randomUUID(), ...(role === 'web' ? {grant_id, user_id: 4} : {})};
        store.attach(socket.socket);
        await store.join(socket.socket, {token: encode_token(key, claims)});
    };
    const node = new MockSocket(); const web = new MockSocket();
    await join(node, 'node', node_permissions);
    store.node_ready(node.socket, {forward_session_id: scope.forward_session_id});
    await join(web, 'web', web_permissions);
    const request = async (method: RelayRequest['method'], payload: unknown = {}, request_id = randomUUID()): Promise<ForwardAck<unknown>> => {
        try {return await store.request(web.socket, {method, payload, request_id});}
        catch (error) {return error_ack(error);}
    };
    return {store, node, web, scope, grant_id, join, request};
}
function code(value: ForwardAck<unknown>, expected: string) {
    assert.equal(value.status, 'error');
    if (value.status === 'error') assert.equal(value.data.code, expected);
}
const control = () => ({command_id: randomUUID(), context_generation: randomUUID(), action: 'run_once', workflow_id: 'a'});
const command = (input: ReturnType<typeof control>) => ({command_id: input.command_id, action: input.action, workflow_id: input.workflow_id, state: 'accepted'});

test('mock relay forwards all workflow RPCs with private correlated deadlines and safe ACKs', async t => {
    const f = await fixture(t);
    const input = control(); const run_id = randomUUID();
    const summary: WorkflowRunSnapshot = {run_id, workflow_id: 'a', status: 'completed', steps: 1, active_nodes: [],
        started_at: 0, ended_at: 1, stop_reason: null, error: null, trigger_available: true};
    const detail: WorkflowRunSnapshot = {...summary, trigger: {disabled: false, count: 0, empty: null, nested: {values: [0, false, null]}},
        branches: [{trigger: {}, vars: {}, outputs: {}, inputs: {}, retry_counts: {}, last_result: null}]};
    const trigger_detail = {...detail};
    delete trigger_detail.branches;
    const legacy = {...summary};
    delete legacy.trigger_available;
    const runtime = {source_config_revision: 0, workflows: [] as JsonValue[], commands: [command(input)]};
    const catalog = {catalog_version: 1, nodes: [] as JsonValue[], start_nodes: [] as JsonValue[], events: [] as JsonValue[], queries: [] as JsonValue[],
        common: {node_fields: [] as JsonValue[], policy_fields: [] as JsonValue[], result: {}, edge_input: {}, references: [] as JsonValue[]}};
    const requests: Array<[RelayRequest['method'], unknown, unknown]> = [
        ['workflow.start', {workflow_id: 'a', trigger: {x: 1}}, {status: 'started', run_id}],
        ['workflow.stop', {run_id}, {stopped: true}], ['workflow.snapshot', {}, []],
        ['workflow.snapshot', {}, [summary]], ['workflow.snapshot', {run_id}, detail],
        ['workflow.snapshot', {run_id, include_branches: false}, trigger_detail],
        ['workflow.snapshot', {run_id, include_branches: true}, detail],
        ['workflow.snapshot', {run_id}, legacy],
        ['workflow.snapshot', {run_id}, null], ['workflow.catalog', {}, catalog],
        ['workflow.control', input, command(input)],
        ['workflow.command', {command_id: input.command_id, context_generation: input.context_generation}, command(input)],
        ['workflow.runtime', {}, runtime],
    ];
    let received = 0;
    f.web.on('relay_request', () => received++);
    for (const [method, payload, data] of requests) {
        f.node.respond = () => success(data);
        const before = Date.now();
        assert.deepEqual(await f.request(method, payload), success(data));
        const forwarded = f.node.requests.at(-1)!;
        assert.equal(forwarded.method, method);
        assert.deepEqual(forwarded.payload, payload);
        assert.equal(forwarded.forward_session_id, f.scope.forward_session_id);
        assert.notEqual(forwarded.correlation_id, forwarded.request_id);
        assert.ok(forwarded.deadline_at >= before && forwarded.deadline_at <= Date.now() + DEFAULT_LIMITS.rpc_timeout_ms);
    }
    for (const action of WORKFLOW_ACTIONS) {
        const payload = action === 'stop_run' ? {command_id: randomUUID(), context_generation: randomUUID(), action, run_id}
            : {...control(), action, expected_generation: 1, ...(action === 'run_once' ? {input: {value: true}} : {})};
        const data = {command_id: payload.command_id, action, workflow_id: action === 'stop_run' ? null : 'a', state: 'accepted'};
        f.node.respond = () => success(data);
        assert.deepEqual(await f.request('workflow.control', payload), success(data));
        assert.deepEqual(f.node.requests.at(-1)!.payload, payload);
    }
    assert.equal(received, 0);
    assert.equal(f.node.acks.size, 0);
});

test('mock workflow read/write permissions apply to web grants and Node tokens', async t => {
    for (const [web, node] of [
        [['bot.snapshot'], [...RELAY_METHODS]], [['workflow.control'], [...RELAY_METHODS]],
        [[...RELAY_METHODS], ['bot.snapshot']], [[...RELAY_METHODS], ['workflow.control']],
    ] as RelayMethod[][][]) {
        const f = await fixture(t, {}, web, node);
        for (const [method, payload] of [['workflow.start', {workflow_id: 'a'}], ['workflow.stop', {run_id: randomUUID()}],
            ['workflow.control', control()]] as const) {
            f.node.respond = request => success(method === 'workflow.start' ? {status: 'skipped'} : method === 'workflow.stop' ? {stopped: true}
                : command(request.payload as ReturnType<typeof control>));
            const result = await f.request(method, payload);
            if (web.includes('workflow.control') && node.includes('workflow.control')) assert.equal(result.status, 'success');
            else code(result, 'FORBIDDEN');
        }
        for (const method of ['workflow.catalog', 'workflow.snapshot', 'workflow.command', 'workflow.runtime'] as const) {
            const before = f.node.requests.length;
            const payload = method === 'workflow.command' ? {command_id: randomUUID(), context_generation: randomUUID()} : {};
            await f.request(method, payload);
            assert.equal(f.node.requests.length - before, web.includes('bot.snapshot') && node.includes('bot.snapshot') ? 1 : 0);
        }
    }
});

test('mock execute_function workflow category requires both capabilities including on Node', async t => {
    for (const [web, node] of [
        [['execute_function'], [...RELAY_METHODS]], [['workflow.control'], [...RELAY_METHODS]],
        [[...RELAY_METHODS], ['execute_function']], [[...RELAY_METHODS], ['workflow.control']],
        [['execute_function', 'workflow.control'], ['execute_function', 'workflow.control']],
    ] as RelayMethod[][][]) {
        const f = await fixture(t, {}, web, node);
        f.node.respond = () => success(true);
        for (const category of ['workflow', 'workflow/nested']) {
            const result = await f.request('execute_function', {function: category, name: 'run_once', args: []});
            if ([web, node].every(permissions => permissions.includes('execute_function') && permissions.includes('workflow.control')))
                assert.deepEqual(result, success(true));
            else code(result, 'FORBIDDEN');
        }
        const result = await f.request('execute_function', {function: 'movement', name: 'go', args: []});
        if (web.includes('execute_function') && node.includes('execute_function')) assert.deepEqual(result, success(true));
        else code(result, 'FORBIDDEN');
    }
});

test('mock workflow mutations cannot replay and reads refresh with current Node permission', async t => {
    const f = await fixture(t);
    for (const [method, payload, data] of [
        ['workflow.start', {workflow_id: 'a'}, {status: 'skipped'}], ['workflow.stop', {run_id: randomUUID()}, {stopped: true}],
        ['workflow.control', control(), null],
    ] as const) {
        f.node.respond = request => success(method === 'workflow.control' ? command(request.payload as ReturnType<typeof control>) : data);
        const id = randomUUID();
        assert.equal((await f.request(method, payload, id)).status, 'success');
        const before = f.node.requests.length;
        code(await f.request(method, payload, id), 'CANCELLED');
        assert.equal(f.node.requests.length, before);
    }
    const id = randomUUID(); let revision = 0;
    f.node.respond = () => success({source_config_revision: revision++, workflows: [], commands: []});
    assert.deepEqual(await f.request('workflow.runtime', {}, id), success({source_config_revision: 0, workflows: [], commands: []}));
    assert.deepEqual(await f.request('workflow.runtime', {}, id), success({source_config_revision: 1, workflows: [], commands: []}));
    const replacement = new MockSocket();
    await f.join(replacement, 'node', ['workflow.control']);
    f.store.node_ready(replacement.socket, {forward_session_id: f.scope.forward_session_id});
    code(await f.request('workflow.runtime', {}, id), 'FORBIDDEN');
    assert.equal(replacement.requests.length, 0);
});

test('mock control availability, method rate limits and malformed context prevent dispatch', async t => {
    const f = await fixture(t);
    f.store.set_control_available(false);
    for (const [method, payload] of [['workflow.start', {workflow_id: 'a'}], ['workflow.stop', {run_id: randomUUID()}],
        ['workflow.control', control()]] as const) code(await f.request(method, payload), 'RELAY_UNAVAILABLE');
    f.node.respond = () => success({source_config_revision: 0, workflows: [], commands: []});
    assert.equal((await f.request('workflow.runtime')).status, 'success');
    f.store.set_control_available(true);
    const before = f.node.requests.length;
    code(await f.request('workflow.control', {...control(), context_generation: 'bad'}), 'INVALID_REQUEST');
    code(await f.request('workflow.command', {command_id: randomUUID()}), 'INVALID_REQUEST');
    for (const include_branches of [false, true]) {
        code(await f.request('workflow.snapshot', {include_branches}), 'INVALID_REQUEST');
    }
    for (const include_branches of [null, 0, 1, 'false', 'true', [], {}]) {
        code(await f.request('workflow.snapshot', {run_id: randomUUID(), include_branches}), 'INVALID_REQUEST');
    }
    for (const action of WORKFLOW_ACTIONS) {
        const payload = action === 'stop_run' ? {command_id: randomUUID(), context_generation: randomUUID(), action, run_id: randomUUID()}
            : {...control(), action};
        for (const [field, value] of Object.entries({workflow_id: 'a', run_id: randomUUID(), expected_generation: 1, input: {value: true}})) {
            const allowed = field === 'workflow_id' || field === 'expected_generation' ? action !== 'stop_run'
                : field === 'run_id' ? action === 'stop_run' : action === 'run_once';
            if (!allowed) code(await f.request('workflow.control', {...payload, [field]: value}), 'INVALID_REQUEST');
        }
    }
    assert.equal(f.node.requests.length, before);
    const limited = await fixture(t, {requests_per_second: 1});
    limited.node.respond = f.node.respond;
    // Authentication consumes this connection's first request in the same rate window.
    code(await limited.request('workflow.runtime'), 'RATE_LIMITED');
    assert.equal(limited.node.requests.length, 0);
});

test('mock relay rejects malformed and mismatched workflow ACKs without admitting domain envelope errors', async t => {
    const f = await fixture(t); const input = control();
    for (const data of [{...command(input), command_id: randomUUID()}, {...command(input), action: 'stop'},
        {...command(input), workflow_id: 'other'}, {...command(input), state: 'invalid'}]) {
        f.node.respond = () => success(data);
        code(await f.request('workflow.control', input), 'INVALID_REQUEST');
    }
    f.node.respond = () => success({...command(input), command_id: randomUUID()});
    code(await f.request('workflow.command', {command_id: input.command_id, context_generation: input.context_generation}), 'INVALID_REQUEST');
    f.node.respond = () => success(null);
    code(await f.request('workflow.snapshot'), 'INVALID_REQUEST');
    f.node.respond = () => success([]);
    code(await f.request('workflow.snapshot', {run_id: randomUUID()}), 'INVALID_REQUEST');
    const run: WorkflowRunSnapshot = {run_id: randomUUID(), workflow_id: 'a', status: 'running', steps: 0, active_nodes: [],
        started_at: 0, ended_at: null, stop_reason: null, error: null, trigger_available: true, trigger: {value: false}, branches: []};
    f.node.respond = () => success(run);
    code(await f.request('workflow.snapshot', {run_id: randomUUID()}), 'INVALID_REQUEST');
    code(await f.request('workflow.snapshot'), 'INVALID_REQUEST');
    f.node.respond = () => success([run]);
    code(await f.request('workflow.snapshot', {run_id: run.run_id}), 'INVALID_REQUEST');
    for (const data of [{...run, trigger: {value: 'x'.repeat(65_536)}}, {...run, trigger: []},
        {...run, trigger: {value: NaN}}, {...run, trigger: JSON.parse('{"nested":{"constructor":1}}')},
        {...run, trigger_available: false}]) {
        f.node.respond = () => success(data);
        code(await f.request('workflow.snapshot', {run_id: run.run_id}), 'INVALID_REQUEST');
    }
    f.node.respond = () => ({status: 'error', data: {code: 'WORKFLOW_NOT_FOUND', message: 'missing', retryable: false}});
    code(await f.request('workflow.runtime'), 'INVALID_REQUEST');
    f.node.respond = () => ({status: 'error', data: {code: 'INVALID_REQUEST', message: 'missing', retryable: false}});
    code(await f.request('workflow.runtime'), 'INVALID_REQUEST');
});

test('mock deadlines, grant revocation and Node replacement release pending workflow ACKs', async t => {
    for (const mode of ['timeout', 'revoke', 'replace', 'close'] as const) {
        const f = await fixture(t);
        t.mock.timers.enable({apis: ['setTimeout']});
        let release!: (value: unknown) => void;
        f.node.respond = () => new Promise(resolve => {release = resolve;});
        const input = control(); const id = randomUUID();
        const result = f.request('workflow.control', input, id);
        await Promise.resolve();
        if (mode === 'timeout') t.mock.timers.tick(DEFAULT_LIMITS.rpc_timeout_ms);
        else if (mode === 'revoke') f.store.revoke_grant({forward_session_id: f.scope.forward_session_id, grant_id: f.grant_id, reason: 'revoked'});
        else if (mode === 'replace') await f.join(new MockSocket(), 'node', [...RELAY_METHODS]);
        else f.store.close();
        code(await result, mode === 'timeout' ? 'TIMEOUT' : mode === 'revoke' ? 'CANCELLED' : mode === 'replace' ? 'NODE_OFFLINE' : 'RELAY_UNAVAILABLE');
        assert.equal(f.node.acks.size, 0);
        release(success(command(input)));
        await Promise.resolve();
        if (mode === 'timeout') code(await f.request('workflow.control', input, id), 'CANCELLED');
        t.mock.timers.reset();
    }
});

test('mock workflow_states info updates fan out only to snapshot readers', async t => {
    const f = await fixture(t, {}, ['bot.snapshot']);
    const received: unknown[] = [];
    f.web.on('info_update', value => received.push(value));
    const summary: WorkflowRunSnapshot = {run_id: randomUUID(), workflow_id: 'a', status: 'completed', steps: 1, active_nodes: [],
        started_at: 0, ended_at: 1, stop_reason: null, error: null, trigger_available: true};
    const update = {forward_session_id: f.scope.forward_session_id, bot_id: f.scope.bot_id, context_generation: randomUUID(),
        stream_id: randomUUID(), seq: 1, emitted_at: Date.now(), data: {name: 'workflow_states', value: {available: true,
            runs: [summary]}}};
    assert.deepEqual(f.store.info(f.node.socket, update), {accepted: true});
    assert.deepEqual(received, [update]);
    const writer = await fixture(t, {}, ['workflow.control']);
    let leaked = false;
    writer.web.on('info_update', () => leaked = true);
    writer.store.info(writer.node.socket, {...update, forward_session_id: writer.scope.forward_session_id});
    assert.equal(leaked, false);
});