import { create } from 'zustand';
import type { TaskSummary, TaskOutputEntry } from '../lib/api';

interface TaskState {
  // Data
  tasks: TaskSummary[];
  selectedTaskId: string | null;
  taskOutputs: Record<string, TaskOutputEntry[]>;

  // Actions
  setTasks: (tasks: TaskSummary[]) => void;
  updateTask: (taskId: string, update: Partial<TaskSummary>) => void;
  selectTask: (taskId: string | null) => void;
  appendOutput: (taskId: string, output: TaskOutputEntry) => void;
  setTaskOutputs: (taskId: string, outputs: TaskOutputEntry[]) => void;
  clearTaskOutputs: (taskId: string) => void;
}

export const useTaskStore = create<TaskState>((set) => ({
  // Initial state
  tasks: [],
  selectedTaskId: null,
  taskOutputs: {},

  // Actions
  setTasks: (tasks) =>
    set({ tasks }),

  updateTask: (taskId, update) =>
    set((state) => ({
      tasks: state.tasks.map((task) =>
        task.id === taskId ? { ...task, ...update } : task
      ),
    })),

  selectTask: (taskId) =>
    set({ selectedTaskId: taskId }),

  appendOutput: (taskId, output) =>
    set((state) => {
      const existing = state.taskOutputs[taskId] ?? [];
      // Avoid duplicates by checking ID
      if (existing.some((o) => o.id === output.id)) {
        return state;
      }
      return {
        taskOutputs: {
          ...state.taskOutputs,
          [taskId]: [...existing, output],
        },
      };
    }),

  setTaskOutputs: (taskId, outputs) =>
    set((state) => ({
      taskOutputs: {
        ...state.taskOutputs,
        [taskId]: outputs,
      },
    })),

  clearTaskOutputs: (taskId) =>
    set((state) => {
      const { [taskId]: _, ...rest } = state.taskOutputs;
      return { taskOutputs: rest };
    }),
}));

// Selectors for common use cases
export const selectSelectedTask = (state: TaskState) =>
  state.tasks.find((t) => t.id === state.selectedTaskId);

export const selectRunningTasks = (state: TaskState) =>
  state.tasks.filter((t) =>
    ['running', 'initializing', 'starting'].includes(t.status)
  );

export const selectPendingTasks = (state: TaskState) =>
  state.tasks.filter((t) => t.status === 'pending');

export const selectCompletedTasks = (state: TaskState) =>
  state.tasks.filter((t) =>
    ['done', 'failed', 'merged'].includes(t.status)
  );