import type { LanguageClient } from "vscode-languageclient/node";
import { CancellationTokenSource } from "vscode-languageclient/node";
import { State as LanguageClientState } from "vscode-languageclient/node";
import type { OptimizationResponse } from "../types";
import { LSP_COMMANDS } from "../constants";
import { Logger } from "../utils";
import PQueue from "p-queue";
import type { QueueTaskItem } from "@codeflash/types";

export interface OptimizationResult {
  success: boolean;
  error?: string;
  patch_file?: string;
  task_id?: string;
  explanation?: string;
  speedup?: string;
  canceled: boolean;
  optimizationReview?: string;
  original_runtime?: number;
  optimized_runtime?: number;
  original_line_profiler?: string;
  optimized_line_profiler?: string;
}

const DEMO_ID = "demo_optimization_task";

export class OptimizationService {
  private logger: Logger;
  private optimizationQueue: PQueue;
  private abortControllerMap = new Map<string, AbortController>();
  private cancelationTokenMap = new Map<string, CancellationTokenSource>();
  constructor(private optimizationClient: LanguageClient) {
    this.logger = new Logger("Codeflash Optimization");
    this.optimizationQueue = new PQueue({ concurrency: 1, autoStart: true });
  }

  async addOptimizationTask(
    id: string,
    functionName: string,
    options: {
      onTaskStart: (taskId: string) => Promise<boolean>;
      onTaskCompleted: (
        result: OptimizationResult,
        id: string,
      ) => Promise<void>;
      onTaskError: (error: Error, id: string) => Promise<void>;
      updateQueueTask: (payload: Partial<Omit<QueueTaskItem, "id">>) => void;
      abortController: AbortController;
      cancelationToken: CancellationTokenSource;
    },
  ): Promise<string> {
    this.abortControllerMap.set(id, options.abortController);
    this.cancelationTokenMap.set(id, options.cancelationToken);
    await this.optimizationQueue.add(
      async () => {
        // STEP 1: Initialize function optimization
        const initialized = await options.onTaskStart(id);
        if (!initialized) {
          return;
        }
        if (options.abortController.signal.aborted) {
          this.logger.info(`Optimization task ${id} was aborted before start.`);
          await options.onTaskCompleted({ success: false, canceled: true }, id);
          return;
        }
        options.updateQueueTask({ status: "optimizing" });

        try {
          const result = await this.optimizeFunction(
            id,
            functionName,
            options.updateQueueTask,
          );

          if (result.success) {
            await options.onTaskCompleted(result, id);
          } else {
            await options.onTaskError(new Error(result.error), id);
          }
        } catch (error) {
          await options.onTaskError(error as Error, id);
        }
      },
      { id, signal: options.abortController.signal },
    );
    return id;
  }

  abortTask(taskId: string): void {
    const abortController = this.abortControllerMap.get(taskId);
    if (!abortController) {
      return;
    }
    try {
      abortController.abort();
    } catch {}
    this.abortControllerMap.delete(taskId);
  }

  cancelTask(task: QueueTaskItem): void {
    const { id: taskId } = task;
    if (task.status === "initializing") {
      // task is cancelled quickly before it starts optimization, clear the session
      this.optimizationClient.sendRequest<Boolean>(
        LSP_COMMANDS.CLEANUP_CURRENT_OPTIMIZER_SESSION,
        {},
      );
    }
    const cancelationToken = this.cancelationTokenMap.get(taskId);
    if (!cancelationToken) {
      return;
    }
    cancelationToken.cancel();
    cancelationToken.dispose();
    this.cancelationTokenMap.delete(taskId);
  }

  private async optimizeFunction(
    taskId: string,
    functionName: string,
    updateTask: (payload: Partial<Omit<QueueTaskItem, "id">>) => void,
  ): Promise<OptimizationResult> {
    if (this.optimizationClient.state !== LanguageClientState.Running) {
      return {
        success: false,
        canceled: false,
        error: "Language Server not running.",
      };
    }

    this.logger.info(
      `Starting optimization process for function '${functionName}'`,
    );

    try {
      updateTask({ description: "Performing function optimization" });
      // Step 2: Perform function optimization
      const optimizationResult = await this.performFunctionOptimization(
        taskId,
        functionName,
      );

      // Check if optimization failed due to threshold not being met or other reasons
      if (optimizationResult.status === "error") {
        const errorMessage =
          optimizationResult.message || "Unknown optimization error";

        return {
          success: false,
          canceled: false,
          error: errorMessage,
        };
        // // Handle threshold errors gracefully
        // if (
        //   errorMessage.toLowerCase().includes("threshold") &&
        //   errorMessage.toLowerCase().includes("coverage")
        // ) {
        //   // updateTask({ description: "Optimization completed but no improvements found due to insufficient test coverage"})
        //   this.logger.info(
        //     `Optimization completed for ${functionName} but no improvements found due to insufficient test coverage`,
        //   );

        //   // Show user-friendly message
        //   vscode.window
        //     .showInformationMessage(
        //       `Codeflash analyzed '${functionName}' but couldn't optimize it due to insufficient test coverage. Consider adding more tests.`,
        //       "Learn More",
        //     )
        //     .then((selection) => {
        //       if (selection === "Learn More") {
        //         vscode.env.openExternal(
        //           vscode.Uri.parse("https://docs.codeflash.ai/test-coverage"),
        //         );
        //       }
        //     });

        //   return {
        //     success: false,
        //     canceled: false,
        //     error: errorMessage,
        //   };
        // } else {
        //   this.logger.error(
        //     `Optimization failed for ${functionName}: ${errorMessage}`,
        //   );
        //   return {
        //     success: false,
        //     canceled: false,
        //     error: errorMessage,
        //   };
        // }
      }

      const isCanceled = optimizationResult.status === "canceled";
      const optimizationReview =
        optimizationResult.optimizationReview as string;
      const msg = isCanceled
        ? `Optimization was canceled`
        : optimizationReview
          ? `Optimization completed successfully. \n 🎯Optimization Quality: ${optimizationReview}`
          : `Optimization completed successfully`;
      updateTask({ description: msg });
      return {
        success: true,
        canceled: isCanceled,
        patch_file: optimizationResult.patch_file as string,
        explanation: optimizationResult.explanation as string,
        task_id: optimizationResult.task_id as string,
        speedup: optimizationResult.extra as string,
        optimizationReview: optimizationResult.optimizationReview as string,
        original_runtime: optimizationResult.original_runtime as number,
        optimized_runtime: optimizationResult.optimized_runtime as number,
        original_line_profiler:
          optimizationResult.original_line_profiler as string,
        optimized_line_profiler:
          optimizationResult.optimized_line_profiler as string,
      };
    } catch (error) {
      this.logger.error(
        `Optimization failed for ${functionName}`,
        error instanceof Error ? error : undefined,
      );
      return {
        success: false,
        canceled: true,
        error: "Optimization failed for " + functionName,
      };
    }
  }

  async initializeFunctionOptimization(
    taskId: string,
    functionName: string,
    uriString: string,
  ): Promise<OptimizationResponse> {
    this.logger.info(`Initializing function optimization for ${functionName}`);

    const result =
      await this.optimizationClient.sendRequest<OptimizationResponse>(
        LSP_COMMANDS.INITIALIZE_FUNCTION_OPTIMIZATION,
        {
          textDocument: { uri: uriString },
          functionName: functionName,
          task_id: taskId,
        },
      );

    this.logger.info(`Initialization Response: ${JSON.stringify(result)}`);
    return result;
  }

  async performFunctionOptimization(
    taskId: string,
    functionName: string,
  ): Promise<OptimizationResponse> {
    this.logger.info(`Performing optimization for function ${functionName}`);
    const result =
      await this.optimizationClient.sendRequest<OptimizationResponse>(
        LSP_COMMANDS.PERFORM_FUNCTION_OPTIMIZATION,
        {
          textDocument: { uri: "" },
          functionName: functionName,
          task_id: taskId,
        },
        this.cancelationTokenMap.get(taskId)?.token,
      );

    this.logger.debug(`Optimization Result: ${JSON.stringify(result)}`);
    return result;
  }

  async runDemoOptimizationTask(
    functionName: string,
  ): Promise<OptimizationResponse> {
    const cancelationToken = new CancellationTokenSource();
    this.cancelationTokenMap.set(DEMO_ID, cancelationToken);

    return await this.optimizationClient.sendRequest<OptimizationResponse>(
      LSP_COMMANDS.START_DEMO_OPTIMIZATION,
      {
        functionName: functionName,
      },
      cancelationToken.token,
    );
  }

  cancelDemoTask(): void {
    const cancelationToken = this.cancelationTokenMap.get(DEMO_ID);
    if (!cancelationToken) {
      return;
    }
    cancelationToken.cancel();
    cancelationToken.dispose();
    this.cancelationTokenMap.delete(DEMO_ID);
  }

  dispose(): void {
    this.logger.dispose();
  }
}
