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;
  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, 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;
  setActiveTab: (tabId: string) => void;
  completeWorkflow: (result: unknown) => void;
  failWorkflow: (reason: string) => void;
  handleEvent: (event: WorkflowEvent) => void;
  reset: () => void;
}

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

  initialize: (requestId) => set({ 
      requestId, 
      status: ConnectionStatus.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: 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]
                  }
              }
          };
      });
  },

  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();
      console.log('Handling Event:', event.type, event);

      switch (event.type) {
          case 'WorkflowStarted':
              state.setDag(event.payload.task_graph);
              break;
          case 'TaskStateChanged': {
              // Explicit typing to help TS
              const p = event.payload; 
              // @ts-ignore - input_commit/output_commit added
              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 - TaskLog is manually added to schema.gen.ts
          case 'TaskLog': {
              const p = event.payload;
              const time = new Date(p.timestamp).toLocaleTimeString();
              const log = `[${time}] [${p.level}] ${p.message}`;
              state.appendTaskLog(p.task_id, log);
              break;
          }
          case 'WorkflowCompleted': {
              state.completeWorkflow(event.payload.result_summary);
              break;
          }
          case 'WorkflowFailed': {
              state.failWorkflow(event.payload.reason);
              break;
          }
          case 'WorkflowStateSnapshot': {
              // Re-hydrate state
              if (event.payload.task_graph) {
                  state.setDag(event.payload.task_graph);
              }
              
              const currentTasks = get().tasks;
              const newTasks = { ...currentTasks };
              
              if (event.payload.tasks_status) {
                  Object.entries(event.payload.tasks_status).forEach(([taskId, status]) => {
                      if (newTasks[taskId] && status) {
                          newTasks[taskId] = { ...newTasks[taskId], status: status as TaskStatus };
                      }
                  });
              }
              
               if (event.payload.tasks_output) {
                  Object.entries(event.payload.tasks_output).forEach(([taskId, outputCommit]) => {
                      if (newTasks[taskId] && outputCommit) {
                          // Correctly mapping outputCommit, not content
                          newTasks[taskId] = { ...newTasks[taskId], outputCommit: outputCommit };
                      }
                  });
              }
              
              set({ tasks: newTasks });
              break;
          }
      }
  },

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