freezmo/lib/providers-client.ts

"use client";

import { useSyncExternalStore } from "react";

import { PROVIDER_BASE_HEADER, PROVIDER_HEADER, PROVIDER_KEY_HEADER, PROVIDERS } from "./providers";

/**
 * Per-provider API keys, kept in this browser only — like the Ollama key,
 * they are never stored on the server, never exported in a backup and never
 * part of the sync payload. They travel with each model call so a self-hosted
 * Freezmo can reach a cloud provider on the user's behalf.
 */

const STORAGE_KEY = "freezmo.providers.v1";

export type GenerationChoice = {
  provider: string;
  /** Model for decks, homework and grading. */
  textModel: string;
  /** Model for reading photos of notes (falls back to textModel). */
  visionModel: string;
};

export type ProvidersState = {
  /** provider id -> API key */
  keys: Record<string, string>;
  /** provider id -> last successfully listed models */
  models: Record<string, string[]>;
  /** Base URL for the "custom" OpenAI-compatible provider. */
  customBaseUrl: string;
  /** Provider used to generate decks and homework (null = server default). */
  generation: GenerationChoice | null;
};

const EMPTY: ProvidersState = { keys: {}, models: {}, customBaseUrl: "", generation: null };

let cached: ProvidersState = EMPTY;
let hydrated = false;
const listeners = new Set<() => void>();

function notify() {
  cached = { ...cached };
  listeners.forEach((listener) => listener());
}

function read(): ProvidersState {
  try {
    const raw = window.localStorage.getItem(STORAGE_KEY);
    if (!raw) return EMPTY;
    const parsed = JSON.parse(raw) as Partial<ProvidersState>;
    const keys: Record<string, string> = {};
    for (const [id, value] of Object.entries(parsed.keys ?? {})) {
      if (typeof value === "string" && value) keys[id] = value.slice(0, 512);
    }
    const models: Record<string, string[]> = {};
    for (const [id, list] of Object.entries(parsed.models ?? {})) {
      if (Array.isArray(list)) models[id] = list.filter((entry) => typeof entry === "string").slice(0, 400);
    }
    return {
      keys,
      models,
      customBaseUrl: typeof parsed.customBaseUrl === "string" ? parsed.customBaseUrl.slice(0, 300) : "",
      generation: generationOrNull(parsed.generation),
    };
  } catch {
    return EMPTY;
  }
}

function generationOrNull(value: unknown): GenerationChoice | null {
  if (!value || typeof value !== "object") return null;
  const record = value as Record<string, unknown>;
  if (typeof record.provider !== "string" || typeof record.textModel !== "string") return null;
  return {
    provider: record.provider,
    textModel: record.textModel,
    visionModel: typeof record.visionModel === "string" && record.visionModel ? record.visionModel : record.textModel,
  };
}

function persist(state: ProvidersState) {
  try {
    window.localStorage.setItem(STORAGE_KEY, JSON.stringify(state));
  } catch {
    /* private mode */
  }
}

export const providersStore = {
  subscribe: (listener: () => void) => {
    listeners.add(listener);
    return () => listeners.delete(listener);
  },
  getSnapshot: () => {
    if (!hydrated && typeof window !== "undefined") {
      cached = read();
      hydrated = true;
    }
    return cached;
  },
  getServerSnapshot: () => EMPTY,
};

function update(patch: Partial<ProvidersState>) {
  cached = { ...providersStore.getSnapshot(), ...patch };
  hydrated = true;
  persist(cached);
  notify();
}

export function useProviders(): ProvidersState {
  return useSyncExternalStore(providersStore.subscribe, providersStore.getSnapshot, providersStore.getServerSnapshot);
}

export function saveProviderKey(id: string, key: string) {
  update({ keys: { ...providersStore.getSnapshot().keys, [id]: key.trim() } });
}

export function clearProviderKey(id: string) {
  const current = providersStore.getSnapshot();
  const keys = { ...current.keys };
  delete keys[id];
  const models = { ...current.models };
  delete models[id];
  const generation = current.generation?.provider === id ? null : current.generation;
  update({ keys, models, generation });
}

export function setProviderModels(id: string, models: string[]) {
  update({ models: { ...providersStore.getSnapshot().models, [id]: models } });
}

export function setCustomBaseUrl(url: string) {
  update({ customBaseUrl: url.trim() });
}

export function setGeneration(choice: GenerationChoice | null) {
  update({ generation: choice });
}

/** Providers with a saved key, in registry order. */
export function enabledProviders(state: ProvidersState): string[] {
  return PROVIDERS.filter((provider) => Boolean(state.keys[provider.id])).map((provider) => provider.id);
}

/**
 * The headers that route a model call to a provider instead of the server's
 * Ollama. Empty for Ollama targets, so existing calls stay untouched.
 */
export function providerHeaders(provider: string | undefined): Record<string, string> {
  if (!provider || provider === "ollama") return {};
  const state = providersStore.getSnapshot();
  const key = state.keys[provider];
  if (!key) return {};
  const headers: Record<string, string> = { [PROVIDER_HEADER]: provider, [PROVIDER_KEY_HEADER]: key };
  if (provider === "custom" && state.customBaseUrl) headers[PROVIDER_BASE_HEADER] = state.customBaseUrl;
  return headers;
}

/** The models a provider offers: its live list, or the built-in seeds. */
export function modelsFor(state: ProvidersState, provider: string): string[] {
  const listed = state.models[provider];
  if (listed && listed.length > 0) return listed;
  return PROVIDERS.find((entry) => entry.id === provider)?.models ?? [];
}