store.ts

import { readFile, readdir, writeFile, mkdir, rm, stat } from "node:fs/promises";
import { join, basename } from "node:path";
import { existsSync } from "node:fs";
import { createHash } from "node:crypto";
import * as lancedb from "@lancedb/lancedb";
import type { VectorQuery } from "@lancedb/lancedb";
import { EmbeddingEngine } from "./embeddings.js";
import {
  type MemoryEntry,
  type MemoryCatalogEntry,
  type MemoryCategory,
  type MemoryStatus,
  type MemoryFrontmatter,
  VALID_CATEGORIES,
  MemoryMCPError,
} from "./types.js";

function vectorRecord(entry: MemoryEntry, vector: number[], contentHash: string): Record<string, unknown> {
  return {
    id: entry.id,
    title: entry.title,
    category: entry.category,
    date: entry.date,
    tags: entry.tags.join(","),
    status: entry.status,
    snippet: entry.content.length > 200 ? entry.content.slice(0, 200) + "..." : entry.content,
    content_hash: contentHash,
    vector,
  };
}

export class MemoryStore {
  private memoriesPath: string;
  private catalog: MemoryEntry[] = [];
  private embeddings: EmbeddingEngine;
  private db: lancedb.Connection | null = null;
  private table: lancedb.Table | null = null;
  private loaded = false;
  private loadingPromise: Promise<void> | null = null;

  constructor(memoriesPath: string, embeddingModel?: string) {
    this.memoriesPath = memoriesPath;
    this.embeddings = new EmbeddingEngine(embeddingModel);
  }

  async load(): Promise<void> {
    this.loadingPromise = this.doLoad();
    return this.loadingPromise;
  }

  private async doLoad(): Promise<void> {
    this.catalog = [];

    await this.embeddings.init();

    if (!existsSync(this.memoriesPath)) {
      await mkdir(this.memoriesPath, { recursive: true });
    }

    for (const category of VALID_CATEGORIES) {
      const catDir = join(this.memoriesPath, category);
      if (!existsSync(catDir)) continue;

      const files = await readdir(catDir);
      const mdFiles = files.filter((f) => f.endsWith(".md"));

      for (const file of mdFiles) {
        const filePath = join(catDir, file);
        const entry = await this.parseMemoryFile(filePath, category);
        if (entry) this.catalog.push(entry);
      }
    }

    const rootFiles = await readdir(this.memoriesPath).catch(() => []);
    const rootMdFiles = (rootFiles as string[]).filter(
      (f) => f.endsWith(".md") && f !== "README.md"
    );

    for (const file of rootMdFiles) {
      const filePath = join(this.memoriesPath, file);
      const fileStat = await stat(filePath);
      if (!fileStat.isFile()) continue;
      const entry = await this.parseMemoryFile(filePath);
      if (entry) this.catalog.push(entry);
    }

    this.catalog.sort(
      (a, b) => new Date(b.date).getTime() - new Date(a.date).getTime()
    );

    if (this.embeddings.isReady()) {
      await this.syncVectorStore();
    }

    this.loaded = true;
  }

  private async syncVectorStore(): Promise<void> {
    const dbPath = join(this.memoriesPath, ".lancedb");
    this.db = await lancedb.connect(dbPath);

    const records: Record<string, unknown>[] = [];

    for (const entry of this.catalog) {
      const text = this.buildEmbeddingText(entry);
      const vector = await this.embeddings.embed(text);
      records.push(vectorRecord(entry, vector, this.contentHash(text)));
    }

    if (records.length > 0) {
      this.table = await this.db.createTable("memories", records, {
        mode: "overwrite",
      });
      console.error(`Indexed ${records.length} memories in LanceDB`);
    } else {
      this.table = null;
    }
  }

  private buildEmbeddingText(entry: MemoryEntry): string {
    const parts = [
      entry.title,
      `Category: ${entry.category}`,
      entry.tags.length > 0 ? `Tags: ${entry.tags.join(", ")}` : "",
      entry.content,
    ];
    return parts.filter(Boolean).join("\n");
  }

  private contentHash(text: string): string {
    return createHash("sha256").update(text).digest("hex").slice(0, 16);
  }

  private parseFrontmatter(raw: string): { data: Partial<MemoryFrontmatter>; content: string } {
    const match = raw.match(/^---\n([\s\S]*?)\n---\n?([\s\S]*)$/);
    if (!match) return { data: {}, content: raw };

    const data: Record<string, unknown> = {};
    for (const line of match[1].split("\n")) {
      const idx = line.indexOf(":");
      if (idx === -1) continue;
      const key = line.slice(0, idx).trim();
      let value: unknown = line.slice(idx + 1).trim();

      if (typeof value === "string" && value.startsWith("[") && value.endsWith("]")) {
        value = value.slice(1, -1).split(",").map((s) => s.trim());
      }
      data[key] = value;
    }

    return { data: data as Partial<MemoryFrontmatter>, content: match[2] };
  }

  private buildFrontmatter(data: Record<string, unknown>): string {
    const lines = ["---"];
    for (const [key, value] of Object.entries(data)) {
      if (value === undefined) continue;
      if (Array.isArray(value)) {
        lines.push(`${key}: [${value.join(", ")}]`);
      } else {
        lines.push(`${key}: ${value}`);
      }
    }
    lines.push("---");
    return lines.join("\n");
  }

  private async parseMemoryFile(
    filePath: string,
    fallbackCategory?: string
  ): Promise<MemoryEntry | null> {
    try {
      const raw = await readFile(filePath, "utf-8");
      const { data: fm, content } = this.parseFrontmatter(raw);

      const id = basename(filePath, ".md");
      const category = (fm.category || fallbackCategory || "domain") as MemoryCategory;

      return {
        id,
        path: filePath,
        title: fm.title || id,
        category,
        date: fm.date || new Date().toISOString().split("T")[0],
        author: fm.author,
        tags: Array.isArray(fm.tags) ? fm.tags : [],
        status: (fm.status as MemoryStatus) || "active",
        content: content.trim(),
      };
    } catch {
      return null;
    }
  }

  private async ensureLoaded(): Promise<void> {
    if (this.loaded) return;
    if (this.loadingPromise) {
      await this.loadingPromise;
      return;
    }
    throw new MemoryMCPError("Store not loaded. Call load() first.", "NOT_LOADED");
  }

  async search(
    query: string,
    opts?: { category?: MemoryCategory; status?: MemoryStatus; limit?: number }
  ): Promise<(MemoryCatalogEntry & { score: number })[]> {
    await this.ensureLoaded();

    const status = opts?.status || "active";
    const limit = opts?.limit || 10;

    if (this.embeddings.isReady() && this.table) {
      return this.semanticSearch(query, { category: opts?.category, status, limit });
    }

    let filtered = this.catalog.filter((m) => m.status === status);
    if (opts?.category) {
      filtered = filtered.filter((m) => m.category === opts.category);
    }
    return this.keywordSearch(query, filtered, limit);
  }

  private async semanticSearch(
    query: string,
    opts: { category?: MemoryCategory; status?: string; limit: number }
  ): Promise<(MemoryCatalogEntry & { score: number })[]> {
    const queryVector = await this.embeddings.embed(query);

    const filters: string[] = [];
    if (opts.status) filters.push(`status = '${opts.status}'`);
    if (opts.category) filters.push(`category = '${opts.category}'`);

    const searchQuery = (this.table!.search(queryVector) as VectorQuery).distanceType("cosine").limit(opts.limit);

    if (filters.length > 0) {
      searchQuery.where(filters.join(" AND "));
    }

    const results = await searchQuery.toArray();

    const MIN_RELEVANCE_SCORE = -0.2;

    return results
      .map((row: Record<string, unknown>) => ({
        id: row.id as string,
        title: row.title as string,
        category: row.category as MemoryCategory,
        date: row.date as string,
        tags: (row.tags as string).split(",").filter(Boolean),
        status: row.status as MemoryStatus,
        snippet: row.snippet as string,
        score: round(1 - ((row._distance as number) || 0)),
      }))
      .filter((r) => r.score >= MIN_RELEVANCE_SCORE);
  }

  private keywordSearch(
    query: string,
    entries: MemoryEntry[],
    limit: number
  ): (MemoryCatalogEntry & { score: number })[] {
    const queryTerms = query.toLowerCase().split(/\s+/).filter(Boolean);

    const scored = entries.map((entry) => {
      let score = 0;

      for (const term of queryTerms) {
        if (entry.title.toLowerCase().includes(term)) score += 10;
        if (entry.tags.some((t) => t.toLowerCase().includes(term))) score += 8;
        if (entry.category.toLowerCase().includes(term)) score += 5;
        if (entry.content.toLowerCase().includes(term)) score += 3;
        if (entry.author?.toLowerCase().includes(term)) score += 2;
      }

      return { entry, score: score / 10 };
    });

    return scored
      .filter((s) => s.score > 0)
      .sort((a, b) => b.score - a.score)
      .slice(0, limit)
      .map((s) => ({ ...this.toCatalogEntry(s.entry), score: round(s.score) }));
  }

  async list(opts?: {
    category?: MemoryCategory;
    status?: MemoryStatus;
    limit?: number;
  }): Promise<MemoryCatalogEntry[]> {
    await this.ensureLoaded();

    let filtered = [...this.catalog];

    if (opts?.category) {
      filtered = filtered.filter((m) => m.category === opts.category);
    }
    if (opts?.status) {
      filtered = filtered.filter((m) => m.status === opts.status);
    }

    const limit = opts?.limit || 50;
    return filtered.slice(0, limit).map((e) => this.toCatalogEntry(e));
  }

  async get(id: string): Promise<MemoryEntry | null> {
    await this.ensureLoaded();
    return this.catalog.find((m) => m.id === id) || null;
  }

  async add(params: {
    title: string;
    category: MemoryCategory;
    content: string;
    tags?: string[];
    author?: string;
  }): Promise<MemoryEntry> {
    const catDir = join(this.memoriesPath, params.category);
    await mkdir(catDir, { recursive: true });

    const date = new Date().toISOString().split("T")[0];
    const slug = params.title
      .toLowerCase()
      .replace(/[^a-z0-9]+/g, "-")
      .replace(/^-|-$/g, "");
    const id = `${date}-${slug}`;
    const filePath = join(catDir, `${id}.md`);

    const fm: Record<string, unknown> = {
      title: params.title,
      category: params.category,
      date,
      status: "active",
    };
    if (params.author) fm.author = params.author;
    if (params.tags?.length) fm.tags = params.tags;

    const fileContent = `${this.buildFrontmatter(fm)}\n\n${params.content}\n`;
    await writeFile(filePath, fileContent, "utf-8");

    const entry: MemoryEntry = {
      id,
      path: filePath,
      title: params.title,
      category: params.category,
      date,
      author: params.author,
      tags: params.tags || [],
      status: "active",
      content: params.content,
    };

    this.catalog.unshift(entry);

    if (this.embeddings.isReady() && this.db) {
      const text = this.buildEmbeddingText(entry);
      const vector = await this.embeddings.embed(text);
      const record = vectorRecord(entry, vector, this.contentHash(text));

      if (this.table) {
        await this.table.add([record]);
      } else {
        this.table = await this.db.createTable("memories", [record], { mode: "overwrite" });
      }
    }

    return entry;
  }

  async archive(id: string): Promise<MemoryEntry> {
    const entry = await this.get(id);
    if (!entry) {
      throw new MemoryMCPError(`Memory "${id}" not found`, "NOT_FOUND");
    }

    entry.status = "archived";

    const raw = await readFile(entry.path, "utf-8");
    const { data, content } = this.parseFrontmatter(raw);
    data.status = "archived";
    const updated = `${this.buildFrontmatter(data as Record<string, unknown>)}\n${content}`;
    await writeFile(entry.path, updated, "utf-8");

    if (this.table) {
      await this.table.update({ where: `id = '${id}'`, values: { status: "archived" } });
    }

    return entry;
  }

  async remove(id: string): Promise<void> {
    const entry = await this.get(id);
    if (!entry) {
      throw new MemoryMCPError(`Memory "${id}" not found`, "NOT_FOUND");
    }

    await rm(entry.path);
    this.catalog = this.catalog.filter((m) => m.id !== id);

    if (this.table) {
      await this.table.delete(`id = '${id}'`);
    }
  }

  private toCatalogEntry(entry: MemoryEntry): MemoryCatalogEntry {
    const snippet =
      entry.content.length > 200
        ? entry.content.slice(0, 200) + "..."
        : entry.content;

    return {
      id: entry.id,
      title: entry.title,
      category: entry.category,
      date: entry.date,
      author: entry.author,
      tags: entry.tags,
      status: entry.status,
      snippet,
    };
  }
}

function round(n: number): number {
  return Math.round(n * 1000) / 1000;
}