import type {
  FileInWorkspace,
  LogEntry,
  QueueTaskItem,
  RestoreStateMessage,
  SidebarStatus,
  SidebarTab,
  Suggestion,
  UpdateQueueTasksMessage,
  UpdateStateMessage,
} from "@codeflash/types";
import { create } from "zustand";
import { handleUpdateState } from "./actions/updateSidebarState";
import { handleUpdateQueueTasks } from "./actions/updateQueueTasks";
import { handleRestoreStateFromCache } from "./actions/restoreStateFromMemory";
import { handleAddLog } from "./actions/addLog";
import { handleAddLogForDemo } from "./actions/addLogForDemo";
import { updateDemoTask } from "./actions/updateDemoTask";

export type State = {
  moduleRoot: string;
  status: SidebarStatus;
  activeTab: SidebarTab;
  optimizationRunning: boolean;
  triedCodeflashDemo: boolean;
  queueTasks: QueueTaskItem[];
  functionsInCurrentFile: string[];
  focusedTaskId: string;
  currentFilePath: string;
  filesInWorkspace: FileInWorkspace[];
  functionsInCurrentContextFile: { file: string; functions: string[] };
  isValidatingApiKey: boolean;
  projectConfigSuggestions: Record<string, Suggestion>;
  configPath?: string;
  demoOptTask: QueueTaskItem | null;
};

export type Action = {
  // main action for updating the state in the sidebar from the extension
  updateSidebarState: (msg: UpdateStateMessage) => void;
  updateQueueTasks: (msg: UpdateQueueTasksMessage) => void;
  restoreStateFromCache: (msg: RestoreStateMessage) => void;
  apiKeyEnterValid: () => void;
  updateFilesInWorkspace: (files: FileInWorkspace[]) => void;
  updateFunctionsInCurrentContextFile: (
    file: string,
    functions: string[],
  ) => void;
  changeTaskFocus: (id: string) => void;
  setActiveTab: (tab: SidebarTab) => void;
  setApiKeyLoadingState: (loading: boolean) => void;
  initDemoOptTask: (task: QueueTaskItem) => void;
  addLog: (log: LogEntry) => void;
  addLogForDemo: (log: LogEntry) => void;
  clearDemoOptTask: () => void;
  updateDemoOptTask: (taskUpdate: Partial<QueueTaskItem>) => void;
};

const initialState: State = {
  moduleRoot: "",
  status: "idle",
  activeTab: "optimization",
  optimizationRunning: false,
  queueTasks: [],
  functionsInCurrentFile: [],
  focusedTaskId: "",
  currentFilePath: "",
  filesInWorkspace: [],
  functionsInCurrentContextFile: { file: "", functions: [] },
  isValidatingApiKey: false,
  projectConfigSuggestions: {},
  triedCodeflashDemo: false,
  demoOptTask: null,
};

export const useStore = create<State & Action>((_set) => ({
  ...initialState,
  setActiveTab: (tab: SidebarTab) =>
    _set((oldState) => ({ ...oldState, activeTab: tab })),
  updateSidebarState: (message: UpdateStateMessage) =>
    _set((oldState) => handleUpdateState(oldState, message)),
  updateQueueTasks: (message: UpdateQueueTasksMessage) =>
    _set((oldState) => handleUpdateQueueTasks(oldState, message)),
  restoreStateFromCache: (message: RestoreStateMessage) =>
    _set((oldState) => handleRestoreStateFromCache(oldState, message)),
  apiKeyEnterValid: () => _set((oldState) => ({ ...oldState, status: "idle" })),
  updateFilesInWorkspace: (files: FileInWorkspace[]) =>
    _set((oldState) => ({ ...oldState, filesInWorkspace: files })),
  updateFunctionsInCurrentContextFile: (file: string, functions: string[]) =>
    _set((oldState) => ({
      ...oldState,
      functionsInCurrentContextFile: {
        file,
        functions,
      },
    })),
  changeTaskFocus: (id: string) =>
    _set((oldState) => ({ ...oldState, focusedTaskId: id })),
  addLog: (log: LogEntry) => _set((oldState) => handleAddLog(oldState, log)),
  setApiKeyLoadingState: (loading: boolean) =>
    _set((oldState) => ({ ...oldState, isValidatingApiKey: loading })),
  initDemoOptTask: (task: QueueTaskItem) =>
    _set((oldState) => ({ ...oldState, demoOptTask: task })),
  addLogForDemo: (log: LogEntry) =>
    _set((oldState) => handleAddLogForDemo(oldState, log)),
  updateDemoOptTask: (payload: Partial<QueueTaskItem>) =>
    _set((oldState) => updateDemoTask(oldState, payload)),
  clearDemoOptTask: () =>
    _set((oldState) => ({ ...oldState, demoOptTask: null })),
}));
