import type { LogEntry, LSPCodeMessage } from "@codeflash/types";
import * as vscode from "vscode";
import { GlobalStateKey, type GlobalState } from "../globalState";
import { isTaskRunning } from "@codeflash/shared";
import { Logger } from "./logger";

export class LogsEventEmitter extends vscode.EventEmitter<LogEntry> {
  logger = new Logger("LogsEventEmitter");

  private currentTaskId: string | null = null;
  private currentBestCandidateCode: string | null = null;
  private logsBuffer: Map<string, LogEntry[]> = new Map();

  constructor(private readonly globalState: GlobalState) {
    super();
  }

  setCurrrentTaskId(taskId: string) {
    if (!!this.currentTaskId && this.currentTaskId !== taskId) {
      // current task changed, flush previous task logs
      this.flush();
    }
    this.logger.info(`Setting current task id to ${taskId}`);
    this.currentTaskId = taskId;
  }

  getCurrentTaskId() {
    return this.currentTaskId;
  }

  public addSingleBufferedLogEntry(logEntry: LogEntry): LogEntry[] | undefined {
    if (!this.currentTaskId) {
      return undefined;
    }
    const currentTaskLogs = this.logsBuffer.get(this.currentTaskId) || [];
    const newLogs = [...currentTaskLogs, logEntry];
    this.logsBuffer.set(this.currentTaskId, newLogs);
    return newLogs;
  }

  public flush(resetTaskId: boolean = true) {
    this.logger.info(`Flushing logs for task ${this.currentTaskId}`);
    if (!this.currentTaskId) {
      return;
    }
    const currentLogs = this.logsBuffer.get(this.currentTaskId) || [];

    if (currentLogs.length === 0) {
      return;
    }
    const currentTasks = this.globalState.get(GlobalStateKey.QueueTasks, [])!;

    const targetTaskIdx = currentTasks.findIndex(
      (task) => task.id === this.currentTaskId,
    );

    if (targetTaskIdx === -1 || !currentTasks[targetTaskIdx]) {
      this.logger.warn(
        `No task found with the id ${this.currentTaskId}, skipping flushing ${currentLogs.length} buffered log entries.`,
      );
      this.logsBuffer.delete(this.currentTaskId);
      this.logger.info(
        `Cleared buffered log entries for unknown task ${this.currentTaskId}. setting to null.`,
      );
      this.currentTaskId = null;
      return;
    }
    const targetTask = currentTasks[targetTaskIdx];
    targetTask.logs.push(...currentLogs);

    if (this.currentBestCandidateCode) {
      targetTask.bestCandidateCode = this.currentBestCandidateCode;
      this.currentBestCandidateCode = null;
    }

    currentTasks[targetTaskIdx] = targetTask;
    this.globalState.set(GlobalStateKey.QueueTasks, currentTasks);
    this.logger.debug(
      `Flushed ${currentLogs.length} buffered log entries to task ${targetTask.id}.`,
    );
    this.logsBuffer.delete(this.currentTaskId);
    // reset
    if (resetTaskId) {
      this.currentTaskId = null;
    }
  }

  private handleMessageId = (logEntry: LogEntry) => {
    const messageId = logEntry.message_id;
    if (!messageId) {
      return;
    }
    if (messageId === "best_candidate") {
      this.currentBestCandidateCode = (logEntry as LSPCodeMessage).code;
    }
  };
  private logMssageHandler = (message: LogEntry) => {
    const currentRunningTask = this.globalState
      .get(GlobalStateKey.QueueTasks, [])!
      .find((t) => t.id === this.currentTaskId);
    if (!currentRunningTask || !isTaskRunning(currentRunningTask)) {
      return;
    }
    if (message.task_id !== this.currentTaskId) {
      this.logger.warn(
        `Received log entry for task ${message.task_id} while current task is ${this.currentTaskId}, ignoring log entry.`,
      );
      return;
    }
    this.handleMessageId(message);
    this.addSingleBufferedLogEntry(message);
  };

  public registerLogsListener(): vscode.Disposable {
    const optimizationEventListener = this.event(
      this.logMssageHandler.bind(this),
    );
    return optimizationEventListener;
  }
}
