geo/apps/browser-extension/src/storage.ts

import type { MonitoringLeaseTask } from "@geo/shared-types";
import { browser } from "wxt/browser";

import {
  createDefaultMonitoringPlatformStates,
  supportedMonitoringPlatforms,
  type StoredMonitoringPlatformState,
} from "./monitoring-platforms";
import { createDefaultPlatformStates, supportedPlatforms, type StoredPlatformState } from "./platforms";

export interface DoubaoMonitoringRuntimeState {
  fp: string | null;
  ms_token: string | null;
  device_id: string | null;
  web_id: string | null;
  tea_uuid: string | null;
  synced_at: string | null;
  source_url: string | null;
}

export interface MonitoringPlatformRuntimeState {
  doubao: DoubaoMonitoringRuntimeState | null;
}

export interface MonitoringPlatformRuntimeStatePatch {
  doubao?: Partial<DoubaoMonitoringRuntimeState> | null;
}

export interface StoredMonitoringScheduledTask extends MonitoringLeaseTask {
  pacing_mode: "immediate" | "delayed";
  scheduled_at: string;
  next_execute_at: string;
}

export interface ExtensionStorageState {
  installation_key: string;
  plugin_installation_id: number | null;
  installation_token: string | null;
  api_base_url: string | null;
  platforms: StoredPlatformState[];
  monitoring_platforms: StoredMonitoringPlatformState[];
  monitoring_last_heartbeat_at: string | null;
  monitoring_last_heartbeat_status: string | null;
  monitoring_last_cycle_at: string | null;
  monitoring_last_cycle_status: string | null;
  monitoring_task_queue: StoredMonitoringScheduledTask[];
  monitoring_runtime: MonitoringPlatformRuntimeState;
  last_updated_at: string | null;
}

const STORAGE_KEY = "geo.publisher.state";

function nextInstallationKey(): string {
  if (typeof crypto !== "undefined" && "randomUUID" in crypto) {
    return crypto.randomUUID();
  }
  return `geo-extension-${Date.now()}-${Math.random().toString(36).slice(2)}`;
}

function mergePlatformStates(platforms?: StoredPlatformState[]): StoredPlatformState[] {
  const defaults = createDefaultPlatformStates();
  if (!platforms?.length) {
    return defaults;
  }

  const incoming = new Map(platforms.map((platform) => [platform.platform_id, platform]));
  return defaults.map((platform) => {
    const stored = incoming.get(platform.platform_id);
    return stored ? { ...platform, ...stored } : platform;
  });
}

function mergeMonitoringPlatformStates(platforms?: StoredMonitoringPlatformState[]): StoredMonitoringPlatformState[] {
  const defaults = createDefaultMonitoringPlatformStates();
  if (!platforms?.length) {
    return defaults;
  }

  const incoming = new Map(platforms.map((platform) => [platform.ai_platform_id, platform]));
  return defaults.map((platform) => {
    const stored = incoming.get(platform.ai_platform_id);
    return stored ? { ...platform, ...stored } : platform;
  });
}

function mergeMonitoringScheduledTasks(tasks?: StoredMonitoringScheduledTask[]): StoredMonitoringScheduledTask[] {
  if (!tasks?.length) {
    return [];
  }

  return tasks
    .filter((task) => typeof task?.task_id === "number" && task.task_id > 0)
    .map((task) => ({
      ...task,
      pacing_mode: task.pacing_mode === "immediate" ? "immediate" : "delayed",
      scheduled_at: task.scheduled_at ?? new Date().toISOString(),
      next_execute_at: task.next_execute_at ?? new Date().toISOString(),
    }));
}

function normalizeDoubaoMonitoringRuntimeState(
  state?: Partial<DoubaoMonitoringRuntimeState> | null,
): DoubaoMonitoringRuntimeState | null {
  if (!state) {
    return null;
  }

  return {
    fp: state.fp ?? null,
    ms_token: state.ms_token ?? null,
    device_id: state.device_id ?? null,
    web_id: state.web_id ?? null,
    tea_uuid: state.tea_uuid ?? null,
    synced_at: state.synced_at ?? null,
    source_url: state.source_url ?? null,
  };
}

function mergeMonitoringRuntimeState(
  runtime?: MonitoringPlatformRuntimeStatePatch | null,
): MonitoringPlatformRuntimeState {
  return {
    doubao: normalizeDoubaoMonitoringRuntimeState(runtime?.doubao),
  };
}

function resolveNextMonitoringRuntimeState(
  current: MonitoringPlatformRuntimeState,
  patch?: MonitoringPlatformRuntimeStatePatch | null,
): MonitoringPlatformRuntimeState {
  if (patch == null) {
    return current;
  }

  const nextDoubao =
    typeof patch.doubao === "undefined"
      ? current.doubao
      : patch.doubao === null
        ? null
        : {
            ...(current.doubao ?? {}),
            ...patch.doubao,
          };

  return mergeMonitoringRuntimeState({
    ...current,
    ...patch,
    doubao: nextDoubao,
  });
}

async function writeState(nextState: ExtensionStorageState): Promise<void> {
  await browser.storage.local.set({
    [STORAGE_KEY]: nextState,
  });
}

export async function getExtensionState(): Promise<ExtensionStorageState> {
  const storageValue = await browser.storage.local.get(STORAGE_KEY);
  const stored = storageValue[STORAGE_KEY] as Partial<ExtensionStorageState> | undefined;

  const nextState: ExtensionStorageState = {
    installation_key: stored?.installation_key || nextInstallationKey(),
    plugin_installation_id: stored?.plugin_installation_id ?? null,
    installation_token: stored?.installation_token ?? null,
    api_base_url: stored?.api_base_url ?? null,
    platforms: mergePlatformStates(stored?.platforms),
    monitoring_platforms: mergeMonitoringPlatformStates(stored?.monitoring_platforms),
    monitoring_last_heartbeat_at: stored?.monitoring_last_heartbeat_at ?? null,
    monitoring_last_heartbeat_status: stored?.monitoring_last_heartbeat_status ?? null,
    monitoring_last_cycle_at: stored?.monitoring_last_cycle_at ?? null,
    monitoring_last_cycle_status: stored?.monitoring_last_cycle_status ?? null,
    monitoring_task_queue: mergeMonitoringScheduledTasks(stored?.monitoring_task_queue),
    monitoring_runtime: mergeMonitoringRuntimeState(stored?.monitoring_runtime),
    last_updated_at: stored?.last_updated_at ?? null,
  };

  const shouldPersist =
    !stored?.installation_key ||
    !stored?.platforms ||
    stored.platforms.length !== supportedPlatforms.length ||
    !stored?.monitoring_platforms ||
    stored.monitoring_platforms.length !== supportedMonitoringPlatforms.length ||
    typeof stored?.monitoring_task_queue === "undefined" ||
    !stored?.monitoring_runtime;

  if (shouldPersist) {
    await writeState(nextState);
  }

  return nextState;
}

export async function patchExtensionState(
  patch: Partial<Omit<ExtensionStorageState, "platforms" | "monitoring_platforms" | "monitoring_runtime">> & {
    platforms?: StoredPlatformState[];
    monitoring_platforms?: StoredMonitoringPlatformState[];
    monitoring_task_queue?: StoredMonitoringScheduledTask[];
    monitoring_runtime?: MonitoringPlatformRuntimeStatePatch | null;
  },
): Promise<ExtensionStorageState> {
  const current = await getExtensionState();
  const nextState: ExtensionStorageState = {
    ...current,
    ...patch,
    platforms: patch.platforms ? mergePlatformStates(patch.platforms) : current.platforms,
    monitoring_platforms: patch.monitoring_platforms
      ? mergeMonitoringPlatformStates(patch.monitoring_platforms)
      : current.monitoring_platforms,
    monitoring_task_queue:
      typeof patch.monitoring_task_queue !== "undefined"
        ? mergeMonitoringScheduledTasks(patch.monitoring_task_queue)
        : current.monitoring_task_queue,
    monitoring_runtime:
      typeof patch.monitoring_runtime !== "undefined"
        ? resolveNextMonitoringRuntimeState(current.monitoring_runtime, patch.monitoring_runtime)
        : current.monitoring_runtime,
    last_updated_at: new Date().toISOString(),
  };
  await writeState(nextState);
  return nextState;
}

export async function updatePlatformState(
  platformId: string,
  updater: Partial<StoredPlatformState>,
): Promise<ExtensionStorageState> {
  const current = await getExtensionState();
  const platforms = current.platforms.map((platform) =>
    platform.platform_id === platformId
      ? {
          ...platform,
          ...updater,
        }
      : platform,
  );
  return patchExtensionState({ platforms });
}

export async function resetPlatformStates(): Promise<ExtensionStorageState> {
  return patchExtensionState({ platforms: createDefaultPlatformStates() });
}

export async function updateMonitoringPlatformState(
  platformId: string,
  updater: Partial<StoredMonitoringPlatformState>,
): Promise<ExtensionStorageState> {
  const current = await getExtensionState();
  const monitoringPlatforms = current.monitoring_platforms.map((platform) =>
    platform.ai_platform_id === platformId
      ? {
          ...platform,
          ...updater,
        }
      : platform,
  );
  return patchExtensionState({ monitoring_platforms: monitoringPlatforms });
}

export async function resetMonitoringPlatformStates(): Promise<ExtensionStorageState> {
  return patchExtensionState({ monitoring_platforms: createDefaultMonitoringPlatformStates() });
}