import { create } from 'zustand';
import { WorkflowDag, TaskState, TaskStatus } from '../types/workflow';

interface WorkflowStoreState {
  requestId: string | null;
  status: 'IDLE' | 'CONNECTING' | 'RUNNING' | 'COMPLETED' | 'ERROR';
  dag: WorkflowDag | null;
  tasks: Record<string, TaskState>;
  error: string | null;
  activeTab: string; // For UI linking
  
  // Actions
  initialize: (requestId: string) => void;
  setDag: (dag: WorkflowDag) => void;
  updateTaskStatus: (taskId: string, status: TaskStatus, message?: string, progress?: number) => void;
  updateTaskContent: (taskId: string, delta: string) => void; // Stream content
  appendTaskLog: (taskId: string, log: string) => void;
  setActiveTab: (tabId: string) => void;
  completeWorkflow: (result: any) => void;
  failWorkflow: (reason: string) => void;
  reset: () => void;
}

export const useWorkflowStore = create<WorkflowStoreState>((set, get) => ({
  requestId: null,
  status: 'IDLE',
  dag: null,
  tasks: {},
  error: null,
  activeTab: 'overview',

  initialize: (requestId) => set({ 
      requestId, 
      status: 'CONNECTING', 
      error: null,
      tasks: {},
      activeTab: 'overview'
  }),

  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: 'RUNNING' });
  },

  updateTaskStatus: (taskId, status, message, progress) => {
      set(state => {
          const task = state.tasks[taskId];
          if (!task) return state;

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

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

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

  appendTaskLog: (taskId, log) => {
      set(state => {
          const task = state.tasks[taskId];
          if (!task) return state;
          return {
              tasks: {
                  ...state.tasks,
                  [taskId]: {
                      ...task,
                      logs: [...task.logs, log]
                  }
              }
          };
      });
  },

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

  completeWorkflow: (_result) => set({ status: 'COMPLETED' }),
  failWorkflow: (reason) => set({ status: 'ERROR', error: reason }),
  
  reset: () => set({ 
      requestId: null, 
      status: 'IDLE', 
      dag: null, 
      tasks: {}, 
      error: null,
      activeTab: 'overview'
  })
}));
