import * as vscode from "vscode";
import { Disposable } from "../utils/dispose";
import { GlobalStateKey, type GlobalState } from "../globalState";
import type { NavigationService } from "../services";
import { Logger } from "../utils";
import { diffLines } from "diff";
import type {
  LSPCodeMessage,
  QueueTaskItem,
  RejectOptimizationMessage,
  ViewPatchMessage,
} from "@codeflash/types";

export class CommentThreadProvider extends Disposable {
  private controller: vscode.CommentController;
  private refreshTimer?: NodeJS.Timeout;
  // a map between task id and comment thread
  private threads: Map<
    string,
    {
      thread: vscode.CommentThread;
      line: number;
    }
  > = new Map();
  private decorationTypes: Map<string, vscode.TextEditorDecorationType> =
    new Map();
  private readonly logger = new Logger("CommentThreadProvider");
  private refreshing: boolean = false;
  constructor(
    private readonly context: vscode.ExtensionContext,
    private readonly globalState: GlobalState,
    private readonly navigationService: NavigationService,
  ) {
    super();
    const controller = vscode.comments.createCommentController(
      "codeflash-comments",
      "Codeflash suggestions",
    );
    controller.commentingRangeProvider = {
      provideCommentingRanges: async (document) => {
        if (document.languageId !== "python") {
          return [];
        }
        // vscode has a limitation where in order to expand the comment thread, the line should be in the commenting ranges,
        // this is for vscode only, other editors like cursor do not have this.
        const successfulOptimizations =
          this.getSuccessfulOptimizations().filter(
            (t) => t.filepath === document.uri.fsPath,
          );
        if (successfulOptimizations.length === 0) {
          return [];
        }

        const ranges: vscode.Range[] = [];

        for (const task of successfulOptimizations) {
          const filePath = task.filepath!;
          if (filePath !== document.uri.fsPath) {
            continue;
          }

          const pos =
            await this.navigationService.manuallyFindTheFunctionInFile(
              task.functionName,
              document.uri,
            );
          if (pos) {
            ranges.push(new vscode.Range(pos, pos));
          }
        }
        return ranges;
      },
    };

    this._disposables.push(controller);
    this._disposables.push(
      vscode.workspace.onDidChangeTextDocument((e) => {
        if (e.document.languageId === "python") {
          this.scheduleRefresh();
        }
      }),
      vscode.window.onDidChangeActiveTextEditor((editor) => {
        if (editor && editor.document.languageId === "python") {
          this.scheduleRefresh();
        }
      }),
      vscode.window.onDidChangeTextEditorVisibleRanges((e) => {
        if (e.textEditor.document.languageId === "python") {
          this.scheduleRefresh();
        }
      }),
    );
    this.controller = controller;
  }

  private getEditorForFileUri(
    fileUri: vscode.Uri | string,
  ): vscode.TextEditor | undefined {
    const visibleEditors = vscode.window.visibleTextEditors;
    return visibleEditors.find(
      (e) =>
        e.document.uri.fsPath ===
        (fileUri instanceof vscode.Uri
          ? fileUri.fsPath
          : vscode.Uri.file(fileUri).fsPath),
    );
  }

  public scheduleRefresh(delay: number = 500): void {
    clearTimeout(this.refreshTimer);
    this.refreshTimer = setTimeout(async () => {
      await this.refresh();
    }, delay);
  }

  private async refresh(): Promise<void> {
    if (this.refreshing) {
      return;
    }

    if (!this.controller) {
      throw new Error("Comment controller is not initialized");
    }

    this.logger.debug("Refreshing comment threads... ");
    this.refreshing = true;

    const removedTaskIds = this.detectRemovedTasks();
    for (const taskId of removedTaskIds) {
      this.cleanupThread(taskId);
      this.cleanUpDecoration(taskId);
    }

    const successfulOptimizations = this.getSuccessfulOptimizations();

    if (successfulOptimizations.length === 0) {
      this.logger.debug("No new tasks to add comment threads for.");
      this.refreshing = false;
      return;
    }

    for (const task of successfulOptimizations) {
      const fileUri = vscode.Uri.file(task.filepath!);
      const editor = this.getEditorForFileUri(fileUri);

      if (!editor) {
        this.logger.debug(`No visible editor for file ${task.filepath}`);
        continue;
      }

      const pos = await this.navigationService.manuallyFindTheFunctionInFile(
        task.functionName,
        fileUri,
      );
      if (!pos) {
        this.logger.debug(
          `Could not find position for function ${task.functionName} in file ${task.filepath}`,
        );
        if (this.threads.has(task.id)) {
          // cleanup if we can't find the function anymore
          this.cleanupThread(task.id);
          this.cleanUpDecoration(task.id);
        }
        continue;
      }

      this.drawDecorationForThread(editor, task.id, pos.line);

      const existingThread = this.threads.get(task.id);
      if (existingThread && existingThread.line === pos.line) {
        // no change in position, keep the existing thread
        continue;
      }

      this.cleanupThread(task.id);

      const thread = this.controller.createCommentThread(
        fileUri,
        new vscode.Range(pos, pos),
        [],
      );

      const acceptPatchPayload: ViewPatchMessage["payload"] = {
        id: task.id,
        functionName: task.functionName,
        patchFile: task.patchFile!,
        explanation: task.explanation || "",
        speedupStr: task.speedupStr || "",
        optimizationReview: task.optimizationReview || "",
        originalRuntime: task.originalRuntime || 0,
        optimizedRuntime: task.optimizedRuntime || 0,
        originalLineProfiler: task.originalLineProfiler || "",
        optimizedLineProfiler: task.optimizedLineProfiler || "",
      };
      const rejectPatchPayload: RejectOptimizationMessage["payload"] = {
        id: task.id,
      };

      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (thread as any).viewPatchPayload = acceptPatchPayload;
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      (thread as any).rejectPayload = rejectPatchPayload;

      // thread.state = vscode.CommentThreadState.Unresolved;
      thread.contextValue = "codeflash";
      thread.collapsibleState = vscode.CommentThreadCollapsibleState.Collapsed;
      thread.label = task.speedupStr
        ? `Optimization found - ${task.speedupStr}`
        : "Optimization found";
      thread.canReply = false;

      const botComment: vscode.Comment = {
        author: {
          name: "Codeflash",
          iconPath: vscode.Uri.file(
            this.context.asAbsolutePath("media/codeflash.svg"),
          ),
        },
        body: this.getMarkdownStringForTask(task),
        mode: vscode.CommentMode.Preview,
      };

      thread.comments = [botComment];

      this.threads.set(task.id, { thread, line: pos.line });
      this._disposables.push(thread);
    }
    this.refreshing = false;
  }

  private drawDecorationForThread(
    editor: vscode.TextEditor,
    taskId: string,
    line: number,
  ) {
    this.cleanUpDecoration(taskId);
    const pos = new vscode.Position(line, 0);
    const range = new vscode.Range(pos, pos);

    const decorationType = vscode.window.createTextEditorDecorationType({
      gutterIconPath: vscode.Uri.file(
        this.context.asAbsolutePath("media/codeflash.svg"),
      ),
      gutterIconSize: "contain",
    });

    editor.setDecorations(decorationType, [
      {
        range: range,
      },
    ]);
    this.decorationTypes.set(taskId, decorationType);
    this._disposables.push(decorationType);
  }

  private getSuccessfulOptimizations(): QueueTaskItem[] {
    return this.globalState
      .get(GlobalStateKey.QueueTasks, [])!
      .filter((t) => t.status === "completed");
  }

  private withNewLine(str: string): string {
    return str.endsWith("\n") ? str : str + "\n";
  }

  private fullUnifiedDiff(oldStr: string, newStr: string): string {
    const parts = diffLines(oldStr, newStr);

    const lines: string[] = [];

    for (const part of parts) {
      const prefix = part.added ? "+" : part.removed ? "-" : " ";
      const contentLines = part.value.split("\n");

      // Remove final empty line caused by trailing newline
      if (contentLines[contentLines.length - 1] === "") {
        contentLines.pop();
      }

      for (const line of contentLines) {
        lines.push(prefix + line);
      }
    }

    return lines.join("\n");
  }

  private getMarkdownStringForTask(task: QueueTaskItem): vscode.MarkdownString {
    const md = new vscode.MarkdownString();
    md.supportHtml = true;
    md.isTrusted = true;
    if (task.speedupStr) {
      md.appendMarkdown(`**Estimated Speedup:** ${task.speedupStr}<br/><br/>`);
    }
    if (task.bestCandidateCode) {
      // original code is always the first code block in the logs
      const originalCode = (
        task.logs.find((l) => l.type === "code") as LSPCodeMessage
      ).code;

      const patch = this.fullUnifiedDiff(
        this.withNewLine(originalCode),
        this.withNewLine(task.bestCandidateCode),
      );
      md.appendCodeblock(patch, "python");
    }
    if (task.explanation) {
      md.appendMarkdown(
        `<br/><details><br/><summary>**Explanation:**</summary><br/>${task.explanation}<br/></details><br/>`,
      );
    }
    return md;
  }

  private detectRemovedTasks(): string[] {
    const currentTaskIds = new Set(
      this.getSuccessfulOptimizations().map((t) => t.id),
    );
    const removedTaskIds: string[] = [];
    for (const taskId of this.threads.keys()) {
      if (!currentTaskIds.has(taskId)) {
        removedTaskIds.push(taskId);
      }
    }
    return removedTaskIds;
  }

  private cleanupThread(taskId: string) {
    this.logger.debug(`Cleaning up comment thread for task ${taskId}...`);
    const thread = this.threads.get(taskId);
    if (thread) {
      thread.thread.dispose();
      this.threads.delete(taskId);
    }
  }

  private cleanUpDecoration(taskId: string) {
    const decoration = this.decorationTypes.get(taskId);
    if (decoration) {
      decoration.dispose();
      this.decorationTypes.delete(taskId);
    }
  }
}
