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';
  
  // 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>;

interface FlowState {
  nodes: LLMNode[];
  edges: Edge[];
  selectedNodeId: string | null;

  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;

  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,

  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();
  },

  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;