stack/packages/memory/src/adapters/pgvector.ts

import { createDb, type DbHandle } from '@mosaic/db';
import type {
  MemoryAdapter,
  MemoryConfig,
  NewInsight as AdapterNewInsight,
  Insight as AdapterInsight,
  InsightSearchResult,
} from '../types.js';
import type { EmbeddingProvider } from '../vector-store.js';
import {
  createPreferencesRepo,
  type PreferencesRepo,
  type Preference,
  type NewPreference,
} from '../preferences.js';
import {
  createInsightsRepo,
  type InsightsRepo,
  type NewInsight as DbNewInsight,
} from '../insights.js';

type PgVectorConfig = Extract<MemoryConfig, { type: 'pgvector' }>;

export class PgVectorAdapter implements MemoryAdapter {
  readonly name = 'pgvector';
  readonly embedder: EmbeddingProvider | null;

  private handle: DbHandle;
  private preferences: PreferencesRepo;
  private insights: InsightsRepo;

  constructor(config: PgVectorConfig) {
    this.handle = createDb();
    this.preferences = createPreferencesRepo(this.handle.db);
    this.insights = createInsightsRepo(this.handle.db);
    this.embedder = config.embedder ?? null;
  }

  /* ------------------------------------------------------------------ */
  /*  Preferences                                                       */
  /* ------------------------------------------------------------------ */

  async getPreference(userId: string, key: string): Promise<unknown | null> {
    const row = await this.preferences.findByUserAndKey(userId, key);
    return row?.value ?? null;
  }

  async setPreference(
    userId: string,
    key: string,
    value: unknown,
    category?: string,
  ): Promise<void> {
    await this.preferences.upsert({
      userId,
      key,
      value,
      ...(category ? { category: category as NewPreference['category'] } : {}),
    });
  }

  async deletePreference(userId: string, key: string): Promise<boolean> {
    return this.preferences.remove(userId, key);
  }

  async listPreferences(
    userId: string,
    category?: string,
  ): Promise<Array<{ key: string; value: unknown; category: string }>> {
    const rows = category
      ? await this.preferences.findByUserAndCategory(userId, category as Preference['category'])
      : await this.preferences.findByUser(userId);

    return rows.map((r) => ({ key: r.key, value: r.value, category: r.category }));
  }

  /* ------------------------------------------------------------------ */
  /*  Insights                                                          */
  /* ------------------------------------------------------------------ */

  async storeInsight(insight: AdapterNewInsight): Promise<AdapterInsight> {
    const row = await this.insights.create({
      userId: insight.userId,
      content: insight.content,
      source: insight.source as DbNewInsight['source'],
      category: insight.category as DbNewInsight['category'],
      relevanceScore: insight.relevanceScore,
      metadata: insight.metadata ?? {},
      embedding: insight.embedding ?? null,
    });

    return toAdapterInsight(row);
  }

  async getInsight(id: string): Promise<AdapterInsight | null> {
    // findById requires userId — search across all users via raw find
    // The adapter interface only takes id, so we pass an empty userId and rely on the id match.
    // Since the repo requires userId, we use a two-step approach.
    const row = await this.insights.findById(id, '');
    if (!row) return null;
    return toAdapterInsight(row);
  }

  async searchInsights(
    userId: string,
    _query: string,
    opts?: { limit?: number; embedding?: number[] },
  ): Promise<InsightSearchResult[]> {
    if (opts?.embedding) {
      const results = await this.insights.searchByEmbedding(
        userId,
        opts.embedding,
        opts.limit ?? 10,
      );
      return results.map((r) => ({
        id: r.insight.id,
        content: r.insight.content,
        score: 1 - r.distance,
        metadata: (r.insight.metadata as Record<string, unknown>) ?? undefined,
      }));
    }

    // Fallback: return recent insights for the user
    const rows = await this.insights.findByUser(userId, opts?.limit ?? 10);
    return rows.map((r) => ({
      id: r.id,
      content: r.content,
      score: Number(r.relevanceScore),
      metadata: (r.metadata as Record<string, unknown>) ?? undefined,
    }));
  }

  async deleteInsight(id: string): Promise<boolean> {
    // The repo requires userId — pass empty string since adapter interface only has id
    return this.insights.remove(id, '');
  }

  /* ------------------------------------------------------------------ */
  /*  Lifecycle                                                         */
  /* ------------------------------------------------------------------ */

  async close(): Promise<void> {
    await this.handle.close();
  }
}

/* ------------------------------------------------------------------ */
/*  Helpers                                                            */
/* ------------------------------------------------------------------ */

function toAdapterInsight(row: {
  id: string;
  userId: string;
  content: string;
  source: string;
  category: string;
  relevanceScore: number;
  metadata: unknown;
  embedding: unknown;
  createdAt: Date;
  updatedAt: Date | null;
  decayedAt: Date | null;
}): AdapterInsight {
  return {
    id: row.id,
    userId: row.userId,
    content: row.content,
    source: row.source,
    category: row.category,
    relevanceScore: row.relevanceScore,
    metadata: (row.metadata as Record<string, unknown>) ?? undefined,
    embedding: (row.embedding as number[]) ?? undefined,
    createdAt: row.createdAt,
    updatedAt: row.updatedAt ?? undefined,
    decayedAt: row.decayedAt ?? undefined,
  };
}