import { color } from "../utils/termcolors.js";
import { SimpleMachineError } from "./error.js";
import { conditionalEdge, edgeToJSON, isRegularEdge, regularEdge, } from "./types.js";
import { StatelogClient } from "../statelogClient.js";
export class GoToNode {
    to;
    data;
    constructor(to, data) {
        this.to = to;
        this.data = data;
    }
}
export function goToNode(to, data) {
    return new GoToNode(to, data);
}
export class SimpleMachine {
    nodes = {};
    edges = {};
    config;
    statelogClient = null;
    constructor(config = {}) {
        this.config = config;
        if (config.statelog) {
            this.statelogClient = new StatelogClient({
                host: config.statelog.host,
                apiKey: config.statelog.apiKey,
                projectId: config.statelog.projectId,
                traceId: config.statelog.traceId,
                debugMode: config.statelog.debugMode ?? false,
                observability: config.statelog.observability,
                logFile: config.statelog.logFile,
            });
        }
    }
    node(id, func) {
        this.nodes[id] = func;
    }
    edge(from, to) {
        if (!this.edges[from]) {
            this.edges[from] = regularEdge(to);
        }
        else {
            throw new SimpleMachineError(` ${from} already has an edge, which leads to ${this.edges[from]}.`);
        }
    }
    conditionalEdge(from, adjacentNodes, to) {
        if (!this.edges[from]) {
            this.edges[from] = conditionalEdge(to, adjacentNodes);
        }
        else {
            throw new SimpleMachineError(` ${from} already has an edge, which leads to ${this.edges[from]}.`);
        }
    }
    debug(message, data) {
        let debugStr = `${color.magenta("[DEBUG]")}: ${message}`;
        if (this.config.debug?.logData && data !== undefined) {
            debugStr += ` | Data: ${color.green(JSON.stringify(data))}`;
        }
        if (this.config.debug?.log) {
            console.log(debugStr);
        }
        //this.statelogClient?.debug(message, data || {});
    }
    async run(startId, input, options) {
        const client = options?.statelogClient ?? this.statelogClient;
        const jsonEdges = {};
        for (const from in this.edges) {
            jsonEdges[from] = edgeToJSON(this.edges[from]);
        }
        client?.graph({
            nodes: Object.keys(this.nodes),
            edges: jsonEdges,
            startNode: startId,
        });
        let currentId = startId;
        let data = input;
        while (currentId) {
            options?.onNodeEnter?.(currentId);
            const nodeFunc = this.nodes[currentId];
            if (!nodeFunc) {
                throw new SimpleMachineError(`Node function for ${currentId} not found.`);
            }
            if (this.config.hooks?.beforeNode) {
                this.debug(`Before hook for node: ${color.green(currentId)}`, data);
                const startData = data;
                const startTime = performance.now();
                data = await this.config.hooks.beforeNode(currentId, data);
                const endTime = performance.now();
                client?.beforeHook({
                    nodeId: currentId,
                    startData,
                    endData: data,
                    timeTaken: endTime - startTime,
                });
            }
            this.debug(`Executing node: ${color.green(currentId)}`, data);
            const nodeSpanId = client?.startSpan("nodeExecution");
            client?.enterNode({ nodeId: currentId, data });
            const startTime = performance.now();
            let nextNode;
            try {
                const result = await this.runAndValidate(nodeFunc, currentId, data);
                const endTime = performance.now();
                if (result instanceof GoToNode) {
                    nextNode = result.to;
                    data = result.data;
                }
                else {
                    data = result;
                }
                client?.exitNode({
                    nodeId: currentId,
                    data,
                    timeTaken: endTime - startTime,
                });
            }
            finally {
                client?.endSpan(nodeSpanId);
            }
            this.debug(`Completed node: ${color.green(currentId)}`, data);
            if (this.config.hooks?.afterNode) {
                this.debug(`After hook for node: ${color.green(currentId)}`, data);
                const startData = data;
                const startTime = performance.now();
                data = await this.config.hooks.afterNode(currentId, data);
                const endTime = performance.now();
                client?.afterHook({
                    nodeId: currentId,
                    startData,
                    endData: data,
                    timeTaken: endTime - startTime,
                });
            }
            const edge = this.edges[currentId];
            if (edge === undefined) {
                currentId = null;
                continue;
            }
            if (nextNode && edge) {
                const isValidTarget = this.validateGoToNodeTarget(nextNode, edge);
                if (!isValidTarget) {
                    throw new SimpleMachineError(`${currentId} tried to go to ${nextNode}, but did not specify a conditional edge to it. Use graph.conditionalEdge("${currentId}", ["${nextNode}"]) to define the edge.`);
                }
                client?.followEdge({
                    fromNodeId: currentId,
                    toNodeId: nextNode,
                    isConditionalEdge: false,
                    data,
                });
                this.debug(`Following goto edge to: ${color.green(nextNode)}`, data);
                currentId = nextNode;
                continue;
            }
            if (isRegularEdge(edge)) {
                client?.followEdge({
                    fromNodeId: currentId,
                    toNodeId: edge.to,
                    isConditionalEdge: false,
                    data,
                });
                this.debug(`Following regular edge to: ${color.green(edge.to)}`);
                currentId = edge.to;
            }
            else {
                if (edge.condition) {
                    const nextId = await edge.condition(data);
                    client?.followEdge({
                        fromNodeId: currentId,
                        toNodeId: nextId,
                        isConditionalEdge: true,
                        data,
                    });
                    this.debug(`Following conditional edge to: ${color.green(nextId)}`, data);
                    currentId = nextId;
                }
                else {
                    this.debug(`Exiting graph from node: ${color.green(currentId)}`);
                    currentId = null;
                    /* throw new SimpleMachineError(
                      `Expected ${currentId} to return a GoToNode, as no function was specified for the conditional edges to ${edge.adjacentNodes.join(", ")}.`,
                    ); */
                }
            }
        }
        return data;
    }
    async runAndValidate(nodeFunc, currentId, _data, retries = 0) {
        const result = await nodeFunc(_data);
        let data;
        if (result instanceof GoToNode) {
            data = result.data;
        }
        else {
            data = result;
        }
        if (this.config.validation?.func) {
            const maxRetries = this.config.validation.maxRetries ?? 0;
            const isValid = await this.config.validation.func(data);
            while (!isValid) {
                if (retries >= maxRetries) {
                    throw new SimpleMachineError(`Validation failed for node ${currentId} after ${maxRetries} retries.`);
                }
                this.debug(`Validation failed for node ${color.green(currentId)}, retrying... (${retries + 1}/${maxRetries})`, data);
                return this.runAndValidate(nodeFunc, currentId, _data, retries + 1);
            }
        }
        return result;
    }
    prettyPrint() {
        for (const from in this.edges) {
            const to = this.edges[from];
            if (!to)
                continue;
            console.log(`${from} -> ${this.prettyPrintEdge(to)}`);
        }
    }
    prettyPrintEdge(edge) {
        if (isRegularEdge(edge)) {
            return edge.to;
        }
        else {
            return edge.adjacentNodes.join(" | ");
        }
    }
    toMermaid() {
        let mermaid = "graph TD\n";
        for (const from in this.edges) {
            const to = this.edges[from];
            if (!to)
                continue;
            if (isRegularEdge(to)) {
                mermaid += `  ${from} --> ${to.to}\n`;
            }
            else {
                to.adjacentNodes.forEach((adjNode) => {
                    mermaid += `  ${from} --> ${adjNode}\n`;
                });
            }
        }
        return mermaid;
    }
    merge(another) {
        for (const nodeId in another.nodes) {
            if (this.nodes[nodeId]) {
                throw new SimpleMachineError(`Node ${nodeId} already exists in the current SimpleMachine.`);
            }
            this.nodes[nodeId] =
                another.nodes[nodeId];
        }
        for (const from in another.edges) {
            if (this.edges[from]) {
                throw new SimpleMachineError(`Edge from ${from} already exists in the current SimpleMachine.`);
            }
            this.edges[from] =
                another.edges[from];
        }
    }
    toJSON() {
        const edges = {};
        for (const from in this.edges) {
            const edge = this.edges[from];
            if (!edge)
                continue;
            if (isRegularEdge(edge)) {
                edges[from] = [edge.to];
            }
            else {
                edges[from] = edge.adjacentNodes;
            }
        }
        return {
            nodes: Object.keys(this.nodes),
            edges,
            config: {
                debug: this.config.debug,
            },
        };
    }
    validateGoToNodeTarget(to, edge) {
        if (!isRegularEdge(edge)) {
            if (edge.adjacentNodes.includes(to)) {
                return true;
            }
        }
        return false;
    }
}
