geo/apps/desktop-client/src/main/session-registry.ts

import { randomUUID } from 'node:crypto'
import { mkdirSync, readFileSync, writeFileSync } from 'node:fs'
import { dirname, join } from 'node:path'

import type { Session } from 'electron/main'
import { app, session } from 'electron/main'

import { observeSessionRequests } from './network-observer'
import { STANDARD_ACCEPT_LANGUAGES, STANDARD_USER_AGENT } from './user-agent'

export interface SessionHandle {
  accountId: string
  partition: string
  session: Session
}

export interface SessionHandleSnapshot {
  accountId: string
  partition: string
}

const registry = new Map<string, SessionHandle>()
const sessionsWithUA = new WeakSet<Session>()
let persistedPartitionsCache: Record<string, string> | null = null

function persistedPartitionsPath(): string {
  return join(app.getPath('userData'), 'desktop-session-partitions.json')
}

function readPersistedPartitions(): Record<string, string> {
  if (persistedPartitionsCache) {
    return persistedPartitionsCache
  }

  try {
    persistedPartitionsCache = JSON.parse(
      readFileSync(persistedPartitionsPath(), 'utf8'),
    ) as Record<string, string>
  } catch {
    persistedPartitionsCache = {}
  }

  return persistedPartitionsCache
}

function writePersistedPartitions(next: Record<string, string>): void {
  persistedPartitionsCache = next
  const target = persistedPartitionsPath()
  mkdirSync(dirname(target), { recursive: true })
  writeFileSync(target, JSON.stringify(next, null, 2), 'utf8')
}

function persistedPartitionFor(accountId: string): string | null {
  return readPersistedPartitions()[accountId] ?? null
}

export function getPersistedPartition(accountId: string): string | null {
  return persistedPartitionFor(accountId)
}

function rememberPersistedPartition(accountId: string, partition: string): void {
  const current = readPersistedPartitions()
  if (current[accountId] === partition) {
    return
  }

  writePersistedPartitions({
    ...current,
    [accountId]: partition,
  })
}

function forgetPersistedPartition(accountId: string): string | null {
  const current = readPersistedPartitions()
  const partition = current[accountId] ?? null
  if (!partition) {
    return null
  }

  const next = { ...current }
  delete next[accountId]
  writePersistedPartitions(next)
  return partition
}

function partitionFor(accountId: string): string {
  return `persist:acc-${accountId}`
}

function pendingPartitionFor(seed: string): string {
  return `persist:pending-${seed}-${randomUUID()}`
}

function applyStandardUserAgent(target: Session): Session {
  if (!sessionsWithUA.has(target)) {
    target.setUserAgent(STANDARD_USER_AGENT, STANDARD_ACCEPT_LANGUAGES)
    sessionsWithUA.add(target)
  }
  return target
}

function prepareSession(target: Session, options: { label: string; partition: string }): Session {
  applyStandardUserAgent(target)
  observeSessionRequests(target, {
    label: options.label,
    partition: options.partition,
  })
  target.clearHostResolverCache().catch((error) => {
    console.warn('[desktop-session] clearHostResolverCache failed', error)
  })
  return target
}

export async function initSessionRegistry(): Promise<void> {
  registry.clear()
}

export function createSessionHandle(accountId?: string): SessionHandle {
  const resolvedAccountID = accountId ?? randomUUID()
  const existing = registry.get(resolvedAccountID)
  if (existing) {
    return existing
  }

  const partition = persistedPartitionFor(resolvedAccountID) ?? partitionFor(resolvedAccountID)
  const handle: SessionHandle = {
    accountId: resolvedAccountID,
    partition,
    session: prepareSession(session.fromPartition(partition), {
      label: resolvedAccountID,
      partition,
    }),
  }
  registry.set(resolvedAccountID, handle)
  return handle
}

export function createSessionHandleForPartition(
  accountId: string,
  partition: string,
): SessionHandle {
  const existing = registry.get(accountId)
  if (existing && existing.partition === partition) {
    return existing
  }

  const handle: SessionHandle = {
    accountId,
    partition,
    session: prepareSession(session.fromPartition(partition), {
      label: accountId,
      partition,
    }),
  }
  registry.set(accountId, handle)
  rememberPersistedPartition(accountId, partition)
  return handle
}

export function createPendingSessionHandle(seed = 'bind'): SessionHandle {
  const accountId = `pending:${seed}:${randomUUID()}`
  const partition = pendingPartitionFor(seed)
  const handle: SessionHandle = {
    accountId,
    partition,
    session: prepareSession(session.fromPartition(partition), {
      label: accountId,
      partition,
    }),
  }
  registry.set(accountId, handle)
  return handle
}

export function attachSessionHandle(accountId: string, handle: SessionHandle): SessionHandle {
  const existing = registry.get(accountId)
  if (existing && existing !== handle) {
    registry.delete(existing.accountId)
  }

  if (handle.accountId !== accountId) {
    registry.delete(handle.accountId)
  }

  const next: SessionHandle = {
    accountId,
    partition: handle.partition,
    session: handle.session,
  }
  registry.set(accountId, next)
  rememberPersistedPartition(accountId, handle.partition)
  return next
}

export function forgetSessionHandle(accountId: string): void {
  registry.delete(accountId)
}

export async function clearSessionHandle(accountId: string): Promise<void> {
  const active = registry.get(accountId)
  const persistedPartition = forgetPersistedPartition(accountId)
  const partition = active?.partition ?? persistedPartition
  registry.delete(accountId)

  if (!partition) {
    return
  }

  const target = prepareSession(session.fromPartition(partition), {
    label: accountId,
    partition,
  })

  try {
    await target.clearStorageData()
  } catch (error) {
    console.warn('[desktop-session] clearStorageData failed', { accountId, partition, error })
  }

  try {
    await target.clearCache()
  } catch (error) {
    console.warn('[desktop-session] clearCache failed', { accountId, partition, error })
  }

  try {
    await target.cookies.flushStore()
  } catch (error) {
    console.warn('[desktop-session] flushStore failed', { accountId, partition, error })
  }
}

export function getSessionHandle(accountId: string): SessionHandle | undefined {
  return registry.get(accountId)
}

export function listSessionHandles(): SessionHandle[] {
  return [...registry.values()]
}

export function listSessionHandleSnapshots(): SessionHandleSnapshot[] {
  return [...registry.values()].map(({ accountId, partition }) => ({
    accountId,
    partition,
  }))
}