import {NodesStore} from "./store";
import {TreeNode} from "./TreeNode";
import {EdgeData, FlowData, FlowDataCreator, Identifiable, RenderedNodeData} from "./FlowDataCreator";
import {Changes, VirtualState} from "./VirtualState";
import {TreeNodeData} from "./Grove";
import {StateBranchNode} from "./StateBranchNode";
import {atomsStore, FloatingHeaderHeightAtom, sidebarEdgeOffsetAtom, treeEmitterAtom} from "./atoms";
import {ReactFlowInstance} from "@xyflow/react";
import {findNodeDomElement, getDimensions} from "../Tree";

let stateHistory: {
    [key: TreeNodeData['id']]: {
        renderedNodes: VirtualState,
        edges: VirtualState,
    }
} = {}

export type FlowState = Pick<NodesStore, 'renderedNodes' | 'edges'>;

/**
 * This class is responsible for managing the flow state of the nodes.
 *
 * Specifically, it manages the rendered nodes and edges from the grove
 */
export class Flow {
    constructor(
        public readonly state: FlowState,
    ) {
    }

    renderTreeNode(treeNode: FlowDataCreator & Identifiable) {
        // ok so maybe we could store the data in memory after this as a history cache for this treeNdde
        // now if we cal this method again, we have to check the flow data for this tree node from the previous time
        // and if in the new data there are is missing from the past data,
        // then we need to remove that past data
        // kind of like the virtual dom
        // should be very straightforward but a little more complex than i would've liked
        // but hey it has to be done and tested to prevent bugs
        // get the flow data for this tree node

        //make sure we get the root
        treeNode = treeNode.getRoot?.() || treeNode;
        const flowData = treeNode.createFlowData();
        const history = this.getOrCreateHistory(treeNode.id);

        // rendered nodes and edges
        ;(['renderedNodes', 'edges'] as (keyof typeof Flow.prototype.state)[]).forEach(
            component => {
                // update history
                const flowDataForComponent = flowData[component];
                const historyForComponent = history[component];

                this.handleChanges(historyForComponent.compare(flowDataForComponent, ['position']), component)
                historyForComponent.updateState(flowDataForComponent);
            }
        )
        this.addPositionToRenderedNodesIfNeeded();
    }

    removeTreeNode(treeNodeId: TreeNode['id']) {
        // remove the rendered nodes and edges from the state
        const history = this.getOrCreateHistory(treeNodeId);
        const emptyFlowData: FlowData = {
            renderedNodes: [],
            edges: [],
        } as FlowData

        (['renderedNodes', 'edges'] as (keyof typeof Flow.prototype.state)[]).forEach(
            component => {
                const flowDataForComponent = emptyFlowData[component];
                const historyForComponent = history[component];

                this.handleChanges(historyForComponent.compare(flowDataForComponent, ['position']), component)
                historyForComponent.updateState(flowDataForComponent);
            }
        )

        // remove the history for this tree node
        this.resetTreeNodeHistory(treeNodeId);
    }

    // reset the state
    resetHistory() {
        Flow.resetHistory();
    }

    static resetHistory() {
        stateHistory = {};
    }

    resetTreeNodeHistory(treeNodeId: string) {
        if (stateHistory[treeNodeId]) {
            delete stateHistory[treeNodeId];
        }
    }

    private addEdge(edge: EdgeData) {
        this.state.edges.push(edge);
    }

    private edgeExists(id: EdgeData['id']): boolean {
        return !!this.getEdge(id);
    }

    private getEdge(id: string) {
        return this.state.edges.find((edge) => edge.id === id);
    }

// we need to potentially remove data when a node changes
    // for example, if we change a branch from tiered to select action (for example after the user deletes all tiers)
    // then we need to remove the flow data of the removed tiers. What would be a good api for this?

    private renderedNodeExists(id: string): boolean {
        return !!this.getRenderedNode(id);
    }

    getRenderedNode(id: string) {
        return this.state.renderedNodes.find((node) => node.id === id);
    }

    private addRenderedNode(renderedNode: Omit<RenderedNodeData, 'position'>) {
        this.state.renderedNodes.push({
            position: {x: 0, y: 0},
            ...renderedNode,
        })
    }

    private updateRenderedNode(newRenderedNode: Omit<RenderedNodeData, "position">) {
        // check each PROPERTY
        // make sure to check that the values for all properties are the same
        // if not, update the values (DON'T UPDATE position {x: 0, y: 0})
        const existingRenderedNode = this.getRenderedNode(newRenderedNode.id);

        if (!existingRenderedNode) {
            return;
        }

        for (const key in newRenderedNode) {
            if (key !== 'position' && existingRenderedNode[key] !== newRenderedNode[key]) {
                existingRenderedNode[key] = newRenderedNode[key];
            }
        }
    }

    private getOrCreateHistory(treeNodeId: string): typeof stateHistory[TreeNodeData['id']] {
        if (!stateHistory[treeNodeId]) {
            stateHistory[treeNodeId] = {
                renderedNodes: new VirtualState([]),
                edges: new VirtualState([]),
            }
        }

        return stateHistory[treeNodeId];
    }

    private handleChanges(changes: Changes<any>, component: keyof typeof Flow.prototype.state) {

        // if added, add
        const renderedNodesOrEdges = this.state[component];

        for (const added of changes.added) {
            renderedNodesOrEdges.push(added);
        }
        // if removed, remove
        for (const removed of changes.removed) {
            const index = renderedNodesOrEdges.findIndex((item) => item.id === removed.id);
            if (index !== -1) {
                renderedNodesOrEdges.splice(index, 1);
            }
        }
        // if modified, update
        for (const {old, new: updated} of changes.modified) {
            const index = renderedNodesOrEdges.findIndex((item) => item.id === old.id);
            if (index !== -1) {
                renderedNodesOrEdges[index] = updated;
            }
        }
    }

    private addPositionToRenderedNodesIfNeeded() {
        for (const renderedNode of this.state.renderedNodes) {
            if (!renderedNode.position) {
                renderedNode.position = {x: 0, y: 0};
            }
        }
    }

    centerRootNode(newTreeNode: TreeNode<StateBranchNode>) {
        const emitter = atomsStore.get(treeEmitterAtom)
        const onStabilizedOnce = (reactFlowInstance: ReactFlowInstance) => {
            emitter.off('renderedNodes:stabilized', onStabilizedOnce);
            if (!reactFlowInstance) {
                return
            }
            const edgeOffset = atomsStore.get(sidebarEdgeOffsetAtom)
            const floatingHeaderOffset = atomsStore.get(FloatingHeaderHeightAtom)
            // this is the tre anchor
            const nodeToCenterId = newTreeNode.getSourceRenderNode()!.getId()
            const node = reactFlowInstance.getNode(nodeToCenterId);
            const domNode = findNodeDomElement(newTreeNode.branchNode.getId())

            if (!node) return;

            const {width, height} = getDimensions(nodeToCenterId)
            const {x: nodeX, y: nodeY} = node.position;
            const rightButtons = domNode?.querySelector('.cp-right-context-buttons') as HTMLElement | null;
            const rightButtonsWidth = rightButtons ? rightButtons.offsetWidth : 0;

            // Adjust width to account for right context buttons if they are visible
            const adjustedWidth = width + rightButtonsWidth;
            // Get React Flow container dimensions
            const containerElement = document.querySelector('.react-flow__viewport')?.parentElement;

            if (!containerElement) return;

            const containerRect = containerElement.getBoundingClientRect();
            const containerWidth = containerRect.width;

            // Get computed styles to account for any padding/borders
            const computedStyle = window.getComputedStyle(containerElement);
            const paddingLeft = parseFloat(computedStyle.paddingLeft) || 0;
            const paddingRight = parseFloat(computedStyle.paddingRight) || 0;
            const usableWidth = containerWidth - paddingLeft - paddingRight;

            const {zoom} = reactFlowInstance.getViewport();

            // Calculate the center position horizontally, but position at top with 40px offset vertically
            const centerX = (usableWidth / 2) - (adjustedWidth / 2);
            const topY = (edgeOffset * 2) + floatingHeaderOffset + edgeOffset; // 40px from top

            // Calculate viewport translation needed
            const viewportX = centerX - (nodeX * zoom);
            const viewportY = topY - (nodeY * zoom);

            reactFlowInstance.setViewport(
                {x: viewportX, y: viewportY, zoom}, {duration: 800}
            )
        };

        emitter.on('renderedNodes:stabilized', onStabilizedOnce);
    }

}
