import { CheckEnvironmentStep } from "../boot/checkEnvironmentStep";
import { Logger } from "../utils";
import { Disposable } from "../utils/dispose";
import * as vscode from "vscode";
import type { StepInfo } from "../providers/InitWebviewProvider";
import type {
  BootCodeflashServerStepResult,
  CheckEnvironmentStepResult,
  CodeflashInitResult,
  StepResult,
} from "../boot/baseStep";
import type BaseStep from "../boot/baseStep";
import { InitStepError, STEP_DEFINITIONS } from "../boot/baseStep";
import { BootCodeflashServerStep } from "../boot/bootCodeflashServer";
import { type GlobalState } from "../globalState";
import { CodeflashInitStep } from "../boot/codeflashInitStep";
import type { LogsEventEmitter } from "../utils/logsEventEmitter";
import { captureException } from "../telemetry/sentry";
import {
  ERROR_CATEGORIES,
  getCodeflashVersionIncompatibleError,
} from "../boot/stepErrors";
import type { Suggestion } from "@codeflash/types";
import { EventTracker } from "../telemetry/posthog";
import type { LanguageClient } from "vscode-languageclient/node";
import { LSP_COMMANDS } from "../constants";
import { ConfigKey, type ConfigService } from "./configService";
import { POSTHOG_EVENTS } from "../telemetry/events";

type CompleteInitResult = {
  environmentResult?: CheckEnvironmentStepResult;
  bootResult?: BootCodeflashServerStepResult;
  initResult?: CodeflashInitResult;
};
export class InitService extends Disposable {
  logger: Logger = new Logger("InitService");
  completeInitResult: CompleteInitResult = {};
  private isRecordingSession: boolean | undefined;

  updateInitMessage?: (message: string) => void;
  updateStep?: (number: number, update: Partial<StepInfo>) => void;
  setSteps?: (setter: (current: StepInfo[]) => StepInfo[]) => void;
  getCurrentSteps?: () => StepInfo[];
  renderSteps?: () => void;
  showLoadingView?: () => void;
  showInitForm?: (
    suggestions: Record<string, Suggestion>,
    error?: string,
    pyprojectPath?: string,
  ) => void;
  setSessionRecording?: (staus: "start" | "stop") => void;
  updateStepProgress?: (currentStep: number, totalSteps: number) => void;

  currentRenderedError?: string;

  private pausePooling: boolean = false;

  constructor(
    private readonly context: vscode.ExtensionContext,
    private readonly globalState: GlobalState,
    private readonly configService: ConfigService,
    private readonly optimizationEventEmitter: LogsEventEmitter,
  ) {
    super();
    this._disposables.push(this.logger);
    this.logger.info("InitService initialized");
    this.isRecordingSession = false;
  }

  public resumePooling() {
    this.pausePooling = false;
  }

  public setPausePooling(pause: boolean) {
    this.pausePooling = pause;
  }

  public setUpdateInitMessage(updateInitMessage: (message: string) => void) {
    this.updateInitMessage = updateInitMessage;
  }

  public setRenderSteps(renderSteps: () => void) {
    this.renderSteps = renderSteps;
  }

  public setSetSessionRecording(
    sessionRecording: (status: "start" | "stop") => void,
  ) {
    this.setSessionRecording = sessionRecording;
  }

  public safeguardSetSessionRecording(status: "start" | "stop") {
    if (!this.setSessionRecording) {
      return;
    }
    if (status === "start" && this.isRecordingSession) {
      return;
    }
    if (status === "stop" && !this.isRecordingSession) {
      return;
    }
    this.isRecordingSession = status === "start";
    this.setSessionRecording(status);
  }

  public setShowInitForm(
    showInitForm: (
      suggestions: Record<string, Suggestion>,
      error?: string,
      pyprojectPath?: string,
    ) => void,
  ) {
    this.showInitForm = showInitForm;
  }
  public setShowLoadingView(showLoadingView: () => void) {
    this.showLoadingView = showLoadingView;
  }
  public setUpdateStep(
    updateStep: (number: number, update: Partial<StepInfo>) => void,
  ) {
    this.updateStep = (number: number, update: Partial<StepInfo>) => {
      if (update && update.status === "failed") {
        this.safeguardSetSessionRecording("start"); // start recording on any step failure
      }
      updateStep(number, update);
    };
  }
  public setGetCurrentSteps(getCurrentSteps: () => StepInfo[]) {
    this.getCurrentSteps = getCurrentSteps;
  }
  public setSetSteps(
    _setSteps: (setter: (current: StepInfo[]) => StepInfo[]) => void,
  ) {
    this.setSteps = _setSteps;
  }

  public setUpdateStepProgress(
    updateStepProgress: (currentStep: number, totalSteps: number) => void,
  ) {
    this.updateStepProgress = updateStepProgress;
  }

  public async initWithRetries(
    maxRetries: number = 3,
  ): Promise<CompleteInitResult | null> {
    try {
      const result = await this.init();
      if (result) {
        vscode.window.showInformationMessage(
          "Codeflash initialized successfully",
        );
      }
      return result;
    } catch (error) {
      if (maxRetries > 0) {
        this.logger.error(
          `InitService init error: ${error} (${maxRetries} retries left)`,
        );

        // Dispose all resources including LSP clients and command registrations
        this.dispose();
        this._disposables = [];

        // Wait longer (1 second) to ensure VS Code fully unregisters commands and disposes resources
        // This prevents "command already exists" errors during retry
        // The delay accounts for:
        // 1. Async disposal of LSP clients
        // 2. VS Code's internal command registry cleanup
        // 3. Event listener teardown
        this.logger.info("Waiting for cleanup to complete before retry...");
        await new Promise((resolve) => setTimeout(resolve, 1000));

        return await this.initWithRetries(maxRetries - 1);
      } else {
        this.logger.error(`InitService init failed: ${error}`);
        vscode.window.showErrorMessage(
          `Codeflash initialization failed: ${error}`,
        );
        throw error;
      }
    }
  }

  private uiMethodsIntialized(): boolean {
    if (
      !this.setSteps ||
      !this.updateInitMessage ||
      !this.updateStep ||
      !this.renderSteps ||
      !this.showLoadingView ||
      !this.getCurrentSteps ||
      !this.showInitForm ||
      !this.setSessionRecording ||
      !this.updateStepProgress
    ) {
      return false;
    }
    return true;
  }

  private async runWithoutPooling(
    step: BaseStep,
    stepNum: number,
  ): Promise<StepResult | InitStepError | null> {
    if (this.isDisposed) {
      return null;
    }
    const result = await step.run();
    if (result instanceof InitStepError) {
      this.configService.listenForConfigChanges();
      captureException(result, {
        ...result,
      });
      this.updateStep!(stepNum, {
        status: "failed",
        error: result,
      });
      this.renderSteps!();
    }
    return result;
  }

  private async runWithPooling(
    step: BaseStep,
    stepsSnapshot: StepInfo[],
    stepNum: number,
    failed: boolean = false,
  ): Promise<StepResult | null> {
    if (this.isDisposed) {
      // Very important, break the loop if the service is disposed.
      return null;
    }

    if (this.pausePooling) {
      await new Promise((resolve) => setTimeout(resolve, 600));
      return this.runWithPooling(step, stepsSnapshot, stepNum, failed);
    }
    const result = await step.run();

    if (result instanceof InitStepError) {
      // maybe the user will try to modify the config, so we need to listen for config changes again.
      this.configService.listenForConfigChanges();
      // failed
      this.updateStep!(stepNum, {
        status: "failed",
        error: result,
      });

      const isInitError = result.category === ERROR_CATEGORIES.CODEFLASH_INIT;
      const stepErrorIdentifier = `${stepNum}-${result.message}-${result.details}`;
      const isDifferentError =
        this.currentRenderedError !== stepErrorIdentifier;
      if (
        isInitError &&
        isDifferentError &&
        step instanceof CodeflashInitStep
      ) {
        // no need to continue pooling if we hit an init error, as the user needs to fix the config first, then we can resume pooling.
        this.pausePooling = true;
        // show the init form for init related errors on the first failure.
        this.showInitForm!(
          step.configSuggestions || {},
          result.details,
          step.pyprojectPath,
        );
      } else if (isDifferentError && !isInitError) {
        // for other errors, we still use the same view, but we only re-render if it's a different error than the last rendered one.
        // this is to avoid error flickering in the UI during pooling retries.
        this.renderSteps!();
      }

      if (isDifferentError) {
        await EventTracker.capture(POSTHOG_EVENTS.INIT_ERROR, {
          stepNumber: stepNum,
          ...result,
        });
      }

      this.currentRenderedError = stepErrorIdentifier;

      step.dispose();
      // retry after 0.5 second
      await new Promise((resolve) => setTimeout(resolve, 500));
      return this.runWithPooling(step, stepsSnapshot, stepNum, true);
    }

    // success
    this.setSteps!(() => stepsSnapshot);
    // success after failing
    if (failed) {
      this.showLoadingView!();
    }
    this.updateStep!(stepNum, { status: "completed" });
    return result;
  }

  private async checkFeaturesCompatibility(
    codeflashVersion: string,
    lspClient: LanguageClient,
    interpreterPath: string,
  ): Promise<InitStepError | undefined> {
    try {
      const lspFeatures = await lspClient.sendRequest<string[]>(
        "server/listFeatures",
        {},
      );
      this.logger.info(`LSP features: ${lspFeatures.join(", ")}`);
      const expectedFeatures = Object.values(LSP_COMMANDS);
      const missingFeatures = expectedFeatures.filter(
        (feature) => !lspFeatures.includes(feature),
      );
      if (missingFeatures.length > 0) {
        return getCodeflashVersionIncompatibleError(
          codeflashVersion,
          interpreterPath,
          false,
        );
      }
      return undefined;
    } catch {
      return getCodeflashVersionIncompatibleError(
        codeflashVersion,
        interpreterPath,
        false,
      );
    }
  }

  private async init(): Promise<CompleteInitResult | null> {
    if (!this.uiMethodsIntialized()) {
      throw new Error("InitService not initialized yet");
    }

    const totalSteps = STEP_DEFINITIONS.length;

    this.setSteps!(() => []);
    this.setSteps!(() =>
      STEP_DEFINITIONS.map((def) => ({
        status: "idle" as const,
        title: def.title,
        description: "",
      })),
    );

    const configuredPythonPath = this.configService.get(
      ConfigKey.DefaultInterpreterPath,
    ) as string | null;
    this.logger.debug(`Configured Python Path: ${configuredPythonPath}`);
    const environmentStep = new CheckEnvironmentStep(this.configService);
    this._disposables.push(
      environmentStep.pythonService!.setupInterpreterWatcher(),
    );
    this._disposables.push(environmentStep);

    // Step 1: Check environment
    this.updateStepProgress!(1, totalSteps);
    this.updateInitMessage!(environmentStep.stepTitle);
    const environmentResult = (await this.runWithPooling(
      environmentStep,
      this.getCurrentSteps!(),
      1,
    )) as CheckEnvironmentStepResult;
    if (!environmentResult) {
      return null;
    }
    this.logger.info(
      `using python interpreter: ${environmentResult.pythonPath}`,
    );
    this.completeInitResult.environmentResult = environmentResult;

    // Step 2: Start server and initialize services
    const bootServerStep = new BootCodeflashServerStep(
      this.context,
      this.globalState,
      this.configService,
      environmentResult.pythonPath,
      this.optimizationEventEmitter,
    );
    this._disposables.push(bootServerStep);
    this.updateStepProgress!(2, totalSteps);
    this.updateInitMessage!(bootServerStep.stepTitle);
    const bootResult = (await this.runWithoutPooling(bootServerStep, 2)) as
      | BootCodeflashServerStepResult
      | InitStepError;
    if (bootResult instanceof InitStepError) {
      return null;
    }
    const featureCompatibilityError = await this.checkFeaturesCompatibility(
      environmentResult.codeflashVersion,
      bootResult.lspService.getClient(),
      environmentResult.pythonPath,
    );
    if (featureCompatibilityError) {
      this.updateStep!(2, {
        status: "failed",
        error: featureCompatibilityError,
      });
      this.renderSteps!();
      return null;
    }

    this.updateStep!(2, { status: "completed" });
    this.completeInitResult.bootResult = bootResult;

    // Step 3: codeflash validation and initialization completed
    const lspClient = bootResult.lspService.getClient();
    const optimizationClient = bootResult.optimizationLspService.getClient();
    const finalizeStep = new CodeflashInitStep(
      lspClient,
      optimizationClient,
      this.globalState,
      this.configService,
    );

    this._disposables.push(finalizeStep);
    this.updateStepProgress!(3, totalSteps);
    this.updateInitMessage!(finalizeStep.stepTitle);
    const finalizeResult = (await this.runWithPooling(
      finalizeStep,
      this.getCurrentSteps!(),
      3,
    )) as CodeflashInitResult;
    if (!finalizeResult) {
      return null;
    }
    this.completeInitResult.initResult = finalizeResult;
    this.updateStepProgress!(0, totalSteps);
    this.safeguardSetSessionRecording("stop");
    bootResult.codeLensProvider.scheduleRefresh();
    bootResult.commentThreadProvider.scheduleRefresh();
    this.configService.listenForConfigChanges();
    bootResult.sidebarProvider.projectConfigs = {
      config: finalizeStep.configSuggestions,
      pyprojectPath: finalizeResult.pyprojectPath!,
    };
    environmentStep.pythonService!.setReloadOnInterpreterChange(true);
    await EventTracker.capture(POSTHOG_EVENTS.INIT_COMPLETED, {});
    return this.completeInitResult;
  }
  override dispose(): void {
    super.dispose();
  }
}
