path: root/frontend/src/store/flowStore.ts

import { create } from 'zustand';
import { 
  addEdge, 
  applyNodeChanges,
  applyEdgeChanges,
  type Connection, 
  type Edge, 
  type EdgeChange, 
  type Node, 
  type NodeChange, 
  type OnNodesChange, 
  type OnEdgesChange, 
  type OnConnect,
  getIncomers,
  getOutgoers
} from 'reactflow';

export type NodeStatus = 'idle' | 'loading' | 'success' | 'error';

export interface Message {
  id?: string;
  role: 'user' | 'assistant' | 'system';
  content: string;
}

export interface Trace {
  id: string;
  sourceNodeId: string;
  color: string;
  messages: Message[];
}

export interface NodeData {
  label: string;
  model: string;
  temperature: number;
  apiKey?: string;
  systemPrompt: string;
  userPrompt: string;
  mergeStrategy: 'raw' | 'smart';
  enableGoogleSearch?: boolean;
  reasoningEffort: 'low' | 'medium' | 'high';  // For OpenAI reasoning models
  disabled?: boolean;  // Greyed out, no interaction
  
  // Traces logic
  traces: Trace[];          // INCOMING Traces
  outgoingTraces: Trace[];  // ALL Outgoing (inherited + self + forks)
  forkedTraces: Trace[];    // Manually created forks from "New" handle
  activeTraceIds: string[]; 
  
  response: string;    
  status: NodeStatus;
  inputs: number; 
  [key: string]: any;
}

export type LLMNode = Node<NodeData>;

// Archived node template (for reuse)
export interface ArchivedNode {
  id: string;
  label: string;
  model: string;
  systemPrompt: string;
  temperature: number;
  reasoningEffort: 'low' | 'medium' | 'high';
}

interface FlowState {
  nodes: LLMNode[];
  edges: Edge[];
  selectedNodeId: string | null;
  archivedNodes: ArchivedNode[];  // Stored node templates

  onNodesChange: OnNodesChange;
  onEdgesChange: OnEdgesChange;
  onConnect: OnConnect;
  
  addNode: (node: LLMNode) => void;
  updateNodeData: (nodeId: string, data: Partial<NodeData>) => void;
  setSelectedNode: (nodeId: string | null) => void;
  
  getActiveContext: (nodeId: string) => Message[]; 
  
  // Actions
  deleteEdge: (edgeId: string) => void;
  deleteNode: (nodeId: string) => void;
  deleteBranch: (startNodeId?: string, startEdgeId?: string) => void;
  
  // Archive actions
  toggleNodeDisabled: (nodeId: string) => void;
  archiveNode: (nodeId: string) => void;
  removeFromArchive: (archiveId: string) => void;
  createNodeFromArchive: (archiveId: string, position: { x: number; y: number }) => void;
  
  // Trace disable
  toggleTraceDisabled: (edgeId: string) => void;
  updateEdgeStyles: () => void;

  propagateTraces: () => void;
}

// Hash string to color
const getStableColor = (str: string) => {
  let hash = 0;
  for (let i = 0; i < str.length; i++) {
    hash = str.charCodeAt(i) + ((hash << 5) - hash);
  }
  const hue = Math.abs(hash % 360);
  return `hsl(${hue}, 70%, 60%)`; 
};

const useFlowStore = create<FlowState>((set, get) => ({
  nodes: [],
  edges: [],
  selectedNodeId: null,
  archivedNodes: [],

  onNodesChange: (changes: NodeChange[]) => {
    set({
      nodes: applyNodeChanges(changes, get().nodes) as LLMNode[],
    });
  },
  onEdgesChange: (changes: EdgeChange[]) => {
    set({
      edges: applyEdgeChanges(changes, get().edges),
    });
    get().propagateTraces();
  },
  onConnect: (connection: Connection) => {
    const { nodes } = get();
    
    // Check if connecting from "new-trace" handle
    if (connection.sourceHandle === 'new-trace') {
       // Logic: Create a new Forked Trace on the source node
       const sourceNode = nodes.find(n => n.id === connection.source);
       if (sourceNode) {
          // Generate the content for this new trace (it's essentially the Self Trace of this node)
          const myResponseMsg: Message[] = [];
          if (sourceNode.data.userPrompt) myResponseMsg.push({ id: `${sourceNode.id}-u`, role: 'user', content: sourceNode.data.userPrompt });
          if (sourceNode.data.response) myResponseMsg.push({ id: `${sourceNode.id}-a`, role: 'assistant', content: sourceNode.data.response });
          
          const newForkId = `trace-${sourceNode.id}-fork-${Date.now()}`;
          const newForkTrace: Trace = {
             id: newForkId,
             sourceNodeId: sourceNode.id,
             color: getStableColor(newForkId), // Unique color for this fork
             messages: [...myResponseMsg]
          };
          
          // Update Source Node to include this fork
          get().updateNodeData(sourceNode.id, {
             forkedTraces: [...(sourceNode.data.forkedTraces || []), newForkTrace]
          });
          
          // Redirect connection to the new handle
          // Note: We must wait for propagateTraces to render the new handle? 
          // ReactFlow might complain if handle doesn't exist yet.
          // But since we updateNodeData synchronously (mostly), it might work.
          // Let's use the new ID for the connection.
          
          set({
            edges: addEdge({ 
              ...connection, 
              sourceHandle: `trace-${newForkId}`, // Redirect!
              style: { stroke: newForkTrace.color, strokeWidth: 2 } 
            }, get().edges),
          });
          
          // Trigger propagation to update downstream
          setTimeout(() => get().propagateTraces(), 0);
          return;
       }
    }
    
    // Normal connection
    set({
      edges: addEdge({ 
        ...connection, 
        style: { stroke: '#888', strokeWidth: 2 } 
      }, get().edges),
    });
    setTimeout(() => get().propagateTraces(), 0);
  },

  addNode: (node: LLMNode) => {
    set((state) => ({ nodes: [...state.nodes, node] }));
    setTimeout(() => get().propagateTraces(), 0);
  },

  updateNodeData: (nodeId: string, data: Partial<NodeData>) => {
    set((state) => ({
      nodes: state.nodes.map((node) => {
        if (node.id === nodeId) {
          return { ...node, data: { ...node.data, ...data } };
        }
        return node;
      }),
    }));
    
    if (data.response !== undefined || data.userPrompt !== undefined) {
      get().propagateTraces();
    }
  },

  setSelectedNode: (nodeId: string | null) => {
    set({ selectedNodeId: nodeId });
  },

  getActiveContext: (nodeId: string) => {
    const node = get().nodes.find(n => n.id === nodeId);
    if (!node) return [];

    // The traces stored in node.data.traces are the INCOMING traces.
    // If we select one, we want its history.
    
    const activeTraces = node.data.traces.filter(t => 
      node.data.activeTraceIds?.includes(t.id)
    );
    
    const contextMessages: Message[] = [];
    activeTraces.forEach(t => {
      contextMessages.push(...t.messages);
    });
    
    return contextMessages;
  },

  deleteEdge: (edgeId: string) => {
    set({
      edges: get().edges.filter(e => e.id !== edgeId)
    });
    get().propagateTraces();
  },

  deleteNode: (nodeId: string) => {
    set({
      nodes: get().nodes.filter(n => n.id !== nodeId),
      edges: get().edges.filter(e => e.source !== nodeId && e.target !== nodeId)
    });
    get().propagateTraces();
  },

  deleteBranch: (startNodeId?: string, startEdgeId?: string) => {
    const { edges, nodes } = get();
    // We ONLY delete edges, NOT nodes.
    const edgesToDelete = new Set<string>();

    // Helper to traverse downstream EDGES based on Trace Dependency
    const traverse = (currentEdge: Edge) => {
      if (edgesToDelete.has(currentEdge.id)) return;
      edgesToDelete.add(currentEdge.id);
      
      const targetNodeId = currentEdge.target;
      // Identify the trace ID carried by this edge
      const traceId = currentEdge.sourceHandle?.replace('trace-', '');
      if (!traceId) return; 
      
      // Look for outgoing edges from the target node that carry the EVOLUTION of this trace.
      // Our logic generates next trace ID as: `${traceId}_${targetNodeId}`
      const expectedNextTraceId = `${traceId}_${targetNodeId}`;
      
      const outgoing = edges.filter(e => e.source === targetNodeId);
      outgoing.forEach(nextEdge => {
        // If the outgoing edge carries the evolved trace, delete it too
        if (nextEdge.sourceHandle === `trace-${expectedNextTraceId}`) {
          traverse(nextEdge);
        }
      });
    };

    if (startNodeId) {
      // If deleting a node, we delete ALL outgoing edges recursively.
      // Because all traces passing through this node are broken.
      // But we can't use `traverse` directly because we don't have a single start edge.
      // We just start traverse on ALL outgoing edges of this node.
      const initialOutgoing = edges.filter(e => e.source === startNodeId);
      initialOutgoing.forEach(e => traverse(e));
      
      // Also delete incoming to this node
      const incomingToNode = edges.filter(e => e.target === startNodeId);
      incomingToNode.forEach(e => edgesToDelete.add(e.id));
      
      set({
        nodes: nodes.filter(n => n.id !== startNodeId),
        edges: edges.filter(e => !edgesToDelete.has(e.id))
      });
    } else if (startEdgeId) {
      const startEdge = edges.find(e => e.id === startEdgeId);
      if (startEdge) {
        traverse(startEdge);
      }
      
      set({
        edges: edges.filter(e => !edgesToDelete.has(e.id))
      });
    }

    get().propagateTraces();
  },

  toggleNodeDisabled: (nodeId: string) => {
    const node = get().nodes.find(n => n.id === nodeId);
    if (node) {
      const newDisabled = !node.data.disabled;
      // Update node data AND draggable property
      set(state => ({
        nodes: state.nodes.map(n => {
          if (n.id === nodeId) {
            return {
              ...n,
              draggable: !newDisabled,  // Disable dragging when node is disabled
              selectable: !newDisabled, // Disable selection when node is disabled
              data: { ...n.data, disabled: newDisabled }
            };
          }
          return n;
        })
      }));
      // Update edge styles to reflect disabled state
      setTimeout(() => get().updateEdgeStyles(), 0);
    }
  },

  archiveNode: (nodeId: string) => {
    const node = get().nodes.find(n => n.id === nodeId);
    if (!node) return;
    
    const archived: ArchivedNode = {
      id: `archive_${Date.now()}`,
      label: node.data.label,
      model: node.data.model,
      systemPrompt: node.data.systemPrompt,
      temperature: node.data.temperature,
      reasoningEffort: node.data.reasoningEffort || 'medium'
    };
    
    set(state => ({
      archivedNodes: [...state.archivedNodes, archived]
    }));
  },

  removeFromArchive: (archiveId: string) => {
    set(state => ({
      archivedNodes: state.archivedNodes.filter(a => a.id !== archiveId)
    }));
  },

  createNodeFromArchive: (archiveId: string, position: { x: number; y: number }) => {
    const archived = get().archivedNodes.find(a => a.id === archiveId);
    if (!archived) return;
    
    const newNode: LLMNode = {
      id: `node_${Date.now()}`,
      type: 'llmNode',
      position,
      data: {
        label: archived.label,
        model: archived.model,
        temperature: archived.temperature,
        systemPrompt: archived.systemPrompt,
        userPrompt: '',
        mergeStrategy: 'smart',
        reasoningEffort: archived.reasoningEffort,
        traces: [],
        outgoingTraces: [],
        forkedTraces: [],
        activeTraceIds: [],
        response: '',
        status: 'idle',
        inputs: 1
      }
    };
    
    get().addNode(newNode);
  },

  toggleTraceDisabled: (edgeId: string) => {
    const { edges, nodes } = get();
    const edge = edges.find(e => e.id === edgeId);
    if (!edge) return;
    
    // Find all nodes connected through this trace (BIDIRECTIONAL)
    const nodesInTrace = new Set<string>();
    const visitedEdges = new Set<string>();
    
    // Traverse downstream (source -> target direction)
    const traverseDownstream = (currentNodeId: string) => {
      nodesInTrace.add(currentNodeId);
      
      const outgoing = edges.filter(e => e.source === currentNodeId);
      outgoing.forEach(nextEdge => {
        if (visitedEdges.has(nextEdge.id)) return;
        visitedEdges.add(nextEdge.id);
        traverseDownstream(nextEdge.target);
      });
    };
    
    // Traverse upstream (target -> source direction)
    const traverseUpstream = (currentNodeId: string) => {
      nodesInTrace.add(currentNodeId);
      
      const incoming = edges.filter(e => e.target === currentNodeId);
      incoming.forEach(prevEdge => {
        if (visitedEdges.has(prevEdge.id)) return;
        visitedEdges.add(prevEdge.id);
        traverseUpstream(prevEdge.source);
      });
    };
    
    // Start bidirectional traversal from clicked edge
    visitedEdges.add(edge.id);
    
    // Go upstream from source (including source itself)
    traverseUpstream(edge.source);
    
    // Go downstream from target (including target itself)  
    traverseDownstream(edge.target);
    
    // Check if any node in this trace is disabled
    const anyDisabled = Array.from(nodesInTrace).some(
      nodeId => nodes.find(n => n.id === nodeId)?.data.disabled
    );
    
    // Toggle: if any disabled -> enable all, else disable all
    const newDisabledState = !anyDisabled;
    
    set(state => ({
      nodes: state.nodes.map(node => {
        if (nodesInTrace.has(node.id)) {
          return { 
            ...node, 
            draggable: !newDisabledState,
            selectable: !newDisabledState,
            data: { ...node.data, disabled: newDisabledState } 
          };
        }
        return node;
      })
    }));
    
    // Update edge styles
    get().updateEdgeStyles();
  },

  updateEdgeStyles: () => {
    const { nodes, edges } = get();
    
    const updatedEdges = edges.map(edge => {
      const sourceNode = nodes.find(n => n.id === edge.source);
      const targetNode = nodes.find(n => n.id === edge.target);
      
      const isDisabled = sourceNode?.data.disabled || targetNode?.data.disabled;
      
      return {
        ...edge,
        style: {
          ...edge.style,
          opacity: isDisabled ? 0.3 : 1,
          strokeDasharray: isDisabled ? '5,5' : undefined
        }
      };
    });
    
    set({ edges: updatedEdges });
  },

  propagateTraces: () => {
    const { nodes, edges } = get();
    
    // We need to calculate traces for each node, AND update edge colors.
    // Topological Sort
    const inDegree = new Map<string, number>();
    const graph = new Map<string, string[]>();
    
    nodes.forEach(node => {
      inDegree.set(node.id, 0);
      graph.set(node.id, []);
    });
    
    edges.forEach(edge => {
      inDegree.set(edge.target, (inDegree.get(edge.target) || 0) + 1);
      graph.get(edge.source)?.push(edge.target);
    });
    
    const topoQueue: string[] = [];
    inDegree.forEach((count, id) => {
      if (count === 0) topoQueue.push(id);
    });
    
    const sortedNodes: string[] = [];
    while (topoQueue.length > 0) {
      const u = topoQueue.shift()!;
      sortedNodes.push(u);
      
      const children = graph.get(u) || [];
      children.forEach(v => {
        inDegree.set(v, (inDegree.get(v) || 0) - 1);
        if (inDegree.get(v) === 0) {
          topoQueue.push(v);
        }
      });
    }
    
    // Map<NodeID, Trace[]>: Traces LEAVING this node
    const nodeOutgoingTraces = new Map<string, Trace[]>();
    // Map<NodeID, Trace[]>: Traces ENTERING this node (to update NodeData)
    const nodeIncomingTraces = new Map<string, Trace[]>();
    
    // Also track Edge updates (Color AND SourceHandle)
    const updatedEdges = [...edges];
    let edgesChanged = false;
    
    // Iterate
    sortedNodes.forEach(nodeId => {
      const node = nodes.find(n => n.id === nodeId);
      if (!node) return;
      
      // 1. Gather Incoming Traces
      const incomingEdges = edges.filter(e => e.target === nodeId);
      const myIncomingTraces: Trace[] = [];
      
      incomingEdges.forEach(edge => {
        const parentOutgoing = nodeOutgoingTraces.get(edge.source) || [];
        
        // Find match based on Handle ID
        // EXACT match first
        // Since we removed 'new-trace' handle, we only look for exact trace matches.
        let matchedTrace = parentOutgoing.find(t => edge.sourceHandle === `trace-${t.id}`);
        
        // If no exact match, try to find a "Semantic Match" (Auto-Reconnect)
        // If edge.sourceHandle was 'trace-X', and now we have 'trace-X_Parent', that's a likely evolution.
        if (!matchedTrace && edge.sourceHandle?.startsWith('trace-')) {
           const oldId = edge.sourceHandle.replace('trace-', '');
           matchedTrace = parentOutgoing.find(t => t.id === `${oldId}_${edge.source}`);
        }
        
        // Fallback: If still no match, and parent has traces, try to connect to the most logical one.
        // If parent has only 1 trace, connect to it.
        // This handles cases where edge.sourceHandle might be null or outdated.
        if (!matchedTrace && parentOutgoing.length > 0) {
           // If edge has no handle ID, default to the last generated trace (usually Self Trace)
           if (!edge.sourceHandle) {
              matchedTrace = parentOutgoing[parentOutgoing.length - 1];
           }
        }
        
        if (matchedTrace) {
           myIncomingTraces.push(matchedTrace);
           
           // Update Edge Visuals & Logical Connection
           const edgeIndex = updatedEdges.findIndex(e => e.id === edge.id);
           if (edgeIndex !== -1) {
             const currentEdge = updatedEdges[edgeIndex];
             const newHandleId = `trace-${matchedTrace.id}`;
             
        // Check if we need to update
        if (currentEdge.sourceHandle !== newHandleId || currentEdge.style?.stroke !== matchedTrace.color) {
            updatedEdges[edgeIndex] = {
            ...currentEdge,
            sourceHandle: newHandleId, // Auto-update handle connection!
            style: { ...currentEdge.style, stroke: matchedTrace.color, strokeWidth: 2 }
            };
            edgesChanged = true;
        }
           }
        }
      });
      
      // Deduplicate incoming traces by ID (in case multiple edges carry same trace)
      const uniqueIncoming = Array.from(new Map(myIncomingTraces.map(t => [t.id, t])).values());
      nodeIncomingTraces.set(nodeId, uniqueIncoming);
      
      // 2. Generate Outgoing Traces
      // Every incoming trace gets appended with this node's response.
      // PLUS, we always generate a "Self Trace" (Start New) that starts here.
      
      const myResponseMsg: Message[] = [];
      if (node.data.userPrompt) {
        myResponseMsg.push({ 
           id: `${node.id}-user`, // Deterministic ID for stability
           role: 'user', 
           content: node.data.userPrompt 
        });
      }
      if (node.data.response) {
        myResponseMsg.push({ 
           id: `${node.id}-assistant`, 
           role: 'assistant', 
           content: node.data.response 
        });
      }
      
      const myOutgoingTraces: Trace[] = [];
      
      // A. Pass-through traces (append history)
      uniqueIncoming.forEach(t => {
        // When a trace passes through a node and gets modified, it effectively becomes a NEW branch of that trace.
        // We must append the current node ID to the trace ID to distinguish branches.
        // e.g. Trace "root" -> passes Node A -> becomes "root_A"
        // If it passes Node B -> becomes "root_B"
        // Downstream Node D can then distinguish "root_A" from "root_B".
        
        // Match Logic:
        // We need to find if this edge was PREVIOUSLY connected to a trace that has now evolved into 'newTrace'.
        // The edge.sourceHandle might be the OLD ID.
        // We need a heuristic: if edge.sourceHandle contains the ROOT ID of this trace, we assume it's a match.
        // But this is risky if multiple branches exist.
        
        // Better heuristic:
        // When we extend a trace t -> t_new (with id t.id + '_' + node.id),
        // we record this evolution mapping.
        
        const newTraceId = `${t.id}_${node.id}`;
        
        myOutgoingTraces.push({
          ...t,
          id: newTraceId,
          messages: [...t.messages, ...myResponseMsg]
        });
      });
      
      // B. Self Trace (New Branch) -> This is the "Default" self trace (always there?)
      // Actually, if we use Manual Forks, maybe we don't need an automatic self trace?
      // Or maybe the "Default" self trace is just one of the outgoing ones.
      // Let's keep it for compatibility if downstream picks it up automatically.
      const selfTrace: Trace = {
        id: `trace-${node.id}`,
        sourceNodeId: node.id,
        color: getStableColor(node.id), 
        messages: [...myResponseMsg]
      };
      myOutgoingTraces.push(selfTrace);
      
      // C. Manual Forks
      if (node.data.forkedTraces) {
         // We need to keep them updated with the latest messages (if prompt changed)
         // But keep their IDs and Colors stable.
         const updatedForks = node.data.forkedTraces.map(fork => ({
            ...fork,
            messages: [...myResponseMsg] // Re-sync messages
         }));
         myOutgoingTraces.push(...updatedForks);
      }
      
      nodeOutgoingTraces.set(nodeId, myOutgoingTraces);
      
      // Update Node Data with INCOMING traces (for sidebar selection)
      // We store uniqueIncoming in node.data.traces
      // Note: We need to update the node in the `nodes` array, but we are inside the loop.
      // We'll do a bulk set at the end.
    });
    
    // Bulk Update Store
    set(state => ({
      edges: updatedEdges,
      nodes: state.nodes.map(n => {
        const traces = nodeIncomingTraces.get(n.id) || [];
        const outTraces = nodeOutgoingTraces.get(n.id) || [];
        return {
          ...n,
          data: {
            ...n.data,
            traces, 
            outgoingTraces: outTraces,
            activeTraceIds: n.data.activeTraceIds
          }
        };
      })
    }));
  }
}));

export default useFlowStore;