cortex/src/core/store.ts

import { randomUUID as uuid } from 'crypto';
import { getDb } from './db';
import { Node, Edge, AddNodeInput, UpdateNodeInput, ListOptions, QueryOptions, SearchResult, EdgeType } from '../types';
import { hybridSearch, deserializeEmbedding } from './search/index';
import { getEmbedding } from './search/ollama';

function rowToNode(row: any): Node {
  return {
    id: row.id,
    kind: row.kind,
    title: row.title,
    content: row.content,
    status: row.status ?? undefined,
    tags: JSON.parse(row.tags || '[]'),
    metadata: JSON.parse(row.metadata || '{}'),
    embedding: row.embedding ? deserializeEmbedding(row.embedding) : null,
    createdAt: row.created_at,
    updatedAt: row.updated_at,
    lastAccessedAt: row.last_accessed_at ?? row.updated_at,
    isStale: !!row.is_stale,
  };
}

function serializeEmbedding(embedding: number[]): Buffer {
  return Buffer.from(new Float32Array(embedding).buffer);
}

function notifyDirty(): void {
  try {
    const { markDirty } = require('../server/heartbeat');
    markDirty();
  } catch {}
}

export async function addNode(input: AddNodeInput): Promise<Node> {
  const db = getDb();
  const id = uuid();
  const now = Date.now();
  const tags = input.tags ?? [];
  const metadata = input.metadata ?? {};
  const content = input.content ?? '';

  // Try to get embedding
  const embedding = await getEmbedding(`${input.title} ${content}`);

  db.prepare(`
    INSERT INTO nodes (id, kind, title, content, status, tags, metadata, embedding, created_at, updated_at, last_accessed_at)
    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
  `).run(
    id, input.kind, input.title, content, input.status ?? null,
    JSON.stringify(tags), JSON.stringify(metadata),
    embedding ? serializeEmbedding(embedding) : null,
    now, now, now
  );

  // Insert tags
  const insertTag = db.prepare('INSERT OR IGNORE INTO node_tags (node_id, tag) VALUES (?, ?)');
  for (const tag of tags) {
    insertTag.run(id, tag);
  }

  notifyDirty();
  return {
    id, kind: input.kind, title: input.title, content, status: input.status,
    tags, metadata, embedding, createdAt: now, updatedAt: now, lastAccessedAt: now, isStale: false,
  };
}

export function getNode(id: string): Node | null {
  const db = getDb();
  const row = db.prepare('SELECT * FROM nodes WHERE id = ?').get(id) as any;
  if (!row) return null;
  db.prepare('UPDATE nodes SET last_accessed_at = ? WHERE id = ?').run(Date.now(), id);
  return rowToNode(row);
}

export function findNodeByPrefix(prefix: string): Node | null {
  const db = getDb();
  const row = db.prepare('SELECT * FROM nodes WHERE id LIKE ?').get(`${prefix}%`) as any;
  return row ? rowToNode(row) : null;
}

export function listNodes(options: ListOptions = {}): Node[] {
  const db = getDb();
  const conditions: string[] = [];
  const params: any[] = [];

  if (!options.includeStale) {
    conditions.push('is_stale = 0');
  }
  if (options.kind) {
    conditions.push('kind = ?');
    params.push(options.kind);
  }
  if (options.status) {
    conditions.push('status = ?');
    params.push(options.status);
  }

  let sql = 'SELECT * FROM nodes';
  if (conditions.length) sql += ' WHERE ' + conditions.join(' AND ');
  sql += ' ORDER BY created_at DESC';
  if (options.limit) {
    sql += ' LIMIT ?';
    params.push(options.limit);
  }

  let nodes = (db.prepare(sql).all(...params) as any[]).map(rowToNode);

  if (options.tags?.length) {
    nodes = nodes.filter(n => options.tags!.some(t => n.tags.includes(t)));
  }

  return nodes;
}

export async function updateNode(id: string, input: UpdateNodeInput): Promise<Node | null> {
  const db = getDb();
  const existing = getNode(id);
  if (!existing) return null;

  const now = Date.now();
  const sets: string[] = ['updated_at = ?'];
  const params: any[] = [now];

  if (input.title !== undefined) { sets.push('title = ?'); params.push(input.title); }
  if (input.content !== undefined) { sets.push('content = ?'); params.push(input.content); }
  if (input.status !== undefined) { sets.push('status = ?'); params.push(input.status); }
  if (input.isStale !== undefined) { sets.push('is_stale = ?'); params.push(input.isStale ? 1 : 0); }
  if (input.tags !== undefined) { sets.push('tags = ?'); params.push(JSON.stringify(input.tags)); }
  if (input.metadata !== undefined) {
    const merged = { ...existing.metadata, ...input.metadata };
    sets.push('metadata = ?');
    params.push(JSON.stringify(merged));
  }

  // Re-embed if title or content changed
  if (input.title !== undefined || input.content !== undefined) {
    const newTitle = input.title ?? existing.title;
    const newContent = input.content ?? existing.content;
    const embedding = await getEmbedding(`${newTitle} ${newContent}`);
    if (embedding) {
      sets.push('embedding = ?');
      params.push(serializeEmbedding(embedding));
    }
  }

  params.push(id);
  db.prepare(`UPDATE nodes SET ${sets.join(', ')} WHERE id = ?`).run(...params);

  // Update tags if changed
  if (input.tags !== undefined) {
    db.prepare('DELETE FROM node_tags WHERE node_id = ?').run(id);
    const insertTag = db.prepare('INSERT OR IGNORE INTO node_tags (node_id, tag) VALUES (?, ?)');
    for (const tag of input.tags) {
      insertTag.run(id, tag);
    }
  }

  notifyDirty();
  return getNode(id);
}

export function removeNode(id: string, hard: boolean = false): boolean {
  const db = getDb();
  if (hard) {
    const result = db.prepare('DELETE FROM nodes WHERE id = ?').run(id);
    return result.changes > 0;
  } else {
    const result = db.prepare('UPDATE nodes SET is_stale = 1, updated_at = ? WHERE id = ?').run(Date.now(), id);
    if (result.changes > 0) notifyDirty();
    return result.changes > 0;
  }
}

export function addEdge(fromId: string, toId: string, type: EdgeType, metadata: Record<string, any> = {}): Edge {
  const db = getDb();
  const id = uuid();
  const now = Date.now();

  db.prepare(`
    INSERT INTO edges (id, from_id, to_id, type, metadata, created_at) VALUES (?, ?, ?, ?, ?, ?)
  `).run(id, fromId, toId, type, JSON.stringify(metadata), now);

  return { id, fromId, toId, type, metadata, createdAt: now };
}

export function removeEdge(id: string): boolean {
  const db = getDb();
  const result = db.prepare('DELETE FROM edges WHERE id = ?').run(id);
  return result.changes > 0;
}

export async function query(text: string, options: QueryOptions = {}): Promise<SearchResult[]> {
  const nodes = listNodes({ kind: options.kind, tags: options.tags, includeStale: options.includeStale });
  return hybridSearch(nodes, text, options);
}