import { canTaskPersistInTheDisk } from "@codeflash/shared";
import type { QueueTaskItem } from "@codeflash/types";
import type { ExtensionContext } from "vscode";
import { window, workspace } from "vscode";

export const enum GlobalStateKey {
  IsExtensionInstalled = "isExtensionInstalled",
  TriedCodeflashDemo = "triedCodeflashDemo",
  Running = "running",
  UserID = "user_id",
  QueueTasks = "queueTasks",
  FocusedTaskId = "focusedTaskId",
  RootDir = "rootDir", // absolute path to the root of the workspace (e.g. where .git is located)
  ModuleRoot = "moduleRoot",
  ContextFiles = "contextFiles",
}

export type ContextFiles = Record<string, string>;

interface GlobalStateKeyMapping {
  [GlobalStateKey.Running]: boolean;
  [GlobalStateKey.IsExtensionInstalled]: boolean;
  [GlobalStateKey.TriedCodeflashDemo]: boolean;
  [GlobalStateKey.UserID]: string | undefined;
  [GlobalStateKey.QueueTasks]: QueueTaskItem[];
  [GlobalStateKey.FocusedTaskId]: string | undefined;
  [GlobalStateKey.RootDir]: string | undefined;
  [GlobalStateKey.ModuleRoot]: string | undefined;
  [GlobalStateKey.ContextFiles]: Record<string, ContextFiles> | undefined;
}

const workspaceSpecificKeys = [
  GlobalStateKey.RootDir,
  GlobalStateKey.Running,
  GlobalStateKey.ModuleRoot,
  GlobalStateKey.ContextFiles,
  GlobalStateKey.QueueTasks,
];

export class GlobalState {
  constructor(private context: ExtensionContext) {
    const localTasks = (this.get(GlobalStateKey.QueueTasks, []) ?? []).filter(
      canTaskPersistInTheDisk,
    );
    const cleanContextFiles = this.get(GlobalStateKey.ContextFiles, {})!;
    for (const taskId of Object.keys(cleanContextFiles)) {
      if (!localTasks.find((task) => task.id === taskId)) {
        delete cleanContextFiles[taskId];
      }
    }
    this.set(GlobalStateKey.ContextFiles, cleanContextFiles);
    this.set(GlobalStateKey.QueueTasks, localTasks);

    // reset state values
    this.set(GlobalStateKey.Running, false);
    this.set(GlobalStateKey.UserID, undefined); // TODO: should we preserve user id?
    this.set(GlobalStateKey.FocusedTaskId, undefined);
    this.set(GlobalStateKey.ModuleRoot, "");
    this.set(GlobalStateKey.RootDir, "");
  }

  get<T extends GlobalStateKeyMapping, E extends keyof T>(
    stateKey: E,
    defaultValue?: T[E],
  ): T[E] | undefined {
    if (workspaceSpecificKeys.includes(stateKey as GlobalStateKey)) {
      return (
        this.context.workspaceState.get(stateKey as string) ?? defaultValue
      );
    }

    return this.context.globalState.get(stateKey as string) ?? defaultValue;
  }

  set<T extends GlobalStateKeyMapping, E extends keyof T>(
    stateKey: E,
    newValue: T[E],
  ): void {
    if (workspaceSpecificKeys.includes(stateKey as GlobalStateKey)) {
      this.context.workspaceState.update(stateKey as string, newValue);
      return;
    }
    this.context.globalState.update(stateKey as string, newValue);
  }

  getStateForWebview(): Partial<{
    [key in keyof GlobalStateKeyMapping]: GlobalStateKeyMapping[key];
  }> {
    return {
      running: this.get(GlobalStateKey.Running, false)!,
      focusedTaskId: this.get(GlobalStateKey.FocusedTaskId)!,
      queueTasks: this.get(GlobalStateKey.QueueTasks, [])!,
      moduleRoot: this.get(GlobalStateKey.ModuleRoot)!,
      triedCodeflashDemo: this.get(GlobalStateKey.TriedCodeflashDemo, false)!,
    };
  }

  /**
   * Run while developing to see the entire global storage contents.
   */
  async showGlobalStateValue() {
    const all: Record<string, unknown> = {};
    for (const key of this.context.globalState.keys()) {
      all[key] = this.get(key as GlobalStateKey);
    }
    for (const key of this.context.workspaceState.keys()) {
      all[key] = this.get(key as GlobalStateKey);
    }

    const document = await workspace.openTextDocument({
      language: "jsonc",
      content: JSON.stringify(all, null, "  "),
    });
    window.showTextDocument(document);
  }

  /**
   * Dev function (clear all global state properties).
   */
  clearGlobalState() {
    for (const key of this.context.globalState.keys()) {
      this.context.globalState.update(key, undefined);
    }
  }
}
