import type {
  AppendLogToDemoOptimizationTaskMessage,
  ChangeTaskFocusFromBackendMessage,
  NewLogEntryMessage,
  RecievedFilesInWorkspace,
  RecievedFunctionsMessage,
  RestoreStateMessage,
  SetActiveSidebarTabMessage,
  SetApiKeyLoadingStateMessage,
  StartDemoOptimizationMessage,
  UpdateDemoOptimizationTaskMessage,
  UpdateQueueTasksMessage,
  UpdateStateMessage,
  WebviewMessage,
} from "@codeflash/types";
import type { Action, State } from "../store/root";

export const messageHandler = (event: MessageEvent, store: State & Action) => {
  const message = event.data;
  if (!message.type) {
    return;
  }
  const webviewMessage = message as WebviewMessage;

  switch (webviewMessage.type) {
    case "updateState":
      store.updateSidebarState(webviewMessage as UpdateStateMessage);
      break;
    case "updateQueueTasks":
      store.updateQueueTasks(webviewMessage as UpdateQueueTasksMessage);
      break;
    case "restoreStateFromCache":
      store.restoreStateFromCache(webviewMessage as RestoreStateMessage);
      break;
    case "apiKeyEnterValid":
      store.apiKeyEnterValid();
      break;
    case "recievedFilesInWorkspace":
      const files = (webviewMessage as RecievedFilesInWorkspace).payload.files;
      store.updateFilesInWorkspace(files);
      break;
    case "recievedFunctions":
      const { filePath, functions } = (
        webviewMessage as RecievedFunctionsMessage
      ).payload;
      store.updateFunctionsInCurrentContextFile(filePath, functions);
      break;
    case "changeTaskFocusFromBackend":
      const { id } = (webviewMessage as ChangeTaskFocusFromBackendMessage)
        .payload;
      store.changeTaskFocus(id);
      break;
    case "newLogEntry":
      const { log } = (webviewMessage as NewLogEntryMessage).payload;
      store.addLog(log);
      break;
    case "setActiveSidebarTab":
      const { tab } = (webviewMessage as SetActiveSidebarTabMessage).payload;
      store.setActiveTab(tab);
      break;
    case "setApiKeyLoadingState":
      const { loading } = (webviewMessage as SetApiKeyLoadingStateMessage)
        .payload;
      store.setApiKeyLoadingState(loading);
      break;
    case "startDemoOptimization":
      const { task } = (webviewMessage as StartDemoOptimizationMessage).payload;
      store.initDemoOptTask(task);
      break;
    case "updateDemoOptimizationTask":
      const payload = (webviewMessage as UpdateDemoOptimizationTaskMessage)
        .payload;
      store.updateDemoOptTask(payload);
      break;
    case "appendLogToDemoOptimizationTask":
      const { log: demoLog } = (
        webviewMessage as AppendLogToDemoOptimizationTaskMessage
      ).payload;
      store.addLogForDemo(demoLog);
      break;
    case "clearDemoOptimizationTask":
      store.clearDemoOptTask();
      break;
    default:
      console.warn("Unhandled message type:", webviewMessage.type);
  }
};
