import {assign, createActor, fromPromise, setup} from 'xstate';
import {randomUUID} from 'node:crypto';
import {isDeepStrictEqual} from 'node:util';
import type {Json} from '../../../types/bot/workflow/workflow.ts';
import type {BranchState, CompiledWorkflow, NodeResult, NodeReturn, RunSnapshot, StopReason,
    WorkflowDependencies, WorkflowError, WorkflowRunEvent} from '../../../types/bot/workflow/runtime.ts';
import {execute_node} from './execute.ts';
import {assert_json} from './value.ts';

interface Token {
    node_id: string;
    edge_id: string | null;
    state: BranchState;
    joins: Record<string, string>;
    bases: Record<string, Record<string, Json>>;
    forced?: NodeReturn;
}
interface Completion {token: Token; returned: NodeReturn; execution: number}
interface ExecutionCompletion {execution: number; completed?: Completion; error?: unknown}
export interface WorkflowRun {
    run_id: string;
    done: Promise<RunSnapshot>;
    get_snapshot(): RunSnapshot;
    stop(reason?: StopReason): Promise<void>;
}

function failure(code: string, message: string, node_id = ''): WorkflowError {return {code, message, node_id};}

function merge_states(tokens: Token[], base: Record<string, Json>): BranchState {
    const state = structuredClone(tokens[0].state);
    state.vars = structuredClone(base);
    state.inputs = {};
    state.outputs = {};
    const writes = new Set<string>();
    for (const token of tokens) {
        for (const [key, value] of Object.entries(token.state.vars)) {
            if (Object.hasOwn(base, key) && isDeepStrictEqual(base[key], value)) continue;
            if (writes.has(key) && !isDeepStrictEqual(state.vars[key], value)) {
                throw Object.assign(new Error(`Conflicting branch variable: ${key}`), {code: 'MERGE_CONFLICT'});
            }
            writes.add(key);
            state.vars[key] = structuredClone(value);
        }
        for (const [id, result] of Object.entries(token.state.outputs)) {
            if (!state.outputs[id] || result.execution > state.outputs[id].execution) state.outputs[id] = structuredClone(result);
        }
        for (const [key, count] of Object.entries(token.state.retry_counts)) {
            state.retry_counts[key] = Math.max(state.retry_counts[key] ?? 0, count);
        }
        Object.assign(state.inputs, structuredClone(token.state.inputs));
    }
    state.last_result = structuredClone(tokens.find(token => token.state.last_result?.status === 'failure')?.state.last_result
        ?? tokens[tokens.length - 1].state.last_result);
    return state;
}

export function create_workflow_run(graph: CompiledWorkflow, deps: WorkflowDependencies,
    trigger: Record<string, Json> = {}, notify: (event: WorkflowRunEvent) => void = () => {}): WorkflowRun {
    assert_json(trigger);
    if (!trigger || Array.isArray(trigger) || typeof trigger !== 'object') throw new Error('Invalid workflow trigger');
    const run_id = randomUUID();
    const controller = new AbortController();
    const initial: BranchState = {
        trigger: structuredClone(trigger), vars: structuredClone(graph.definition.initial_vars ?? {}),
        outputs: {}, last_result: null, inputs: {}, retry_counts: {},
    };
    const snapshot: RunSnapshot = {
        run_id, workflow_id: graph.definition.id, status: 'running', steps: 0, active_nodes: [],
        started_at: Date.now(), ended_at: null, stop_reason: null, error: null,
        trigger: structuredClone(trigger), trigger_available: true, branches: [],
    };
    const emit = (event: WorkflowRunEvent) => {try {notify(structuredClone(event));} catch { /* Observers cannot stop execution. */ }};
    const update = () => emit({type: 'run_updated', snapshot});
    const waiting = new Map<string, Map<string, Token[]>>();
    const ready: Token[] = [{node_id: graph.definition.entry, edge_id: null, state: initial, joins: {}, bases: {}}];
    const active = new Map<number, Promise<void>>();
    const completions: ExecutionCompletion[] = [];
    let wake: (() => void) | null = null;
    const active_nodes = new Map<number, string>();
    let fatal: WorkflowError | null = null;
    const token_count = () => ready.length + active.size + snapshot.branches.length
        + [...waiting.values()].reduce((sum, group) => sum + [...group.values()].reduce((n, queue) => n + queue.length, 0), 0);

    const join_owners = graph.join_owners;

    const enqueue = (token: Token) => {
        const node = graph.nodes.get(token.node_id)!;
        if (node.type !== 'merge' || node.config.mode === 'any') {ready.push(token); return;}
        const incoming = graph.incoming.get(node.id)!;
        const key = `${node.id}:${token.joins[node.id] ?? 'root'}`;
        let group = waiting.get(key);
        if (!group) {group = new Map(); waiting.set(key, group);}
        const edge_id = token.edge_id!;
        group.set(edge_id, [...group.get(edge_id) ?? [], token]);
        if (!incoming.every(edge => (group!.get(edge.id)?.length ?? 0) > 0)) return;
        const tokens = incoming.map(edge => group!.get(edge.id)!.shift()!);
        if ([...group.values()].every(queue => queue.length === 0)) waiting.delete(key);
        let joined: Token;
        try {joined = {...tokens[0], state: merge_states(tokens, tokens[0].bases[node.id] ?? initial.vars),
            joins: {...tokens[0].joins}, bases: {...tokens[0].bases}};}
        catch (error) {
            joined = {...tokens[0], joins: {...tokens[0].joins}, bases: {...tokens[0].bases}, forced: {
                port: 'failure', result: {status: 'failure', output: null,
                    error: failure((error as {code?: string}).code ?? 'MERGE_FAILED', String((error as Error).message), node.id)},
            }};
        }
        // Keep lineage when a retry returns to the join without revisiting its owner.
        // A fresh owner execution replaces the key; FIFO inputs isolate rounds within it.
        joined.bases[node.id] = structuredClone(joined.state.vars);
        ready.push(joined);
    };

    const commit = ({token, returned, execution}: Completion) => {
        const state = structuredClone(token.state);
        state.vars = {...state.vars, ...returned.vars};
        state.retry_counts = {...state.retry_counts, ...returned.retry_counts};
        state.outputs[token.node_id] = {...structuredClone(returned.result), execution};
        state.last_result = {...structuredClone(returned.result), node_id: token.node_id};
        assert_json(state);
        emit({type: 'node_finished', run_id, node_id: token.node_id, execution, result: returned.result});
        const edges = (graph.outgoing.get(token.node_id) ?? []).filter(edge => edge.port === returned.port);
        if (token_count() + Math.max(1, edges.length) > 256) {
            fatal = failure('BRANCH_LIMIT_EXCEEDED', 'Workflow exceeded its live branch limit', token.node_id);
            return;
        }
        if (!edges.length) {
            snapshot.branches.push(state);
            if (returned.result.status === 'failure') fatal ??= returned.result.error
                ?? failure('UNHANDLED_FAILURE', 'A branch ended with an unhandled failure', token.node_id);
            return;
        }
        const joins = {...token.joins};
        const bases = {...token.bases};
        if (edges.length) {
            const fork_id = randomUUID();
            let baseline: Record<string, Json> | undefined;
            for (const [join, owner] of join_owners) if (owner === token.node_id) {
                joins[join] = fork_id;
                // Baselines are immutable; all joins owned by this execution share one copy.
                baseline ??= structuredClone(state.vars);
                bases[join] = baseline;
            }
        }
        for (const edge of edges) {
            const next = structuredClone(state);
            next.inputs = {[edge.id]: {source: edge.source, port: edge.port, result: structuredClone(returned.result) as unknown as Json}};
            enqueue({node_id: edge.target, edge_id: edge.id, state: next, joins: {...joins}, bases: {...bases}});
        }
    };

    const execute_graph = async () => {
        try {
            while ((ready.length || active.size) && !fatal && !controller.signal.aborted) {
                while (ready.length && active.size < 16 && !fatal && !controller.signal.aborted) {
                    const token = ready.shift()!;
                    if (snapshot.steps >= graph.definition.policy.max_steps) {
                        fatal = failure('STEP_LIMIT_EXCEEDED', 'Workflow exceeded its step limit', token.node_id); break;
                    }
                    const execution = ++snapshot.steps;
                    active_nodes.set(execution, token.node_id);
                    snapshot.active_nodes = [...active_nodes.values()];
                    emit({type: 'node_started', run_id, node_id: token.node_id, execution});
                    update();
                    const complete = (value: ExecutionCompletion) => {
                        completions.push(value);
                        wake?.(); wake = null;
                    };
                    const promise = (token.forced ? Promise.resolve(token.forced)
                        : execute_node(graph.nodes.get(token.node_id)!, structuredClone(token.state), controller.signal, deps))
                        .then(returned => complete({execution, completed: {token, returned, execution}}), error => complete({execution, error}));
                    active.set(execution, promise);
                }
                if (fatal || !active.size) break;
                if (!completions.length) await new Promise<void>(resolve => {wake = resolve;});
                const finished = completions.shift()!;
                active.delete(finished.execution);
                active_nodes.delete(finished.execution);
                snapshot.active_nodes = [...active_nodes.values()];
                if (controller.signal.aborted) break;
                if (finished.error) throw finished.error;
                commit(finished.completed!);
                update();
            }
            if (!controller.signal.aborted && !fatal && waiting.size) {
                fatal = failure('MERGE_INCOMPLETE', 'All-input merge has missing inputs; check exclusive branches and loops');
            }
        } catch (error) {
            if (!controller.signal.aborted) fatal = failure((error as {code?: string})?.code ?? 'RUNTIME_ERROR',
                error instanceof Error ? error.message : String(error));
        } finally {
            if (fatal) controller.abort(new Error(fatal.message));
            await Promise.all(active.values());
            completions.length = 0;
            active.clear(); ready.length = 0; waiting.clear();
            snapshot.active_nodes = [];
            snapshot.error = fatal;
        }
        return fatal;
    };

    const machine = setup({
        types: {context: {} as {error: WorkflowError | null}, events: {} as {type: 'STOP'; reason: StopReason}},
        actors: {graph: fromPromise(execute_graph)},
        actions: {
            stop: ({event}) => {
                if (snapshot.stop_reason || fatal) return;
                snapshot.stop_reason = event.reason;
                snapshot.status = 'stopping';
                controller.abort(new Error(`Workflow stopped: ${event.reason}`));
                update();
            },
        },
    }).createMachine({
        id: 'workflow', initial: 'running', context: {error: null},
        states: {
            running: {
                on: {STOP: {actions: 'stop'}},
                invoke: {src: 'graph', onDone: {target: 'ending', actions: assign({error: ({event}) => event.output})},
                    onError: {target: 'failed', actions: assign({error: ({event}) => failure('RUNTIME_ERROR', String(event.error))})}},
            },
            ending: {always: [
                {guard: () => snapshot.stop_reason === 'timeout', target: 'timed_out'},
                {guard: () => snapshot.stop_reason !== null, target: 'cancelled'},
                {guard: ({context}) => context.error !== null, target: 'failed'},
                {target: 'completed'},
            ]},
            completed: {type: 'final'}, failed: {type: 'final'}, cancelled: {type: 'final'}, timed_out: {type: 'final'},
        },
    });
    const actor = createActor(machine);
    let resolve!: (value: RunSnapshot) => void;
    const done = new Promise<RunSnapshot>(finish => {resolve = finish;});
    const timer = setTimeout(() => actor.send({type: 'STOP', reason: 'timeout'}), graph.definition.policy.timeout_ms);
    actor.subscribe({next: value => {
        if (value.status !== 'done') return;
        clearTimeout(timer);
        snapshot.status = value.value as RunSnapshot['status'];
        snapshot.error = value.context.error;
        snapshot.ended_at = Date.now();
        update();
        resolve(structuredClone(snapshot));
    }, error: error => {
        clearTimeout(timer); snapshot.status = 'failed'; snapshot.ended_at = Date.now();
        snapshot.error = failure('ACTOR_ERROR', String(error)); update(); resolve(structuredClone(snapshot));
    }});
    actor.start();
    return {
        run_id, done, get_snapshot: () => structuredClone(snapshot),
        async stop(reason = 'manual') {if (snapshot.ended_at === null) actor.send({type: 'STOP', reason}); await done;},
    };
}