import { create } from 'zustand';
import { schemas } from '../api/schema.gen';
import { WorkflowDag, TaskState, TaskStatus, WorkflowEvent, WorkflowStatus, ConnectionStatus } from '../types/workflow';

interface WorkflowStoreState {
  requestId: string | null;
  status: WorkflowStatus;
  mode: 'realtime' | 'historical';
  dag: WorkflowDag | null;
  tasks: Record<string, TaskState>;
  logs: string[]; // Global realtime logs
  error: string | null;
  activeTab: string; // For UI linking
  
  // Actions
  initialize: (requestId: string) => void;
  setMode: (mode: 'realtime' | 'historical') => void;
  setDag: (dag: WorkflowDag) => void;
  updateTaskStatus: (taskId: string, status: TaskStatus, message?: string, progress?: number, inputCommit?: string, outputCommit?: string) => void;
  updateTaskContent: (taskId: string, delta: string) => void; // Stream content (append)
  setTaskContent: (taskId: string, content: string) => void; // Set full content
  appendTaskLog: (taskId: string, log: string) => void;
  appendGlobalLog: (log: string) => void; // New action for raw global logs
  setActiveTab: (tabId: string) => void;
  completeWorkflow: (result: unknown) => void;
  failWorkflow: (reason: string) => void;
  handleEvent: (event: WorkflowEvent) => void;
  loadFromSnapshot: (snapshotPayload: any) => void;
  reset: () => void;
}

export const useWorkflowStore = create<WorkflowStoreState>((set, get) => ({
  requestId: null,
  status: ConnectionStatus.Idle,
  mode: 'realtime',
  dag: null,
  tasks: {},
  logs: [],
  error: null,
  activeTab: 'overview',

  initialize: (requestId) => set({ 
      requestId, 
      status: ConnectionStatus.Connecting, 
      mode: 'realtime',
      error: null,
      tasks: {},
      logs: [],
      activeTab: 'overview'
  }),

  setMode: (mode) => set({ mode }),

  setDag: (dag) => {
      // Initialize tasks based on DAG
      const initialTasks: Record<string, TaskState> = {};
      dag.nodes.forEach(node => {
          initialTasks[node.id] = {
              status: node.initial_status,
              logs: [],
              progress: 0,
              content: ''
          };
      });
      set({ dag, tasks: initialTasks, status: schemas.TaskStatus.enum.Running });
  },

  updateTaskStatus: (taskId, status, message, progress, inputCommit, outputCommit) => {
      set(state => {
          let task = state.tasks[taskId];
          
          // Create task if it doesn't exist (handle orphan events or pre-DAG events)
          if (!task) {
              task = {
                  status: schemas.TaskStatus.enum.Pending, // Default initial status
                  logs: [],
                  progress: 0,
                  content: ''
              };
          }

          const newLogs = [...task.logs];
          if (message) {
              newLogs.push(`[${new Date().toLocaleTimeString()}] ${message}`);
          }

          return {
              tasks: {
                  ...state.tasks,
                  [taskId]: {
                      ...task,
                      status,
                      message: message || task.message,
                      progress: progress !== undefined ? progress : task.progress,
                      logs: newLogs,
                      inputCommit: inputCommit || task.inputCommit,
                      outputCommit: outputCommit || task.outputCommit
                  }
              }
          };
      });
  },

  updateTaskContent: (taskId, delta) => {
      set(state => {
          const task = state.tasks[taskId];
          if (!task) return state;

          return {
              tasks: {
                  ...state.tasks,
                  [taskId]: {
                      ...task,
                      content: (task.content || "") + delta
                  }
              }
          };
      });
  },

  setTaskContent: (taskId, content) => {
      set(state => {
          const task = state.tasks[taskId];
          if (!task) return state;

          return {
              tasks: {
                  ...state.tasks,
                  [taskId]: {
                      ...task,
                      content
                  }
              }
          };
      });
  },

  appendTaskLog: (taskId, log) => {
      set(state => {
          let task = state.tasks[taskId];
           if (!task) {
              task = {
                  status: schemas.TaskStatus.enum.Pending,
                  logs: [],
                  progress: 0,
                  content: ''
              };
          }
          
          return {
              tasks: {
                  ...state.tasks,
                  [taskId]: {
                      ...task,
                      logs: [...task.logs, log]
                  }
              }
          };
      });
  },

  appendGlobalLog: (log) => {
      set(state => ({
          logs: [...state.logs, log]
      }));
  },

  setActiveTab: (tabId) => set({ activeTab: tabId }),

  completeWorkflow: (_result) => set({ status: schemas.TaskStatus.enum.Completed }),
  failWorkflow: (reason) => set({ status: schemas.TaskStatus.enum.Failed, error: reason }),
  
  handleEvent: (event: WorkflowEvent) => {
      const state = get();
      // Enhanced Logging (Filtered)
      if (event.type !== 'TaskStreamUpdate' && event.type !== 'TaskLog') {
          console.log(`[Store] Handling Event: ${event.type}`, event);
      }

      switch (event.type) {
          case 'WorkflowStarted':
              state.setDag(event.payload.task_graph);
              break;
          case 'TaskStateChanged': {
              const p = event.payload; 
              console.log(`[Store] Task Update: ${p.task_id} -> ${p.status}`);
              // @ts-ignore
              state.updateTaskStatus(
                  p.task_id, 
                  p.status, 
                  p.message || undefined, 
                  p.progress || undefined,
                  p.input_commit,
                  p.output_commit
              );
              break;
          }
          case 'TaskStreamUpdate': {
              const p = event.payload;
              state.updateTaskContent(p.task_id, p.content_delta);
              break;
          }
          // @ts-ignore
          case 'TaskLog': {
              const p = event.payload;
              const time = new Date(p.timestamp).toLocaleTimeString();
              const log = `[${time}] [${p.level}] ${p.message}`;
              
              // Update Task-specific logs
              state.appendTaskLog(p.task_id, log);
              
              // Update Global Raw Logs
              const globalLog = `[${time}] [${p.task_id}] [${p.level}] ${p.message}`;
              state.appendGlobalLog(globalLog);
              break;
          }
          case 'WorkflowCompleted': {
              console.log("[Store] Workflow Completed");
              state.completeWorkflow(event.payload.result_summary);
              break;
          }
          case 'WorkflowFailed': {
              console.log("[Store] Workflow Failed:", event.payload.reason);
              state.failWorkflow(event.payload.reason);
              break;
          }
          case 'WorkflowStateSnapshot': {
              // Used for real-time rehydration (e.g. page refresh)
              console.log("[Store] Processing WorkflowStateSnapshot...", event.payload);
              // First, restore DAG if present
              if (event.payload.task_graph) {
                  // WARNING: setDag resets tasks to initial state!
                  // We must be careful not to lose existing state if we are just updating.
                  // But usually Snapshot means "replace everything".
                  state.setDag(event.payload.task_graph);
              }
              
              const currentTasks = get().tasks; // These are now reset if setDag was called
              const newTasks = { ...currentTasks };
              
              const payload = event.payload as any;

              // NEW: Handle task_states (Comprehensive Snapshot)
              if (payload.task_states) {
                   Object.entries(payload.task_states).forEach(([taskId, stateSnapshot]: [string, any]) => {
                       // Merge or Create
                       const existing = newTasks[taskId] || {
                           status: schemas.TaskStatus.enum.Pending,
                           logs: [],
                           progress: 0,
                           content: ''
                       };
                       
                       newTasks[taskId] = {
                           ...existing,
                           status: stateSnapshot.status,
                           // Prefer snapshot logs if available, they are the full history
                           logs: (stateSnapshot.logs && stateSnapshot.logs.length > 0) ? stateSnapshot.logs : existing.logs,
                           // Prefer snapshot content if available
                           content: stateSnapshot.content !== undefined && stateSnapshot.content !== null ? stateSnapshot.content : existing.content,
                           inputCommit: stateSnapshot.input_commit,
                           outputCommit: stateSnapshot.output_commit,
                           metadata: stateSnapshot.metadata
                       };
                   });
              } else {
                  // Fallback / Compatibility
                  if (payload.tasks_status) {
                      Object.entries(payload.tasks_status).forEach(([taskId, status]) => {
                          if (newTasks[taskId] && status) {
                              newTasks[taskId] = { ...newTasks[taskId], status: status as TaskStatus };
                          }
                      });
                  }
                  
                  if (payload.tasks_output) {
                      Object.entries(payload.tasks_output).forEach(([taskId, outputCommit]) => {
                          if (newTasks[taskId] && outputCommit) {
                              newTasks[taskId] = { ...newTasks[taskId], outputCommit: outputCommit as string };
                          }
                      });
                  }
    
                  if (payload.tasks_metadata) {
                      Object.entries(payload.tasks_metadata).forEach(([taskId, metadata]) => {
                          if (newTasks[taskId] && metadata) {
                              newTasks[taskId] = { ...newTasks[taskId], metadata: metadata };
                          }
                      });
                  }
              }

              // Handle global log replay
              // @ts-ignore
              if (payload.logs && Array.isArray(payload.logs)) {
                  set({ logs: payload.logs });
              }
              
              set({ tasks: newTasks });
              break;
          }
      }
  },

  loadFromSnapshot: (payload: any) => {
      // Used for loading completed/archived sessions
      const dag = payload.task_graph;
      // Check if we have the new `task_states` format in the snapshot
      const taskStates = payload.task_states;

      const tasks_status = payload.tasks_status;
      const tasks_output = payload.tasks_output;
      const tasks_metadata = payload.tasks_metadata;
      
      const newTasks: Record<string, TaskState> = {};
      
      if (dag) {
          dag.nodes.forEach((node: any) => {
              if (taskStates && taskStates[node.id]) {
                  // Use new format
                  const s = taskStates[node.id];
                   newTasks[node.id] = {
                      status: s.status,
                      logs: s.logs || [],
                      progress: s.status === 'Completed' ? 100 : 0,
                      content: s.content || '', 
                      outputCommit: s.output_commit,
                      metadata: s.metadata
                   };
              } else {
                  // Legacy fallback
                  newTasks[node.id] = {
                      status: tasks_status?.[node.id] || node.initial_status,
                      logs: [],
                      progress: 100, 
                      content: '', // Content is not in legacy snapshot
                      outputCommit: tasks_output?.[node.id],
                      metadata: tasks_metadata?.[node.id]
                  };
              }
          });
      }
      
      set({ 
          dag, 
          tasks: newTasks, 
          status: schemas.TaskStatus.enum.Completed, 
          mode: 'historical',
          error: null
      });
      
      if (payload.logs) {
          set({ logs: payload.logs });
      }
  },

  reset: () => set({ 
      requestId: null, 
      status: ConnectionStatus.Idle, 
      mode: 'realtime',
      dag: null, 
      tasks: {}, 
      logs: [],
      error: null,
      activeTab: 'overview'
  })
}));
