import * as vscode from "vscode";
import type { LanguageClient } from "vscode-languageclient/node";
import {
  CancellationTokenSource,
  State as LanguageClientState,
} from "vscode-languageclient/node";
import type { OptimizationResponse } from "../types";
import type {
  AnalysisService,
  OptimizationService,
  NavigationService,
} from "../services";
import { Logger, generateWebviewHtml } from "../utils";
import {
  LSP_COMMANDS,
  SIDEBAR_BUNDLE_DIR,
  SIDEBAR_VIEW_ID,
} from "../constants";
import { randomUUID } from "crypto";
import { GitPatchProvider } from "./GitPatchProvider";
import type {
  IncomingWebviewMessage,
  LogEntry,
  OutgoingWebviewMessage,
  QueueTaskItem,
  RestoreStateMessage,
  SidebarStatus,
  SubmitInitFormMessage,
  Suggestion,
  UpdateStateMessage,
  ViewPatchMessage,
} from "@codeflash/types";
import {
  isTaskRunning,
  canRetryTaskInQueue,
  canDeleteTaskInQueue,
} from "@codeflash/shared";
import type { ContextFiles, GlobalState } from "../globalState";
import { GlobalStateKey } from "../globalState";
import { Disposable } from "../utils/dispose";
import { readFileSync, unlinkSync } from "fs";
import type { LogsEventEmitter } from "../utils/logsEventEmitter";
import { AUTH_TYPE, ERRORS } from "./AuthConfig";
import { EventTracker } from "../telemetry/posthog";
import type { CommentThreadProvider } from "./commentThreadProvider";
import { CodeflashAuthProvider } from "./CodeflashAuthProvider";
import { ConfigKey, type ConfigService } from "../services/configService";
import { getRootWorkspaceFolder } from "../utils/rootWorkspace";
import { join } from "path";
import { POSTHOG_EVENTS } from "../telemetry/events";

export class SidebarProvider
  extends Disposable
  implements vscode.WebviewViewProvider
{
  private _view?: vscode.WebviewView | undefined;
  private readonly _logger: Logger;
  private _currentFocusUri: vscode.Uri | null = null;
  private _currentFunctionCount = 0;
  private initializedOnce = false;
  private runningDemo = false;
  private openedPathches = new Map<string, GitPatchProvider>();
  private configSuggestions?: Record<string, Suggestion> | undefined =
    undefined;
  private pyprojectPath: string | undefined = undefined;

  constructor(
    private readonly _extensionUri: vscode.Uri,
    private readonly _lspClient: LanguageClient,
    private readonly _optimizationService: OptimizationService,
    private readonly _analysisService: AnalysisService,
    private readonly _navigationService: NavigationService,
    private readonly optimizationEventEmitter: LogsEventEmitter,
    private readonly globalState: GlobalState,
    private readonly configService: ConfigService,
    private readonly commentThreadProvider: CommentThreadProvider,
    private readonly scheduleCodelensRefresh: () => void,
  ) {
    super();
    this._logger = new Logger("Codeflash Sidebar");
    const lspStateChangeDisposable = this._lspClient.onDidChangeState(
      this.handleLspStateChange,
    );
    this._disposables.push(lspStateChangeDisposable);
  }

  set projectConfigs(payload: {
    config: Record<string, Suggestion>;
    pyprojectPath: string;
  }) {
    this.configSuggestions = payload.config;
    this.pyprojectPath = payload.pyprojectPath;
    this.syncUIState();
  }

  public resolveWebviewView(
    webviewView: vscode.WebviewView,
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    _context: vscode.WebviewViewResolveContext,
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    _token: vscode.CancellationToken,
  ): void {
    this._view = webviewView;

    webviewView.webview.options = {
      enableScripts: true,
      localResourceRoots: [
        vscode.Uri.joinPath(this._extensionUri, "media"),
        vscode.Uri.joinPath(this._extensionUri, SIDEBAR_BUNDLE_DIR, "assets"),
      ],
    };

    webviewView.webview.html = generateWebviewHtml(
      webviewView.webview,
      this._extensionUri,
    );
    const messageListener = webviewView.webview.onDidReceiveMessage((m) =>
      this.handleWebviewMessage(m),
    );
    const optimizationEventListener = this.optimizationEventEmitter.event(
      (log) => {
        if (log.task_id === this.optimizationEventEmitter.getCurrentTaskId()) {
          this.addLogEntryToUI(log);
        }
      },
    );

    this._disposables.push(
      webviewView.onDidChangeVisibility(this.handleVisibilityChange),
      vscode.window.onDidChangeActiveTextEditor(this.handleEditorChange),
      optimizationEventListener,
      messageListener,
    );
  }

  private addLogEntryToUI(log: LogEntry) {
    this.sendMessage({
      type: this.runningDemo
        ? "appendLogToDemoOptimizationTask"
        : "newLogEntry",
      payload: {
        log,
      },
    });
  }

  sendWebviewApiKeyError() {
    this.updateWebviewState("apiKeyError", "Invalid API key.", null, null);
  }

  private retryUntilLspRunning = async (
    callback: () => Promise<void>,
    maxRetries = 5,
  ): Promise<void> => {
    const client = this._lspClient;

    const isLspRunning = client.state === LanguageClientState.Running;
    if (!isLspRunning && maxRetries > 0) {
      await new Promise((resolve) => setTimeout(resolve, 1000));
      return this.retryUntilLspRunning(callback, maxRetries - 1);
    }
    if (!isLspRunning) {
      throw new Error(
        `LSP is not running for ${callback.name}, current state: ${client.state.toString()}`,
      );
    }
    callback();
  };

  // eslint-disable-next-line @typescript-eslint/require-await
  private handleInitialState = async (): Promise<void> => {
    try {
      if (!this.globalState.get(GlobalStateKey.UserID)) {
        this.sendWebviewApiKeyError();
        return;
      }

      // get the initial active file
      const active = this.mustGetActivePythonFileUri();
      if (active) {
        this._currentFocusUri = active;
        this.handlePotentialAnalysis(active);
      }
    } finally {
      this.initializedOnce = true;
    }
  };

  private mustGetActivePythonFileUri(): vscode.Uri | null {
    const activeEditor = vscode.window.activeTextEditor;
    if (activeEditor && activeEditor.document.languageId === "python") {
      return activeEditor.document.uri;
    }
    const visibleEditors = vscode.window.visibleTextEditors;
    for (const editor of visibleEditors) {
      const uri = this.getPythonFileUri(editor);
      if (!uri) {
        continue;
      }
      return uri;
    }
    return null;
  }

  private handleLspStateChange = (e: {
    newState: LanguageClientState;
  }): void => {
    if (e.newState === LanguageClientState.Running && this._currentFocusUri) {
      this.handlePotentialAnalysis(this._currentFocusUri);
    }
  };

  private handleWebviewMessage = async (
    data: IncomingWebviewMessage,
  ): Promise<void> => {
    if (
      !this.globalState.get(GlobalStateKey.UserID) &&
      !["apiKeyEnter", "webviewReady", "signIn", "cancelAuth"].includes(
        data.type,
      ) &&
      this.initializedOnce
    ) {
      this._logger.info(
        "API key is invalid, not going to process the webview message.",
      );
      this.sendWebviewApiKeyError();
      return;
    }
    switch (data.type) {
      case "requestAnalysis":
        await this.handleRequestAnalysis();
        break;
      case "webviewReady":
        this._logger.debug("Webview reported ready. restoring previous state");
        this.retryUntilLspRunning(this.handleInitialState.bind(this));
        this.syncUIState();
        break;
      case "navigateToFunction":
        await this.handleNavigateToFunction(
          data.payload.functionName,
          data.payload.filePath,
        );
        break;
      case "optimizeFunction":
        this.addFunctionToQueue(
          data.payload.functionName,
          data.payload.filePath,
        );
        break;
      case "createSampleFile":
        await this.handleCreateSampleFile();
        break;
      case "apiKeyEnter":
        await this.handleApiKeyEnter(data.payload.apiKey, "api_key");
        break;
      case "viewPatch":
        this.handleViewPatch(data.payload);
        break;
      case "deleteTask":
        this.handleDeleteTask(data.payload.id);
        break;
      case "cancelTask":
        this.handleCancelTask(data.payload.id);
        break;
      case "cancelDemoTask":
        if (!this.runningDemo) {
          break;
        }
        this.runningDemo = false;
        await this.handleSkipDemoOptimization();
        this._optimizationService.cancelDemoTask();
        break;
      case "optimizeCurrentDiff":
        await this.handleOptimizeCurrentDiff();
        break;
      case "viewLogs":
        this.handleViewLogs();
        break;
      case "filesInWorkspace":
        this.requestFilesInWorkspace();
        break;
      case "requestFunctions":
        const filePath = data.payload.filePath;
        this.requestFunctionsInFile(filePath);
        break;
      case "changeTaskFocus":
        this.globalState.set(GlobalStateKey.FocusedTaskId, data.payload.id);
        break;
      case "showMessage":
        const { type, message } = data.payload;
        if (type === "info") {
          vscode.window.showInformationMessage(message);
        } else if (type === "error") {
          vscode.window.showErrorMessage(message);
        } else if (type === "warning") {
          vscode.window.showWarningMessage(message);
        }
        break;
      case "requestStartDemoOptimization":
        await this.handleStartDemoOptimization();
        break;
      case "requestSkipDemoOptimization":
        this.handleSkipDemoOptimization();
        break;
      case "requestDemoOptimizationSuccess":
        this.handleDemoOptimizationSuccess();
        break;
      case "signIn":
        await this.handleSignIn();
        break;
      case "cancelAuth":
        CodeflashAuthProvider.incommingCancellationTokenSource?.cancel();
        break;
      case "submitInitForm":
        this.handleSubmitInitForm(data);
        break;
      case "optimizeCurrentFile":
        await this.handleOptimizeAllInCurrentFile();
        break;
      case "openExternal":
        if (data.payload?.url) {
          vscode.env.openExternal(vscode.Uri.parse(data.payload.url));
        }
        break;
    }
  };

  private async handleSubmitInitForm(incomingMessage: SubmitInitFormMessage) {
    const { formData } = incomingMessage.payload;

    const workspaceRoot = getRootWorkspaceFolder()?.uri.fsPath || "";
    const defaultPyprojectPath = join(workspaceRoot, "pyproject.toml");

    try {
      const result = await this._lspClient.sendRequest<{
        status: string;
        message?: string;
      }>(LSP_COMMANDS.WRITE_CONFIG, {
        config_file: this.pyprojectPath || defaultPyprojectPath,
        config: formData,
      });
      if (result.status !== "success") {
        vscode.window.showErrorMessage(
          result.message || "Failed to write configuration",
        );
      }
      // restart
      vscode.commands.executeCommand("codeflash.reloadExtension");
    } catch (error) {
      vscode.window.showErrorMessage(
        (error as Error)?.message || "Failed to write configuration",
      );
    }
  }

  private async handleAfterPatchApplied(
    id: string,
    patchFile: string,
    gotoFunction: boolean = true,
  ): Promise<void> {
    const task = this.globalState
      .get(GlobalStateKey.QueueTasks, [])
      ?.find((t) => t.id === id);
    if (!task) {
      return;
    }

    // TODO: get the trace id from the lsp
    await EventTracker.capture(POSTHOG_EVENTS.OPTIMIZATION_ACCEPTED, {});

    if (task.filepath && gotoFunction) {
      this._navigationService.navigateToFunction(
        task.functionName,
        vscode.Uri.file(task.filepath),
      );
    }

    const newContextFiles = this.globalState.get(
      GlobalStateKey.ContextFiles,
      {},
    )!;
    delete newContextFiles[id];
    unlinkSync(patchFile);

    this.globalState.set(GlobalStateKey.ContextFiles, newContextFiles);
    this.updateUIQueueTasks("remove", { id });
    this.commentThreadProvider.scheduleRefresh();
  }

  async applyPatchDirectlyForInlineThread(
    payload: ViewPatchMessage["payload"],
  ): Promise<void> {
    const { id, patchFile } = payload;

    const onPatchApplied = () =>
      this.handleAfterPatchApplied(id, patchFile, false);

    const newPatchProvider = new GitPatchProvider(
      this._extensionUri,
      this.globalState,
      payload,
      this._logger,
      onPatchApplied,
    );

    await newPatchProvider.applyPatch();
  }

  async rejectOptimization(taskId: string): Promise<void> {
    await EventTracker.capture(POSTHOG_EVENTS.OPTIMIZATION_REJECTED, {});
    const task = this.globalState
      .get(GlobalStateKey.QueueTasks, [])
      ?.find((t) => t.id === taskId);
    if (!task || task.status !== "completed") {
      return;
    }

    this.updateUIQueueTasks("update", {
      id: taskId,
      status: "rejected",
      description: "Optimization was rejected",
    });
    this.deleteContextFilesForTask(taskId);
    this.scheduleCodelensRefresh();
    this.commentThreadProvider.scheduleRefresh();
    this.changeTaskFocus("");
  }

  private deleteContextFilesForTask(taskId: string): void {
    const newContextFiles = this.globalState.get(
      GlobalStateKey.ContextFiles,
      {},
    )!;
    if (!newContextFiles[taskId]) {
      return;
    }
    delete newContextFiles[taskId];
    this.globalState.set(GlobalStateKey.ContextFiles, newContextFiles);
  }

  private async handleSignIn(): Promise<void> {
    this.sendMessage({
      type: "authStarted",
    });

    this._logger.info("Starting OAuth authentication flow");

    try {
      const session = await vscode.authentication.getSession(AUTH_TYPE, [], {
        createIfNone: true,
      });

      if (session?.accessToken) {
        await this.handleApiKeyEnter(
          session.accessToken,
          "with_webapp_redirect",
        );

        this.sendMessage({
          type: "authCompleted",
        });
      }
    } catch (error) {
      const errorMessage =
        error instanceof Error ? error.message : String(error);

      if (
        errorMessage.includes("User cancelled") ||
        errorMessage.includes("Cancelled") ||
        errorMessage === ERRORS.USER_CANCELLED
      ) {
        this._logger.info("OAuth authentication cancelled by user");
        this.sendMessage({
          type: "authCancelled",
        });
      } else if (
        errorMessage.includes("timed out") ||
        errorMessage === ERRORS.TIMED_OUT
      ) {
        this._logger.error("OAuth authentication timed out");
        this.sendMessage({
          type: "authFailed",
          payload: { error: "Authentication timed out" },
        });
      } else {
        this._logger.error(`OAuth authentication failed: ${errorMessage}`);
        this.sendMessage({
          type: "authFailed",
          payload: { error: errorMessage },
        });
      }
    }
  }

  async handleViewPatch(payload: ViewPatchMessage["payload"]): Promise<void> {
    const { id, patchFile } = payload;

    if (this.openedPathches.has(patchFile)) {
      this.openedPathches.get(patchFile)?.showPatch();
      return;
    }

    const onPatchApplied = () => this.handleAfterPatchApplied(id, patchFile);

    const newPatchProvider = new GitPatchProvider(
      this._extensionUri,
      this.globalState,
      payload,
      this._logger,
      onPatchApplied,
    );

    this.openedPathches.set(patchFile, newPatchProvider);
    await newPatchProvider.showPatch();
  }
  private async requestFunctionsInFile(
    filePath: string,
    functions?: string[],
  ): Promise<void> {
    try {
      let finalFunctions = functions;

      if (!finalFunctions) {
        const analysisResult =
          await this._analysisService.getOptimizableFunctions(
            vscode.Uri.file(filePath),
          );
        finalFunctions = analysisResult.functions;

        if (
          process.platform === "win32" &&
          (!finalFunctions || finalFunctions.length === 0)
        ) {
          const fallbackFunctions =
            await this.extractTopLevelFunctionsFromFile(filePath);
          if (fallbackFunctions.length > 0) {
            this._logger.debug(
              `Fallback parser detected ${fallbackFunctions.length} functions for ${filePath}`,
            );
            finalFunctions = fallbackFunctions;
          }
        }
      }

      this.sendMessage({
        type: "recievedFunctions",
        payload: {
          filePath,
          functions: finalFunctions,
        },
      });
    } catch (error) {
      this._logger.debug(
        "Failed to get optimizable functions :" + (error as Error)?.message ||
          "",
      );
    }
  }

  private async extractTopLevelFunctionsFromFile(
    filePath: string,
  ): Promise<string[]> {
    try {
      const document = await vscode.workspace.openTextDocument(
        vscode.Uri.file(filePath),
      );
      const text = document.getText();
      const lines = text.split(/\r?\n/);
      const detectedFunctions: string[] = [];
      const seen = new Set<string>();

      for (const line of lines) {
        const trimmedLeft = line.replace(/^\s+/, "");
        const indentLength = line.length - trimmedLeft.length;
        if (indentLength !== 0) {
          continue;
        }

        const match = trimmedLeft.match(
          /^(?:async\s+)?def\s+([A-Za-z_]\w*)\s*\(/,
        );
        if (match) {
          const fnName = match[1]!;
          if (!seen.has(fnName)) {
            seen.add(fnName);
            detectedFunctions.push(fnName);
          }
        }
      }

      return detectedFunctions;
    } catch (error) {
      this._logger.debug(
        `Fallback parsing failed for ${filePath}: ${(error as Error)?.message || error}`,
      );
      return [];
    }
  }

  private async requestFilesInWorkspace(): Promise<void> {
    let pathPrefix = "";
    const moduleRoot = this.globalState.get(GlobalStateKey.ModuleRoot);
    if (moduleRoot) {
      const relativeModuleRoot = vscode.workspace.asRelativePath(
        vscode.Uri.file(moduleRoot),
      );
      if (relativeModuleRoot !== moduleRoot) {
        pathPrefix = relativeModuleRoot + "/";
      }
    }
    this._logger.info(
      `requesting files in workspace with prefix: ${pathPrefix}`,
    );
    const files = await vscode.workspace.findFiles(
      pathPrefix + "**/*.py",
      "**/site-packages/**",
    );
    const filesPayload = files.map((f) => {
      return {
        abs: f.fsPath,
        rel: vscode.workspace.asRelativePath(f),
      };
    });
    this.sendMessage({
      type: "recievedFilesInWorkspace",
      payload: {
        files: filesPayload,
      },
    });
  }

  private async handleOptimizeCurrentDiff(): Promise<void> {
    this._logger.info("Optimizing current diff");
    const result = await this._lspClient.sendRequest<OptimizationResponse>(
      LSP_COMMANDS.GET_OPTIMIZABLE_FUNCTIONS_IN_CURRENT_DIFF,
      {},
    );
    if (result.status === "success") {
      const functions = result.functions as string[];
      let added = 0;
      for (const [path, qualifiedNames] of Object.entries(functions)) {
        for (const name of qualifiedNames) {
          const addedToQueue = this.addFunctionToQueue(name, path);
          if (addedToQueue) {
            added++;
          }
        }
      }

      const msg =
        added > 0
          ? `Detected ${added} functions to optimize.`
          : "No optimizable functions in your current modified code";

      vscode.window.showInformationMessage(msg);
      return;
    }
    vscode.window.showErrorMessage("Failed to optimize current diff");
  }

  private async handleApiKeyEnter(
    apiKey: string,
    method: "api_key" | "with_webapp_redirect" = "api_key",
  ): Promise<void> {
    if (!this._lspClient) {
      vscode.window.showErrorMessage("Language Server not running.");
      return;
    }
    if (!apiKey?.trim()) {
      return;
    }
    this.sendMessage({
      type: "setApiKeyLoadingState",
      payload: {
        loading: true,
      },
    });
    try {
      const result = await this._lspClient.sendRequest<OptimizationResponse>(
        LSP_COMMANDS.PROVIDE_API_KEY,
        { api_key: apiKey },
      );
      if (result.status === "success") {
        this._logger.info("API key is valid.");
        const userId = result.user_id as string;
        this.globalState.set(GlobalStateKey.UserID, userId);
        EventTracker.setUserId(userId);
        await EventTracker.capture(POSTHOG_EVENTS.USER_LOGGED_IN, {
          loginMethod: method,
        });
        // lazy option to restart the extension host, uncomment if we have issues with initializing the state
        // await vscode.commands.executeCommand("workbench.action.restartExtensionHost");
        this.handleInitialState();
        this.sendMessage({
          type: "apiKeyEnterValid",
        });
      } else {
        vscode.window.showErrorMessage("Invalid API key.");
        this.sendMessage({
          type: "apiKeyEnterInvalid",
        });
      }
    } finally {
      this.sendMessage({
        type: "setApiKeyLoadingState",
        payload: {
          loading: false,
        },
      });
    }
  }

  private async handleStartDemoOptimization(): Promise<void> {
    try {
      await EventTracker.capture(POSTHOG_EVENTS.DEMO_STARTED, {});

      this.runningDemo = true;
      const taskId = randomUUID().toString();
      const functionName = "find_common_tags";

      const demoTask: QueueTaskItem = {
        id: taskId,
        status: "optimizing",
        functionName,
        filepath: `${functionName}.py`,
        logs: [],
      };

      this.sendMessage({
        type: "startDemoOptimization",
        payload: {
          task: demoTask,
        },
      });

      const initLog = this.getInitLog(taskId);
      this.sendMessage({
        type: "appendLogToDemoOptimizationTask",
        payload: {
          log: initLog,
        },
      });

      const result =
        await this._optimizationService.runDemoOptimizationTask(functionName);
      if (result.status === "success") {
        this.sendMessage({
          type: "updateDemoOptimizationTask",
          payload: {
            status: "completed",
          },
        });
        this.globalState.set(GlobalStateKey.TriedCodeflashDemo, true);
        await EventTracker.capture(POSTHOG_EVENTS.DEMO_COMPLETED, {});
      } else if (result.status === "canceled") {
        // demo was canceled
      } else {
        // demo failed
      }
    } finally {
      this.runningDemo = false;
    }
  }

  private handleDemoOptimizationSuccess(): void {
    this.globalState.set(GlobalStateKey.TriedCodeflashDemo, true);
    this.syncUIState();
    this.sendMessage({
      type: "clearDemoOptimizationTask",
    });
    this.scheduleCodelensRefresh();
  }

  private async handleSkipDemoOptimization(): Promise<void> {
    await EventTracker.capture(POSTHOG_EVENTS.DEMO_SKIPPED, {});
    this.globalState.set(GlobalStateKey.TriedCodeflashDemo, true);
    this.syncUIState();
    this.scheduleCodelensRefresh();
  }

  private syncUIState(): void {
    const message: RestoreStateMessage = {
      type: "restoreStateFromCache",
      state: {
        ...this.globalState.getStateForWebview(),
        configSuggestions: this.configSuggestions,
        configPath: this.pyprojectPath,
      },
    };
    this.sendMessage(message);
  }

  private async handleRequestAnalysis(): Promise<void> {
    if (this._lspClient.state !== LanguageClientState.Running) {
      this.updateWebviewState(
        "serverError",
        "Cannot analyze: Server not running.",
        null,
        null,
      );
      return;
    }
    if (!this._currentFocusUri) {
      return;
    }
    await this.requestAnalysisFromServer(this._currentFocusUri);
  }

  private getInitLog(taskId: string): LogEntry {
    return {
      id: randomUUID(),
      type: "text",
      text: "initializing codeflash optimization process",
      takes_time: true,
      task_id: taskId,
    };
  }
  private async handleNavigateToFunction(
    functionName: string,
    fileUri?: string,
  ): Promise<void> {
    let finalUri = this._currentFocusUri;
    if (fileUri) {
      finalUri = vscode.Uri.file(fileUri);
    }
    if (!finalUri) {
      vscode.window.showErrorMessage("No active file to navigate to function.");
      return;
    }

    await this._navigationService.navigateToFunction(functionName, finalUri);
  }

  private async handleCreateSampleFile(): Promise<void> {
    this._logger.debug("Creating sample file from sidebar");
    await vscode.commands.executeCommand("codeflash.createSampleFile");
  }

  handleCancelTask(taskId: string): void {
    const task = this.globalState
      .get(GlobalStateKey.QueueTasks, [])
      ?.find((task) => task.id === taskId);
    if (!task) {
      vscode.window.showWarningMessage("Task not found");
      return;
    }
    if (!isTaskRunning(task)) {
      return;
    }
    this._optimizationService.cancelTask(task);
    this.updateUIQueueTasks("update", {
      id: taskId,
      status: "canceled",
      description: "Optimization was canceled",
    });
    this.deleteContextFilesForTask(taskId);
    this.changeTaskFocus("");
  }

  handleDeleteTask(taskId: string): void {
    const task = this.globalState
      .get(GlobalStateKey.QueueTasks, [])
      ?.find((task) => task.id === taskId);

    if (!task) {
      vscode.window.showWarningMessage(`Cannot delete task. Task not found.`);
      return;
    }
    this._logger.debug(`Deleting ${task.status} task: ${taskId}`);

    if (!canDeleteTaskInQueue(task, true)) {
      vscode.window.showWarningMessage(`Cannot delete ${task.status} task`);
      return;
    }

    // Remove from queue
    this._optimizationService.abortTask(taskId);
    this.updateUIQueueTasks("remove", { id: taskId });
    this.deleteContextFilesForTask(taskId);
  }

  private handleViewLogs(): void {
    this._logger.debug("Opening Codeflash output panel");

    // Get the Codeflash output channel and show it
    const outputChannel = this._logger._outputChannel;
    outputChannel.show(true); // true = preserve focus on current editor
  }

  changeTaskFocus(taskId: string): void {
    this.globalState.set(GlobalStateKey.FocusedTaskId, taskId);
    this.sendMessage({
      type: "changeTaskFocusFromBackend",
      payload: {
        id: taskId,
      },
    });
  }

  addFunctionToQueue(
    functionName: string,
    uri?: vscode.Uri | string,
    focusOnView: boolean = false,
  ): boolean {
    let filePathStr = "";
    if (typeof uri === "string") {
      filePathStr = uri;
    } else {
      const _uri = uri ?? this._currentFocusUri;
      filePathStr = _uri?.fsPath || "";
    }
    if (!filePathStr) {
      return false;
    }

    const existingTask = this.globalState
      .get(GlobalStateKey.QueueTasks, [])
      ?.find(
        (task) =>
          task.filepath === filePathStr && task.functionName === functionName,
      );

    if (existingTask) {
      if (!canRetryTaskInQueue(existingTask)) {
        const msg = `Optimization for '${functionName}' in '${filePathStr}' is already ${existingTask.status}`;
        this._logger.debug(msg);
        vscode.window.showInformationMessage(msg);
        return false;
      } else {
        this._optimizationService.abortTask(existingTask.id);
        this.updateUIQueueTasks("remove", { id: existingTask.id });
      }
    }

    if (focusOnView) {
      vscode.commands.executeCommand(
        "workbench.view.extension." + SIDEBAR_VIEW_ID,
      );
      this._view?.show?.();
      const runningTask = this.globalState
        .get(GlobalStateKey.QueueTasks, [])
        ?.find(isTaskRunning);
      this._logger.debug(
        `focusing on sidebar view, running task: ${runningTask}`,
      );

      // set the active tab to the main optimization tab
      this.sendMessage({
        type: "setActiveSidebarTab",
        payload: {
          // if there is an already running task show the queued new task in the tasks tab, if not the newly added task will start right away so show the main tab
          tab: runningTask ? "tasks" : "optimization",
        },
      });
    }

    const taskId = randomUUID().toString();
    const abortController = new AbortController();
    const cancelationToken = new CancellationTokenSource();
    this.updateUIQueueTasks("add", {
      id: taskId,
      status: "queued",
      functionName,
      filepath: filePathStr,
    });

    this._optimizationService.addOptimizationTask(taskId, functionName, {
      abortController,
      cancelationToken,
      onTaskStart: async (id): Promise<boolean> => {
        await EventTracker.capture(POSTHOG_EVENTS.OPTIMIZATION_STARTED, {});
        this.optimizationEventEmitter.setCurrrentTaskId(id);
        this.updateUIQueueTasks("update", {
          id,
          status: "initializing",
        });
        const task = this.globalState
          .get(GlobalStateKey.QueueTasks, [])
          ?.find((t) => t.id === id);
        if (!task) {
          return false;
        }
        this.changeTaskFocus(id);
        const initLog: LogEntry = {
          task_id: id,
          id: randomUUID(),
          type: "text",
          text: "initializing codeflash optimization process",
          takes_time: true,
        };

        this.optimizationEventEmitter.addSingleBufferedLogEntry(initLog);
        this.addLogEntryToUI(initLog);

        this._logger.debug(
          `Starting optimization process for function '${functionName}' via sidebar`,
        );
        this.scheduleCodelensRefresh();
        const initResult =
          await this._optimizationService.initializeFunctionOptimization(
            id,
            functionName,
            filePathStr,
          );

        if (initResult.status === "error") {
          const errorMsg = initResult.message || "";
          vscode.window.showErrorMessage(
            `Failed to initialize optimization ${errorMsg ? " :" + errorMsg : ""}`,
          );
          // remove the task from the queue since it failed to initialize
          this.updateUIQueueTasks("remove", { id });
          if (errorMsg.toLowerCase().includes("not found")) {
            vscode.window.showErrorMessage(
              `Function '${functionName}' not found in file '${filePathStr}, re-analyzing the file...`,
            );
            // function is not found in the file, try to re-analyze the file
            await this.handleRequestAnalysis();
          }
          return false;
        }

        const filesInContext: string[] = initResult.files_inside_context || [];
        // store the initial content of each file in the optimization context to act as a common ancestor when performing 3-way merge later if we get a successful optimization
        this.perserveContextFilesSnapshot(id, filesInContext);

        this.globalState.set(GlobalStateKey.Running, true);
        return true;
      },
      onTaskCompleted: async (result, id) => {
        this.globalState.set(GlobalStateKey.Running, false);

        if (result.canceled) {
          return;
        }
        await EventTracker.capture(POSTHOG_EVENTS.OPTIMIZATION_FOUND, {
          speedup: result.speedup || "",
        });
        const reviewMsg = result.optimizationReview
          ? `Optimization completed successfully for function '${functionName}'.\n 🎯 Optimization Quality: ${result.optimizationReview}`
          : `Optimization completed successfully for function '${functionName}'`;
        vscode.window.showInformationMessage(reviewMsg);
        const resultRes = {
          patchFile: result.patch_file!,
          explanation: result.explanation!,
          speedupStr: result.speedup || "",
          optimizationReview: result.optimizationReview || "",
          originalRuntime: result.original_runtime || 0,
          optimizedRuntime: result.optimized_runtime || 0,
          originalLineProfiler: result.original_line_profiler || "",
          optimizedLineProfiler: result.optimized_line_profiler || "",
        };
        const task = this.updateUIQueueTasks("update", {
          id,
          status: "completed",
          ...resultRes,
        });
        if (!task) {
          return;
        }
        this.commentThreadProvider.scheduleRefresh();
        if (
          this.configService.get(ConfigKey.AutoOpenOptimizationViewOnComplete)
        ) {
          // show the patch immediately after the task is completed
          this.handleViewPatch({
            ...resultRes,
            id: task.id,
            functionName: task.functionName,
          });
        }

        this.scheduleCodelensRefresh();
      },
      onTaskError: async (error, id) => {
        this.globalState.set(GlobalStateKey.Running, false);
        const msg = error instanceof Error ? error.message : String(error);
        let displayedErrorMsg = msg;

        if (msg.includes("NO OPTIMIZATIONS GENERATED")) {
          displayedErrorMsg = "No optimizations are generated";
        } else if (msg.includes("No best optimizations found for")) {
          // optimizations are generated but none are good enough, didn't pass the threshold
          displayedErrorMsg = "No best optimization found";
        } else if (msg.includes("NO TESTS GENERATED")) {
          displayedErrorMsg = "No tests generated";
        }
        // else if (
        //   msg.includes("The threshold for test coverage was not met")
        // ) {
        // }
        await EventTracker.capture(POSTHOG_EVENTS.NO_OPTIMIZATION_FOUND, {
          reason: displayedErrorMsg,
        });
        this.updateUIQueueTasks("update", {
          id,
          status: "skipped",
          description: displayedErrorMsg,
          error: displayedErrorMsg,
        });
      },
      updateQueueTask: (payload) => {
        this.updateUIQueueTasks("update", { id: taskId, ...payload });
      },
    });
    return true;
  }

  // Important: the files array here are paths of worktree, so make sure the worktree still exist before calling this function
  private perserveContextFilesSnapshot(
    taskId: string,
    filesInContext: string[],
  ): void {
    const resolvedContextFiles: ContextFiles = {};
    for (const file of filesInContext) {
      try {
        // The LSP should now return properly resolved paths, but handle edge cases
        const content = readFileSync(file, "utf-8");
        resolvedContextFiles[file] = content;
      } catch (error) {
        this._logger.error(
          `Failed to read context file: ${file}. This may cause issues during optimization.`,
          error instanceof Error ? error : undefined,
        );
        // Continue with other files instead of failing completely
        // The optimization might still work if this is just a helper file
      }
    }
    this.globalState.set(GlobalStateKey.ContextFiles, {
      ...this.globalState.get(GlobalStateKey.ContextFiles, {}),
      [taskId]: resolvedContextFiles,
    });
  }

  private updateUIQueueTasks(
    mode: "add" | "remove" | "update",
    task: Partial<QueueTaskItem>,
  ): QueueTaskItem | null {
    // flush to keep the logs in sync
    this.optimizationEventEmitter.flush(false); // don't reset current task id yet, logs may still come in for this task
    const existingTasks =
      this.globalState.get(GlobalStateKey.QueueTasks, []) || [];

    this._logger.debug(
      `queue state changed (${mode}) : ${JSON.stringify(task)}`,
    );
    let idx = -1;

    if (mode === "add") {
      if (!task.id) {
        return null;
      }
      task.logs = [];
      existingTasks.push(task as QueueTaskItem);
      idx = existingTasks.length - 1;
    } else if (mode === "update") {
      const index = existingTasks.findIndex((t) => t.id === task.id);
      if (index !== -1) {
        idx = index;
        const oldTask = existingTasks[index];
        if (!oldTask) {
          return null;
        }
        existingTasks[index] = {
          ...oldTask,
          ...task,
          id: oldTask.id, // avoid changing id by mistake
        };
      }
    } else if (mode === "remove") {
      this.deleteContextFilesForTask(task.id!);
      const index = existingTasks.findIndex((t) => t.id === task.id);
      if (index !== -1) {
        existingTasks.splice(index, 1);
      }
    }
    this.globalState.set(GlobalStateKey.QueueTasks, existingTasks);

    this.scheduleCodelensRefresh();
    this.sendMessage({
      type: "updateQueueTasks",
      payload: {
        tasks: existingTasks,
      },
    });
    return idx !== -1 ? existingTasks[idx]! : null;
  }

  private handleVisibilityChange = async (): Promise<void> => {
    if (!this._view?.visible) {
      return;
    }
    this._logger.debug("View became visible. Re-evaluating state.");
    await this.handleInitialState();
  };

  private handleEditorChange = async (
    editor: vscode.TextEditor | undefined,
  ): Promise<void> => {
    this._logger.debug(
      `Editor changed. New URI: ${editor?.document.uri.toString() ?? "None"}`,
    );

    const uri = this.getPythonFileUri(editor);
    if (!uri) {
      return;
    }

    this._currentFocusUri = uri;

    // If Python file is open, check LSP state
    const lspState = this._lspClient.state;
    if (lspState === LanguageClientState.Running) {
      this._logger.debug("!!!LSP is running. Requesting analysis.");
      await this.handlePotentialAnalysis(uri);
    } else if (lspState === LanguageClientState.Starting) {
      this.updateWebviewState(
        "connecting",
        "Starting Codeflash Language Server...",
        null,
        null,
      );
    } else {
      this.updateWebviewState(
        "serverError",
        "Codeflash Language Server unavailable.",
        null,
        null,
      );
    }
  };

  private async handlePotentialAnalysis(uri: vscode.Uri | null): Promise<void> {
    if (!this._view) {
      return;
    }

    if (this._lspClient.state !== LanguageClientState.Running) {
      const message =
        this._lspClient.state === LanguageClientState.Starting
          ? "Connecting to server..."
          : "Server unavailable.";
      const status: SidebarStatus =
        this._lspClient.state === LanguageClientState.Starting
          ? "connecting"
          : "serverError";
      this.updateWebviewState(status, message, null, null);
      return;
    }

    if (!uri) {
      this.updateWebviewState(
        "noPythonFile",
        "Open a Python file to see optimizable functions.",
        null,
        null,
      );
      return;
    }

    await this.requestAnalysisFromServer(uri);
  }

  private async requestAnalysisFromServer(uri: vscode.Uri): Promise<void> {
    if (!this._view || this._lspClient.state !== LanguageClientState.Running) {
      this._logger.warn(
        `Analysis skipped for ${uri.fsPath}: View not ready or LSP not running. LSP State: ${this._lspClient.state}`,
      );
      if (this._view?.visible) {
        const message =
          this._lspClient.state === LanguageClientState.Starting
            ? "Language Server is starting up. Please wait..."
            : "Cannot analyze: Language Server not running.";
        this.updateWebviewState("serverError", message, null, null);
      }
      return;
    }

    this._logger.debug(`Requesting analysis for: ${uri.fsPath}`);

    try {
      const result = await this._analysisService.getOptimizableFunctions(uri);
      this._logger.debug(
        `Analysis result for ${uri.fsPath}: Status=${result.status}, Functions=${result.functions.length}, Message=${result.message}`,
      );
      this.requestFunctionsInFile(uri.fsPath, result.functions);

      const stillFocused = this._currentFocusUri?.toString() === uri.toString();
      if (this._view.visible && stillFocused) {
        this.updateWebviewState(
          result.status,
          result.message,
          result.functions.length > 0 ? result.functions : null,
          uri.fsPath,
        );
      } else if (this._view.visible && !stillFocused) {
        this._logger.debug(
          `Analysis complete for ${uri.fsPath}, but focus changed. Discarding results.`,
        );
      }
    } catch (error) {
      this._logger.error(
        `Failed to get analysis for ${uri.fsPath}`,
        error instanceof Error ? error : undefined,
      );
      const errorMessage =
        error instanceof Error ? error.message : String(error);
      this.updateWebviewState(
        "error",
        `Analysis failed: ${errorMessage}`,
        null,
        null,
      );
    }
  }

  private getPythonFileUri(
    editor: vscode.TextEditor | undefined,
  ): vscode.Uri | null {
    return editor?.document.languageId === "python" &&
      editor.document.uri.scheme === "file"
      ? editor.document.uri
      : null;
  }

  sendMessage(message: OutgoingWebviewMessage): void {
    if (!this._view) {
      return;
    }

    this._view.webview.postMessage(message).then(
      (success) => {
        if (!success) {
          this._logger.warn("Webview postMessage failed to send.");
        }
      },
      (error) => {
        this._logger.error(`Webview postMessage error: ${error}`);
      },
    );
  }

  private updateWebviewState(
    status: SidebarStatus,
    message: string,
    functionData: string[] | null,
    filePath: string | null,
  ): void {
    if (!this._view) {
      this._logger.warn(
        "Attempted to update webview state, but view is not available.",
      );
      return;
    }

    this._logger.debug(
      `Updating webview: Status=${status}, File=${filePath ?? "None"}`,
    );

    // Update function count and badge
    const newFunctionCount = functionData?.length || 0;
    if (this._currentFunctionCount !== newFunctionCount) {
      this._currentFunctionCount = newFunctionCount;
    }

    const updateMessage: UpdateStateMessage = {
      type: "updateState",
      payload: {
        status,
        message,
        functionData,
        currentFilePath: filePath || this._currentFocusUri?.fsPath || null,
      },
    };

    this.sendMessage(updateMessage);
  }

  public async refreshAnalysis(): Promise<void> {
    this._logger.debug("Manual analysis refresh requested");
    if (this._view?.visible) {
      await this.handleInitialState();
    }
  }

  public async handleOptimizeAllInCurrentFile(): Promise<void> {
    const activeEditor = vscode.window.activeTextEditor;
    if (!activeEditor || activeEditor.document.languageId !== "python") {
      vscode.window.showWarningMessage(
        "Please open a Python file to optimize all functions",
      );
      return;
    }

    try {
      const analysisResult =
        await this._analysisService.getOptimizableFunctions(
          activeEditor.document.uri,
        );

      if (
        analysisResult.status !== "success" ||
        !analysisResult.functions ||
        analysisResult.functions.length === 0
      ) {
        vscode.window.showInformationMessage(
          "No optimizable functions found in the current file",
        );
        return;
      }

      const functionCount = analysisResult.functions.length;
      const action = await vscode.window.showInformationMessage(
        `Found ${functionCount} optimizable function${functionCount > 1 ? "s" : ""}. Optimize all?`,
        "Yes",
        "No",
      );

      if (action !== "Yes") {
        return;
      }

      for (const functionName of analysisResult.functions) {
        const added = this.addFunctionToQueue(
          functionName,
          activeEditor.document.uri,
        );
        if (!added) {
          this._logger.warn(
            `Failed to add function '${functionName}' to queue`,
          );
        }
      }
    } catch (error) {
      this._logger.error(
        "Error optimizing all functions in current file",
        error instanceof Error ? error : undefined,
      );
      vscode.window.showErrorMessage(
        `Failed to optimize all functions: ${
          error instanceof Error ? error.message : "Unknown error"
        }`,
      );
    }
  }

  override dispose(): void {
    this._view = undefined;
    super.dispose();
  }
}
