import { BinaryOperatorAggregate } from "../channels/binop.js";
import { END, CompiledGraph, Graph, START } from "./graph.js";
import { LastValue } from "../channels/last_value.js";
import { ChannelWrite, PASSTHROUGH, SKIP_WRITE, } from "../pregel/write.js";
import { ChannelRead, PregelNode } from "../pregel/read.js";
import { NamedBarrierValue } from "../channels/named_barrier_value.js";
import { EphemeralValue } from "../channels/ephemeral_value.js";
import { RunnableCallable } from "../utils.js";
import { TAG_HIDDEN } from "../constants.js";
import { InvalidUpdateError } from "../errors.js";
const ROOT = "__root__";
export class StateGraph extends Graph {
    constructor(fields) {
        super();
        Object.defineProperty(this, "channels", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        // TODO: this doesn't dedupe edges as in py, so worth fixing at some point
        Object.defineProperty(this, "waitingEdges", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: new Set()
        });
        this.channels = _getChannels(fields.channels);
        for (const c of Object.values(this.channels)) {
            if (c.lc_graph_name === "BinaryOperatorAggregate") {
                this.supportMultipleEdges = true;
                break;
            }
        }
    }
    get allEdges() {
        return new Set([
            ...this.edges,
            ...Array.from(this.waitingEdges).flatMap(([starts, end]) => starts.map((start) => [start, end])),
        ]);
    }
    addNode(key, action) {
        if (key in this.channels) {
            throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
        }
        return super.addNode(key, action);
    }
    addEdge(startKey, endKey) {
        if (typeof startKey === "string") {
            return super.addEdge(startKey, endKey);
        }
        if (this.compiled) {
            console.warn("Adding an edge to a graph that has already been compiled. This will " +
                "not be reflected in the compiled graph.");
        }
        for (const start of startKey) {
            if (start === END) {
                throw new Error("END cannot be a start node");
            }
            if (!Object.keys(this.nodes).some((node) => node === start)) {
                throw new Error(`Need to add_node ${start} first`);
            }
        }
        if (endKey === END) {
            throw new Error("END cannot be an end node");
        }
        if (!Object.keys(this.nodes).some((node) => node === endKey)) {
            throw new Error(`Need to add_node ${endKey} first`);
        }
        this.waitingEdges.add([startKey, endKey]);
        return this;
    }
    compile({ checkpointer, interruptBefore, interruptAfter, } = {}) {
        // validate the graph
        this.validate([
            ...(Array.isArray(interruptBefore) ? interruptBefore : []),
            ...(Array.isArray(interruptAfter) ? interruptAfter : []),
        ]);
        // prepare output channels
        const stateKeys = Object.keys(this.channels);
        const outputs = stateKeys.length === 1 && stateKeys[0] === ROOT
            ? stateKeys[0]
            : stateKeys;
        // create empty compiled graph
        const compiled = new CompiledStateGraph({
            builder: this,
            checkpointer,
            interruptAfter,
            interruptBefore,
            autoValidate: false,
            nodes: {},
            channels: {
                ...this.channels,
                [START]: new EphemeralValue(),
            },
            inputs: START,
            outputs,
            streamChannels: outputs,
            streamMode: "updates",
        });
        // attach nodes, edges and branches
        compiled.attachNode(START);
        for (const [key, node] of Object.entries(this.nodes)) {
            compiled.attachNode(key, node);
        }
        for (const [start, end] of this.edges) {
            compiled.attachEdge(start, end);
        }
        for (const [starts, end] of this.waitingEdges) {
            compiled.attachEdge(starts, end);
        }
        for (const [start, branches] of Object.entries(this.branches)) {
            for (const [name, branch] of Object.entries(branches)) {
                compiled.attachBranch(start, name, branch);
            }
        }
        return compiled.validate();
    }
}
function _getChannels(schema) {
    const channels = {};
    for (const [name, val] of Object.entries(schema)) {
        if (name === ROOT) {
            channels[name] = getChannel(val);
        }
        else {
            const key = name;
            channels[name] = getChannel(val);
        }
    }
    return channels;
}
function getChannel(reducer) {
    if (typeof reducer === "object" &&
        reducer &&
        "reducer" in reducer &&
        reducer.reducer) {
        return new BinaryOperatorAggregate(reducer.reducer, reducer.default);
    }
    if (typeof reducer === "object" &&
        reducer &&
        "value" in reducer &&
        reducer.value) {
        return new BinaryOperatorAggregate(reducer.value, reducer.default);
    }
    return new LastValue();
}
export class CompiledStateGraph extends CompiledGraph {
    attachNode(key, node) {
        const stateKeys = Object.keys(this.builder.channels);
        function getStateKey(key, input) {
            if (!input) {
                return SKIP_WRITE;
            }
            else if (typeof input !== "object" || Array.isArray(input)) {
                throw new InvalidUpdateError(`Expected dict, got ${typeof input}`);
            }
            else {
                return key in input ? input[key] : SKIP_WRITE;
            }
        }
        // state updaters
        const stateWriteEntries = stateKeys.map((key) => key === ROOT
            ? { channel: key, value: PASSTHROUGH, skipNone: true }
            : {
                channel: key,
                value: PASSTHROUGH,
                mapper: new RunnableCallable({
                    func: getStateKey.bind(null, key),
                    trace: false,
                    recurse: false,
                }),
            });
        // add node and output channel
        if (key === START) {
            this.nodes[key] = new PregelNode({
                tags: [TAG_HIDDEN],
                triggers: [START],
                channels: [START],
                writers: [new ChannelWrite(stateWriteEntries, [TAG_HIDDEN])],
            });
        }
        else {
            this.channels[key] = new EphemeralValue();
            this.nodes[key] = new PregelNode({
                triggers: [],
                // read state keys
                channels: stateKeys.length === 1 && stateKeys[0] === ROOT
                    ? stateKeys
                    : stateKeys.reduce((acc, k) => {
                        acc[k] = k;
                        return acc;
                    }, {}),
                // publish to this channel and state keys
                writers: [
                    new ChannelWrite(stateWriteEntries.concat({ channel: key, value: key }), [TAG_HIDDEN]),
                ],
                bound: node,
            });
        }
    }
    attachEdge(start, end) {
        if (end === END) {
            return;
        }
        if (Array.isArray(start)) {
            const channelName = `join:${start.join("+")}:${end}`;
            // register channel
            this.channels[channelName] =
                new NamedBarrierValue(new Set(start));
            // subscribe to channel
            this.nodes[end].triggers.push(channelName);
            // publish to channel
            for (const s of start) {
                this.nodes[s].writers.push(new ChannelWrite([{ channel: channelName, value: s }], [TAG_HIDDEN]));
            }
        }
        else if (start === START) {
            const channelName = `start:${end}`;
            // register channel
            this.channels[channelName] =
                new EphemeralValue();
            // subscribe to channel
            this.nodes[end].triggers.push(channelName);
            // publish to channel
            this.nodes[START].writers.push(new ChannelWrite([{ channel: channelName, value: START }], [TAG_HIDDEN]));
        }
        else {
            this.nodes[end].triggers.push(start);
        }
    }
    attachBranch(start, name, branch) {
        // attach branch publisher
        this.nodes[start].writers.push(branch.compile(
        // writer
        (dests) => {
            const filteredDests = dests.filter((dest) => dest !== END);
            if (!filteredDests.length) {
                return;
            }
            const writes = filteredDests.map((dest) => ({
                channel: `branch:${start}:${name}:${dest}`,
                value: start,
            }));
            return new ChannelWrite(writes, [TAG_HIDDEN]);
        }, 
        // reader
        (config) => ChannelRead.doRead(config, this.outputs, true)));
        // attach branch subscribers
        const ends = branch.ends
            ? Object.values(branch.ends)
            : Object.keys(this.builder.nodes);
        for (const end of ends) {
            if (end === END) {
                continue;
            }
            const channelName = `branch:${start}:${name}:${end}`;
            this.channels[channelName] =
                new EphemeralValue();
            this.nodes[end].triggers.push(channelName);
        }
    }
}
