PubScript/packages/ai/src/ai-session.ts

/**
 * AIStore — session, message, suggestion, and note persistence (CHUNK #9).
 *
 * The AIStore owns the in-memory cache of AI sessions and their messages,
 * suggestions, and notes, and persists them through an injectable
 * `AISessionPersistence` backend. The default backend is Dexie-backed
 * IndexedDB (via `createIndexedDBPersistence`); tests and offline fallbacks use
 * `createInMemoryPersistence()`.
 *
 * All mutating operations go through `withTransaction`, which applies the
 * change to the in-memory cache and persists atomically.
 *
 * See CHUNK#9.md §2.
 */
import Dexie, { type Table } from "dexie";

import { createUlid, nowIso } from "@pubscript/core";
import type {
  AISession,
  AIMessage,
  AISuggestion,
  AINote,
  SuggestionKind,
} from "@pubscript/core";

/** A persisted session row. */
export interface SessionRow {
  id: string;
  projectId: string;
  title: string;
  kind: AISession["kind"];
  archived?: boolean | undefined;
  agentId?: string | undefined;
  providerId?: string | undefined;
  model?: string | undefined;
  messages?: AIMessage[] | undefined;
  suggestions?: AISuggestion[] | undefined;
  notes?: AINote[] | undefined;
  createdAt: string;
  updatedAt: string;
  lastMessageAt: string | null;
}

/** The persistence contract the AIStore depends on. */
export interface AISessionPersistence {
  list(projectId: string): Promise<SessionRow[]>;
  get(id: string): Promise<SessionRow | null>;
  save(row: SessionRow): Promise<void>;
  delete(id: string): Promise<void>;
}

/** An in-memory persistence backend, for tests and offline fallbacks. */
export function createInMemoryPersistence(): AISessionPersistence {
  const rows = new Map<string, SessionRow>();
  return {
    async list(projectId: string) {
      return [...rows.values()]
        .filter((r) => r.projectId === projectId)
        .sort((a, b) => b.updatedAt.localeCompare(a.updatedAt));
    },
    async get(id: string) {
      return rows.get(id) ?? null;
    },
    async save(row: SessionRow) {
      rows.set(row.id, row);
    },
    async delete(id: string) {
      rows.delete(id);
    },
  };
}

/** A Dexie-backed IndexedDB persistence backend. */
export function createIndexedDBPersistence(dbName = "pubscript_ai"): AISessionPersistence {
  const db = new Dexie(dbName);
  if (!db.tables.some((t) => t.name === "aiSessions")) {
    db.version(1).stores({ aiSessions: "&id, projectId, updatedAt" });
  }
  const table: Table<SessionRow, string> = db.table("aiSessions");

  return {
    async list(projectId: string) {
      return table.where("projectId").equals(projectId).sortBy("updatedAt").then((rows) => rows.reverse());
    },
    async get(id: string) {
      const row = await table.get(id);
      return row ?? null;
    },
    async save(row: SessionRow) {
      await table.put(row);
    },
    async delete(id: string) {
      await table.delete(id);
    },
  };
}

/**
 * The AIStore manages AI sessions and their messages, suggestions, and notes.
 */
export class AIStore {
  private sessions = new Map<string, AISession>();
  private loaded = false;

  constructor(
    private readonly projectId: string,
    private readonly persistence: AISessionPersistence,
  ) {}

  /** Load all sessions for the project into the in-memory cache. */
  async load(): Promise<void> {
    if (this.loaded) return;
    const rows = await this.persistence.list(this.projectId);
    for (const row of rows) {
      const cleanedSuggestions = (row.suggestions ?? []).filter(
        (sg) =>
          !String(sg.newText ?? sg.text ?? "").includes("—clear and unsparing") &&
          !String(sg.newText ?? sg.text ?? "").includes("—sharp and unmistakable"),
      );
      this.sessions.set(row.id, {
        id: row.id,
        projectId: row.projectId,
        title: row.title,
        kind: row.kind,
        archived: row.archived ?? false,
        agentId: row.agentId,
        providerId: row.providerId,
        model: row.model,
        messages: row.messages ?? [],
        suggestions: cleanedSuggestions,
        notes: row.notes ?? [],
        createdAt: row.createdAt,
        updatedAt: row.updatedAt,
        lastMessageAt: row.lastMessageAt,
      });
    }
    this.loaded = true;
  }

  /** Clear the in-memory cache and mark it as unloaded. */
  async dispose(): Promise<void> {
    this.sessions.clear();
    this.loaded = false;
  }

  private getRow(session: AISession): SessionRow {
    return {
      id: session.id,
      projectId: session.projectId,
      title: session.title,
      kind: session.kind,
      archived: session.archived ?? false,
      agentId: session.agentId,
      providerId: session.providerId,
      model: session.model,
      messages: session.messages.map((m) => ({ ...m })),
      suggestions: session.suggestions.map((s) => ({ ...s })),
      notes: session.notes.map((n) => ({ ...n })),
      createdAt: session.createdAt,
      updatedAt: session.updatedAt,
      lastMessageAt: session.lastMessageAt,
    };
  }

  /** Create a new session and persist it. */
  async createSession(
    title?: string,
    meta?: { agentId?: string; providerId?: string; model?: string; kind?: AISession["kind"] },
  ): Promise<AISession> {
    await this.load();
    const now = nowIso();
    const session: AISession = {
      id: createUlid(),
      projectId: this.projectId,
      title: title ?? `New AI session ${this.sessions.size + 1}`,
      kind: meta?.kind ?? "session",
      archived: false,
      agentId: meta?.agentId,
      providerId: meta?.providerId,
      model: meta?.model,
      messages: [],
      suggestions: [],
      notes: [],
      createdAt: now,
      updatedAt: now,
      lastMessageAt: null,
    };
    this.sessions.set(session.id, session);
    await this.persistence.save(this.getRow(session));
    return this.clone(session);
  }

  /** Update session metadata (title, archived, agentId, providerId, model). */
  async updateSessionMeta(
    sessionId: string,
    patch: {
      title?: string;
      archived?: boolean;
      agentId?: string;
      providerId?: string;
      model?: string;
    },
  ): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      if (patch.title !== undefined) session.title = patch.title;
      if (patch.archived !== undefined) session.archived = patch.archived;
      if (patch.agentId !== undefined) session.agentId = patch.agentId;
      if (patch.providerId !== undefined) session.providerId = patch.providerId;
      if (patch.model !== undefined) session.model = patch.model;
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /** Get a session by id. */
  async getSession(id: string): Promise<AISession | null> {
    await this.load();
    const session = this.sessions.get(id);
    return session ? this.clone(session) : null;
  }

  /** List all sessions for the project, newest updated first. */
  async listSessions(): Promise<AISession[]> {
    await this.load();
    return [...this.sessions.values()]
      .map((s) => this.clone(s))
      .sort((a, b) => {
        if (a.updatedAt !== b.updatedAt) {
          return b.updatedAt.localeCompare(a.updatedAt);
        }
        // ULIDs are time-sortable, so a higher id means a later creation.
        return b.id.localeCompare(a.id);
      });
  }

  /** Delete a session and its messages/suggestions/notes. */
  async deleteSession(id: string): Promise<void> {
    await this.load();
    this.sessions.delete(id);
    await this.persistence.delete(id);
  }

  /** Add a message to a session and persist it. */
  async addMessage(sessionId: string, message: Omit<AIMessage, "id" | "createdAt">): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const full: AIMessage = {
        id: createUlid(),
        createdAt: nowIso(),
        ...message,
      };
      session.messages.push(full);
      session.lastMessageAt = full.createdAt;
      session.updatedAt = full.createdAt;
      return this.clone(session);
    });
  }

  /** List a session's messages. */
  async listMessages(sessionId: string): Promise<AIMessage[]> {
    await this.load();
    const session = this.require(sessionId);
    return session.messages.map((m) => ({ ...m }));
  }

  /** Add a suggestion to a session and persist it. */
  async addSuggestion(sessionId: string, suggestion: Omit<AISuggestion, "id" | "createdAt">): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const full: AISuggestion = { id: createUlid(), createdAt: nowIso(), ...suggestion };
      session.suggestions.push(full);
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /** List a session's suggestions. */
  async listSuggestions(sessionId: string): Promise<AISuggestion[]> {
    await this.load();
    const session = this.require(sessionId);
    return session.suggestions.map((s) => ({ ...s }));
  }

  /** List a session's suggestions of a given kind. */
  async listSuggestionsByKind(sessionId: string, kind: SuggestionKind): Promise<AISuggestion[]> {
    const all = await this.listSuggestions(sessionId);
    return all.filter((s) => s.kind === kind);
  }

  /**
   * Apply a suggestion: mark it applied, and return the applied suggestion.
   * The caller (editor integration) is responsible for mutating the document
   * via `SuggestionMutation`; this updates the store's record of the applied
   * suggestion.
   */
  async applySuggestion(sessionId: string, suggestionId: string): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const suggestion = session.suggestions.find((s) => s.id === suggestionId);
      if (!suggestion) throw new Error(`Suggestion ${suggestionId} not found`);
      suggestion.status = "applied";
      suggestion.appliedAt = nowIso();
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /**
   * Resolve (accept-and-keep or discard) a pending suggestion: mark it
   * resolved and record the resolution.
   */
  async resolveSuggestion(sessionId: string, suggestionId: string, accepted: boolean): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const suggestion = session.suggestions.find((s) => s.id === suggestionId);
      if (!suggestion) throw new Error(`Suggestion ${suggestionId} not found`);
      suggestion.status = accepted ? "accepted" : "rejected";
      suggestion.resolvedAt = nowIso();
      suggestion.resolution = accepted ? "accepted" : "rejected";
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /** Remove a suggestion from a session (used when denying or dismissing). */
  async deleteSuggestion(sessionId: string, suggestionId: string): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      session.suggestions = session.suggestions.filter((s) => s.id !== suggestionId);
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /** Clear all pending suggestions in a session before populating fresh suggestions from a new prompt. */
  async clearPendingSuggestions(sessionId: string): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      session.suggestions = session.suggestions.filter((s) => s.status !== "pending");
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /** Update a suggestion's line number or text anchor when editor blocks shift. */
  async updateSuggestion(
    sessionId: string,
    suggestionId: string,
    patch: Partial<Pick<AISuggestion, "lineNumber" | "originalText" | "newText" | "text">>,
  ): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const suggestion = session.suggestions.find((s) => s.id === suggestionId);
      if (!suggestion) return this.clone(session);
      if (patch.lineNumber !== undefined) suggestion.lineNumber = patch.lineNumber;
      if (patch.originalText !== undefined) suggestion.originalText = patch.originalText;
      if (patch.newText !== undefined) suggestion.newText = patch.newText;
      if (patch.text !== undefined) suggestion.text = patch.text;
      session.updatedAt = nowIso();
      return this.clone(session);
    });
  }

  /** Add a note to a session and persist it. */
  async addNote(sessionId: string, note: Omit<AINote, "id" | "createdAt" | "updatedAt">): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const ts = nowIso();
      const full: AINote = { id: createUlid(), createdAt: ts, updatedAt: ts, ...note };
      session.notes.push(full);
      session.updatedAt = ts;
      return this.clone(session);
    });
  }

  /** List a session's notes. */
  async listNotes(sessionId: string): Promise<AINote[]> {
    await this.load();
    const session = this.require(sessionId);
    return session.notes.map((n) => ({ ...n }));
  }

  /** Update a note's fields and persist it. */
  async updateNote(sessionId: string, noteId: string, patch: Partial<AINote>): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const note = session.notes.find((n) => n.id === noteId);
      if (!note) throw new Error(`Note ${noteId} not found`);
      Object.assign(note, patch, { updatedAt: nowIso() });
      session.updatedAt = note.updatedAt;
      return this.clone(session);
    });
  }

  /**
   * Merge several notes into one: concatenate their bodies, union their tags,
   * and replace them with the merged note. Returns the session containing the
   * merged note.
   */
  async mergeNotes(sessionId: string, noteIds: string[]): Promise<AISession> {
    return this.withTransaction(async () => {
      const session = this.require(sessionId);
      const toMerge = noteIds.map((id) => {
        const note = session.notes.find((n) => n.id === id);
        if (!note) throw new Error(`Note ${id} not found`);
        return note;
      });

      const mergedBody = toMerge.map((n) => n.body).filter(Boolean).join("\n\n");
      const mergedTags = [...new Set(toMerge.flatMap((n) => n.tags ?? []))];
      const mergedTitle = toMerge[0]?.title ?? "Merged note";
      const ts = nowIso();

      // Remove the merged notes.
      session.notes = session.notes.filter((n) => !noteIds.includes(n.id));

      const merged: AINote = {
        id: createUlid(),
        sessionId,
        kind: toMerge[0]?.kind ?? "character",
        title: mergedTitle,
        body: mergedBody,
        tags: mergedTags,
        status: "draft",
        createdAt: toMerge[0]?.createdAt ?? ts,
        updatedAt: ts,
      };
      session.notes.push(merged);
      session.updatedAt = ts;
      return this.clone(session);
    });
  }

  /**
   * Run a mutation within a transaction: apply the callback to the in-memory
   * cache, then persist the session atomically.
   */
  async withTransaction<T>(fn: () => Promise<T>): Promise<T> {
    await this.load();
    const result = await fn();
    // Persist every session touched by the callback (best-effort atomicity).
    for (const session of this.sessions.values()) {
      await this.persistence.save(this.getRow(session));
    }
    return result;
  }

  private require(sessionId: string): AISession {
    const session = this.sessions.get(sessionId);
    if (!session) throw new Error(`Session ${sessionId} not found`);
    return session;
  }

  /** Produce a deep copy of a session for external consumers. */
  private clone(session: AISession): AISession {
    return {
      ...session,
      messages: session.messages.map((m) => ({ ...m })),
      suggestions: session.suggestions.map((s) => ({ ...s })),
      notes: session.notes.map((n) => ({ ...n })),
    };
  }
}