PubScript/packages/ai/src/suggestions.ts

/**
 * SuggestionBuilder — turns an agent stream into a structured, pending
 * suggestion (CHUNK #9).
 *
 * Each suggestion kind has a tailored system prompt that shapes the model's
 * output. The streamed text is parsed into a title, summary, and full body,
 * then wrapped in the core `AISuggestion` shape with `status: "pending"`.
 *
 * See CHUNK#9.md §3.
 */
import { createUlid, nowIso } from "@pubscript/core";

import { AgentController } from "./agent.js";
import type { Message } from "./providers/types.js";

/** The kinds of suggestions the builder can generate. */
export type SuggestionKind =
  | "brainstorm"
  | "outline"
  | "expand"
  | "tighten"
  | "continue"
  | "research"
  | "dialogue"
  | "worldbuilding";

/** Context about the current scene, used to tailor the system prompt. */
export interface SuggestionContext {
  chapterTitle?: string;
  characters?: string[];
  setting?: string;
  previous?: string;
}

/** A request to build a suggestion. */
export interface SuggestionPrompt {
  kind: SuggestionKind;
  prompt: string;
  context?: SuggestionContext;
}

/** The parsed, pending suggestion produced by the builder. */
export interface SuggestionResult {
  id: string;
  sessionId: string;
  kind: SuggestionKind;
  title: string;
  text: string;
  summary: string;
  sourceMessageId: string;
  status: "pending";
  createdAt: string;
}

/** The system prompt for each suggestion kind. */
const SYSTEM_PROMPTS: Record<SuggestionKind, (ctx: SuggestionContext) => string> = {
  brainstorm: (ctx) => [
    "You are a creative writing collaborator helping an author brainstorm ideas.",
    ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "",
    ctx.characters && ctx.characters.length ? `Relevant characters: ${ctx.characters.join(", ")}.` : "",
    ctx.setting ? `The story is set in ${ctx.setting}.` : "",
    "Propose several distinct, concrete directions the story could take next. Be specific and evocative, but concise. Avoid prose that reads like finished text — these are seeds, not drafts.",
  ]
    .filter(Boolean)
    .join(" "),
  outline: (ctx) => [
    "You are a story-structuring assistant helping an author outline a scene or chapter.",
    ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "",
    ctx.previous ? `The preceding material: ${ctx.previous}.` : "",
    ctx.characters && ctx.characters.length ? `Characters in the scene: ${ctx.characters.join(", ")}.` : "",
    ctx.setting ? `The story is set in ${ctx.setting}.` : "",
    "Produce a clear, ordered outline (numbered beats or a short bulleted sequence) that advances the scene. Each beat should be one or two sentences.",
  ]
    .filter(Boolean)
    .join(" "),
  expand: (ctx) => [
    "You are a prose-expansion assistant helping an author flesh out a scene.",
    ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "",
    ctx.characters && ctx.characters.length ? `Characters in the scene: ${ctx.characters.join(", ")}.` : "",
    "Expand the author's outline or summary into richer, more detailed prose. Preserve the author's voice and intent; add sensory detail, pacing, and subtext without changing the plot.",
  ]
    .filter(Boolean)
    .join(" "),
  tighten: (ctx) => [
    "You are a line-editing assistant helping an author tighten prose.",
    ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "",
    "Identify redundancy, wordiness, and weak verbs, then return a tightened version of the text. Preserve meaning, voice, and intent. Make only necessary changes.",
  ]
    .filter(Boolean)
    .join(" "),
  continue: (ctx) => [
    "You are a collaborative fiction writer continuing a story in the author's voice.",
    ctx.chapterTitle ? `The current chapter is titled "${ctx.chapterTitle}".` : "",
    ctx.previous ? `What has happened so far: ${ctx.previous}.` : "",
    ctx.characters && ctx.characters.length ? `Characters present: ${ctx.characters.join(", ")}.` : "",
    "Write the next segment of the story, continuing naturally from what precedes it. Match the established voice, tense, and POV. Do not summarize or skip ahead — write the scene.",
  ]
    .filter(Boolean)
    .join(" "),
  research: (ctx) => [
    "You are a research assistant for a fiction author.",
    ctx.setting ? `The story is set in ${ctx.setting}.` : "",
    "Provide accurate, well-organized factual background that supports the author's worldbuilding and plot. Distinguish established fact from plausible inference. Keep it practical and citation-aware.",
  ]
    .filter(Boolean)
    .join(" "),
  dialogue: (ctx) => [
    "You are a dialogue-writing assistant helping an author craft natural, character-specific dialogue.",
    ctx.characters && ctx.characters.length ? `Speaking characters: ${ctx.characters.join(", ")}.` : "",
    ctx.setting ? `The scene is set in ${ctx.setting}.` : "",
    "Write dialogue that sounds natural, reveals character, and carries subtext. Attribute speakers clearly. Avoid on-the-nose exposition.",
  ]
    .filter(Boolean)
    .join(" "),
  worldbuilding: (ctx) => [
    "You are a worldbuilding assistant helping an author develop a consistent, vivid fictional world.",
    ctx.setting ? `The setting under development: ${ctx.setting}.` : "",
    "Develop the world consistently: geography, culture, politics, technology, magic systems, economics, or daily life. Flag any potential inconsistencies with established canon and suggest resolutions.",
  ]
    .filter(Boolean)
    .join(" "),
};

/**
 * Builds structured, pending suggestions from an agent stream.
 */
export class SuggestionBuilder {
  constructor(private readonly agent: AgentController) {}

  /**
   * Build a pending suggestion from a prompt. Streams via the agent, then
   * parses the response into a structured suggestion.
   */
  async build(
    prompt: SuggestionPrompt,
    sessionId: string,
    sourceMessageId: string,
  ): Promise<SuggestionResult> {
    const system = SYSTEM_PROMPTS[prompt.kind](prompt.context ?? {});
    const messages: Message[] = [
      { role: "system", content: system },
      { role: "user", content: prompt.prompt },
    ];

    let text = "";
    await this.agent.collect(messages, {}, { onText: (t) => (text += t) });

    const parsed = this.parse(prompt.kind, text);
    return {
      id: createUlid(),
      sessionId,
      kind: prompt.kind,
      title: parsed.title,
      text: parsed.text,
      summary: parsed.summary,
      sourceMessageId,
      status: "pending",
      createdAt: nowIso(),
    };
  }

  /** Parse streamed text into a title, summary, and body. */
  private parse(kind: SuggestionKind, text: string): { title: string; summary: string; text: string } {
    const cleaned = text.trim();
    const firstLine = cleaned.split(/\r?\n/).find((l) => l.trim().length > 0) ?? "";
    const title = firstLine.replace(/^#+\s*/, "").trim().slice(0, 120) || `${kind} suggestion`;
    const summary = cleaned.split(/\r?\n\r?\n/)[0]!.replace(/^#+\s*/, "").trim().slice(0, 300);
    return { title, summary, text: cleaned };
  }
}