export type MessageType =
  | "showMessage"
  | "updateState"
  | "restoreStateFromCache"
  | "updateQueueTasks"
  | "apiKeyEnterInvalid"
  | "apiKeyEnterValid"
  | "requestAnalysis"
  | "navigateToFunction"
  | "optimizeFunction"
  | "webviewReady"
  | "createSampleFile"
  | "apiKeyEnter"
  | "viewDiff"
  | "viewPatch"
  | "rejectOptimization"
  | "changeTaskFocus"
  | "deleteTask"
  | "viewLogs"
  | "optimizeCurrentDiff"
  | "optimizeCurrentFile"
  | "newLogEntry"
  | "filesInWorkspace"
  | "requestFunctions"
  | "recievedFunctions"
  | "recievedFilesInWorkspace"
  | "changeTaskFocusFromBackend"
  | "setActiveSidebarTab"
  | "setApiKeyLoadingState"
  | "cancelTask"
  | "cancelDemoTask"
  | "submitInitForm"
  | "appendLogToDemoOptimizationTask"
  | "startDemoOptimization"
  | "clearDemoOptimizationTask"
  | "requestStartDemoOptimization"
  | "updateDemoOptimizationTask"
  | "requestSkipDemoOptimization"
  | "requestDemoOptimizationSuccess"
  | "signIn"
  | "authStarted"
  | "authCompleted"
  | "authFailed"
  | "authCancelled"
  | "submitInitForm"
  | "cancelAuth"
  | "openExternal";

export type QueueTaskItemStatus =
  | "queued"
  | "initializing"
  | "optimizing"
  | "completed"
  | "failed"
  | "canceled"
  | "skipped"
  | "rejected";

export type SidebarTab = "optimization" | "tasks" | "projectConfig";

export type QueueTaskItem = {
  id: string;
  status: QueueTaskItemStatus;
  functionName: string;
  filepath?: string;
  description?: string;
  error?: string;
  patchFile?: string;
  explanation?: string;
  speedupStr?: string;
  bestCandidateCode?: string;
  optimizationReview?: string;
  originalRuntime?: number;
  optimizedRuntime?: number;
  originalLineProfiler?: string;
  optimizedLineProfiler?: string;
  logs: LogEntry[];
  prUrl?: string; // URL of the created pull request
};

export type SidebarStatus =
  | "connecting"
  | "idle"
  | "noPythonFile"
  | "running"
  | "success"
  | "error"
  | "serverError"
  | "empty"
  | "apiKeyError";

export interface WebviewMessage {
  type: MessageType;
  payload?: unknown;
}

export interface ShowMessageMessage extends WebviewMessage {
  type: "showMessage";
  payload: {
    type: "info" | "error" | "warning";
    message: string;
  };
}
export interface AuthStartedMessage {
  type: "authStarted";
}

export interface AuthCompletedMessage {
  type: "authCompleted";
}

export interface AuthCancelledMessage {
  type: "authCancelled";
}

export interface AuthFailedMessage {
  type: "authFailed";
  payload: {
    error: string;
  };
}
export interface SubmitInitFormMessage extends WebviewMessage {
  type: "submitInitForm";
  payload: {
    formData: Record<string, unknown>;
    pyprojectPath?: string;
  };
}

export interface UpdateStateMessage extends WebviewMessage {
  type: "updateState";
  payload: {
    status: SidebarStatus;
    message: string;
    functionData: string[] | null;
    currentFilePath: string | null;
  };
}

export interface RestoreStateMessage extends WebviewMessage {
  type: "restoreStateFromCache";
  state: unknown;
}

export interface UpdateQueueTasksMessage extends WebviewMessage {
  type: "updateQueueTasks";
  payload: {
    tasks: QueueTaskItem[];
  };
}

export interface AppendLogToDemoOptimizationTaskMessage extends WebviewMessage {
  type: "appendLogToDemoOptimizationTask";
  payload: {
    log: LogEntry;
  };
}
export interface UpdateDemoOptimizationTaskMessage extends WebviewMessage {
  type: "updateDemoOptimizationTask";
  payload: Partial<QueueTaskItem>;
}

export interface ClearDemoOptimizationTaskMessage extends WebviewMessage {
  type: "clearDemoOptimizationTask";
}
export interface StartDemoOptimizationMessage extends WebviewMessage {
  type: "startDemoOptimization";
  payload: {
    task: QueueTaskItem;
  };
}

export interface SetApiKeyLoadingStateMessage extends WebviewMessage {
  type: "setApiKeyLoadingState";
  payload: {
    loading: boolean;
  };
}

export interface ApiKeyEnterInvalidMessage extends WebviewMessage {
  type: "apiKeyEnterInvalid";
}

export interface ApiKeyEnterValidMessage extends WebviewMessage {
  type: "apiKeyEnterValid";
}

export interface NewLogEntryMessage extends WebviewMessage {
  type: "newLogEntry";
  payload: {
    log: LogEntry;
  };
}

export interface RequestAnalysisMessage extends WebviewMessage {
  type: "requestAnalysis";
}

export interface NavigateToFunctionMessage extends WebviewMessage {
  type: "navigateToFunction";
  payload: {
    functionName: string;
    filePath?: string;
  };
}

export interface OptimizeFunctionMessage extends WebviewMessage {
  type: "optimizeFunction";
  payload: {
    functionName: string;
    filePath?: string;
  };
}

export interface WebviewReadyMessage extends WebviewMessage {
  type: "webviewReady";
}

export interface CreateSampleFileMessage extends WebviewMessage {
  type: "createSampleFile";
}

export interface RequestStartDemoOptimizationMessage extends WebviewMessage {
  type: "requestStartDemoOptimization";
}

export interface RequestSkipDemoOptimizationMessage extends WebviewMessage {
  type: "requestSkipDemoOptimization";
}

export interface RequestDemoOptimizationSuccessMessage extends WebviewMessage {
  type: "requestDemoOptimizationSuccess";
}

export interface ChangeTaskFocusMessage extends WebviewMessage {
  type: "changeTaskFocus";
  payload: {
    id: string;
  };
}

export interface FilesInWorkspace extends WebviewMessage {
  type: "filesInWorkspace";
}

export interface RecievedFilesInWorkspace extends WebviewMessage {
  type: "recievedFilesInWorkspace";
  payload: {
    files: FileInWorkspace[];
  };
}

export interface RequestFunctionsMessage extends WebviewMessage {
  type: "requestFunctions";
  payload: {
    filePath: string;
  };
}

export interface RecievedFunctionsMessage extends WebviewMessage {
  type: "recievedFunctions";
  payload: {
    filePath: string;
    functions: string[];
  };
}

export interface ChangeTaskFocusFromBackendMessage extends WebviewMessage {
  type: "changeTaskFocusFromBackend";
  payload: {
    id: string;
  };
}

export interface SetActiveSidebarTabMessage extends WebviewMessage {
  type: "setActiveSidebarTab";
  payload: {
    tab: SidebarTab;
  };
}

export interface ApiKeyEnterMessage extends WebviewMessage {
  type: "apiKeyEnter";
  payload: {
    apiKey: string;
  };
}
export interface SignInMessage extends WebviewMessage {
  type: "signIn";
}

export interface CancelAuthMessage extends WebviewMessage {
  type: "cancelAuth";
}

export interface ViewDiffMessage extends WebviewMessage {
  type: "viewDiff";
  payload: {
    functionName: string;
  };
}

export interface ViewPatchMessage extends WebviewMessage {
  type: "viewPatch";
  payload: {
    id: string;
    patchFile: string;
    functionName: string;
    speedupStr?: string;
    explanation?: string;
    optimizationReview?: string;
    originalRuntime: number;
    optimizedRuntime: number;
    originalLineProfiler?: string;
    optimizedLineProfiler?: string;
  };
}

export interface RejectOptimizationMessage extends WebviewMessage {
  type: "rejectOptimization";
  payload: {
    id: string;
  };
}

export interface DeleteTaskMessage extends WebviewMessage {
  type: "deleteTask";
  payload: {
    id: string;
  };
}

export interface CancelTaskMessage extends WebviewMessage {
  type: "cancelTask";
  payload: {
    id: string;
  };
}

export interface CancelDemoTaskMessage extends WebviewMessage {
  type: "cancelDemoTask";
}

export interface ViewLogsMessage extends WebviewMessage {
  type: "viewLogs";
}

export interface OptimizeCurrentDiffMessage extends WebviewMessage {
  type: "optimizeCurrentDiff";
}

export interface OptimizeCurrentFileMessage extends WebviewMessage {
  type: "optimizeCurrentFile";
}

export interface OpenExternalMessage extends WebviewMessage {
  type: "openExternal";
  payload: {
    url: string;
  };
}

export type FileInWorkspace = { abs: string; rel: string };

/** Status of an execution step */
export type ExecutionStatus =
  | "running"
  | "success"
  | "error"
  | "warning"
  | "info";

/** Type of content in a log entry */
export type LogContentType = "text" | "code" | "markdown";

export type MessageID = "best_candidate" | "candidate";
interface LspMessage {
  type: LogContentType;
  id: string;
  takes_time?: boolean;
  task_id: string;
  message_id?: MessageID;
}
export interface LSPTextMessage extends LspMessage {
  text: string;
}

export interface LSPCodeMessage extends LspMessage {
  code: string;
  collapsed?: boolean;
  function_name?: string;
  file_name?: string;
}

export interface LSPMarkdownMessage extends LspMessage {
  markdown: string;
}

export interface LSPErrorMessage extends LspMessage {
  text: string;
}

export type LogEntry =
  | LSPTextMessage
  | LSPCodeMessage
  | LSPMarkdownMessage
  | LSPErrorMessage;

export const LSPMessageType = {
  CODE: "code",
  TEXT: "text",
  MARKDOWN: "markdown",
};

export type IncomingWebviewMessage =
  | ShowMessageMessage
  | RequestAnalysisMessage
  | NavigateToFunctionMessage
  | OptimizeFunctionMessage
  | WebviewReadyMessage
  | CreateSampleFileMessage
  | ApiKeyEnterMessage
  | ViewDiffMessage
  | ViewPatchMessage
  | RejectOptimizationMessage
  | DeleteTaskMessage
  | CancelTaskMessage
  | CancelDemoTaskMessage
  | ViewLogsMessage
  | OptimizeCurrentDiffMessage
  | OptimizeCurrentFileMessage
  | OpenExternalMessage
  | FilesInWorkspace
  | RequestFunctionsMessage
  | SignInMessage
  | ChangeTaskFocusMessage
  | SubmitInitFormMessage
  | RequestStartDemoOptimizationMessage
  | RequestSkipDemoOptimizationMessage
  | RequestDemoOptimizationSuccessMessage
  | CancelAuthMessage;

export type OutgoingWebviewMessage =
  | UpdateStateMessage
  | RestoreStateMessage
  | UpdateQueueTasksMessage
  | ApiKeyEnterInvalidMessage
  | ApiKeyEnterValidMessage
  | NewLogEntryMessage
  | RecievedFilesInWorkspace
  | RecievedFunctionsMessage
  | ChangeTaskFocusFromBackendMessage
  | SetActiveSidebarTabMessage
  | SetApiKeyLoadingStateMessage
  | AppendLogToDemoOptimizationTaskMessage
  | StartDemoOptimizationMessage
  | UpdateDemoOptimizationTaskMessage
  | ClearDemoOptimizationTaskMessage
  | AuthStartedMessage
  | AuthCompletedMessage
  | AuthCancelledMessage
  | AuthFailedMessage
  | SetApiKeyLoadingStateMessage;

export type Suggestion = {
  choices: string[];
  default: string;
};
