STCourses/packages/feature-ai/src/generator/generate.ts

import { stat } from "node:fs/promises";
import { join } from "node:path";
import type { AiProvider } from "@stc/contracts";
import { resolveInside, type CourseRepository, type LoadedCourse } from "@stc/course-format";
import { draftLocation, safeRelPath, type DraftFile, type DraftSink } from "./draft-sink";
import { OUTPUT_FORMAT, REQUIRED, VOICE, styleTemplate, type Kind, type Style } from "./templates";

export type GenerateSource =
  | { kind: "lesson"; lessonId: string }
  | { kind: "file"; path: string }
  | { kind: "cache"; path: string }
  | { kind: "text"; label: string; text: string };

export interface GenerateRequest {
  courseId: string;
  kind: Kind;
  style: Style;
  topic: string;
  sources: GenerateSource[];
  instructions?: string;
}

/** `location` is a folder (filesystem sink) or `draft/<author> @ <commit>` (studio writer). */
export interface GenerateResult { location: string; branch?: string; commit?: string; paths: string[]; warnings: string[] }

const MAX_MATERIAL = 400_000;

export function slugify(s: string): string {
  return s.normalize("NFKD").replace(/[̀-ͯ]/g, "").toLowerCase().replace(/[^a-z0-9]+/g, "-").replace(/^-+|-+$/g, "").slice(0, 60) || "draft";
}

export function targetPaths(req: GenerateRequest): string[] {
  const slug = slugify(req.topic);
  if (req.kind === "exercise") return [`exercises/${slug}/README.md`, `exercises/${slug}/starter.js`, `exercises/${slug}/tests.js`, `exercises/${slug}/solution.js`, `exercises/${slug}/fixtures/${slug}.json`, `exercises/${slug}/manifest-snippet.yaml`];
  if (req.kind === "questions") return [`questions/${slug}.md`];
  return req.style === "systemic-design" ? [`course/${slug}/lesson.md`, `course/${slug}/scenarios.md`] : [`course/${slug}.md`];
}

/** Reads the chosen inputs. Authors may use any course file, including solutions; cache text comes from <dataDir>/cache/<courseId>/. */
export async function gatherMaterial(repo: CourseRepository, course: LoadedCourse, sources: GenerateSource[], dataDir: string): Promise<{ text: string; warnings: string[] }> {
  const parts: string[] = [];
  const warnings: string[] = [];
  let total = 0;
  const push = (label: string, body: string) => {
    if (total >= MAX_MATERIAL) { warnings.push(`Left out "${label}": the material is already ${MAX_MATERIAL} characters.`); return; }
    let text = body;
    if (total + text.length > MAX_MATERIAL) {
      text = text.slice(0, MAX_MATERIAL - total);
      warnings.push(`"${label}" was shortened to fit the material limit.`);
    }
    total += text.length;
    parts.push(`<source label="${label.replace(/"/g, "'")}">\n${text}\n</source>`);
  };
  const readText = async (abs: string | null, label: string) => {
    const info = abs ? await stat(abs).catch(() => null) : null;
    if (!abs || !info?.isFile()) { warnings.push(`Source not found: ${label}`); return; }
    const text = await Bun.file(abs).text();
    if (text.includes("\u0000")) { warnings.push(`Skipped binary file: ${label}`); return; }
    push(label, text);
  };
  for (const s of sources) {
    if (s.kind === "text") push(s.label || "Pasted text", s.text);
    else if (s.kind === "lesson") {
      const e = course.lessonIndex.get(s.lessonId);
      const md = e ? await repo.readLessonMarkdown(course.id, s.lessonId) : null;
      if (!e || md === null) warnings.push(`Lesson has no readable text: ${s.lessonId}`);
      else push(`${e.lesson.type === "video" ? e.lesson.transcript : e.lesson.file} (lesson ${s.lessonId}: ${e.lesson.title})`, md);
    } else if (s.kind === "file") await readText(repo.resolvePath(course.id, s.path), s.path);
    else await readText(resolveInside(join(dataDir, "cache", course.id), s.path), `cache/${s.path}`);
  }
  return { text: parts.join("\n\n"), warnings };
}

function outline(course: LoadedCourse): string {
  return (course.manifest?.modules ?? []).map((m, i) => `${i + 1}. ${m.title} (${m.lessons.length} lessons)`).join("\n");
}

export function buildGeneratorPrompt(req: GenerateRequest, material: string, course: LoadedCourse): { system: string; user: string } {
  const what = req.kind === "document" ? "an explanatory lesson document" : req.kind === "exercise" ? "a small exercise with tests and hint tiers" : "check-your-understanding questions";
  const system = [
    `You draft course material for STCourses authors. You write ${what} for the course "${course.manifest?.title ?? course.id}", in the house style below, from the source material the author selected.`,
    VOICE,
    styleTemplate(req.kind, req.style),
    OUTPUT_FORMAT,
  ].join("\n\n");
  const tracks = course.manifest?.tracks?.map((t) => `${t.label} (${t.id})`).join(", ") || "one track (JavaScript)";
  const user = [
    `Course: ${course.manifest?.title ?? course.id}`,
    `Language tracks: ${tracks}`,
    `Existing modules:\n${outline(course) || "(none yet)"}`,
    `Topic: ${req.topic}`,
    req.instructions?.trim() ? `Author's instructions:\n${req.instructions.trim()}` : "",
    `Write these files:\n${targetPaths(req).map((p) => `- ${p}`).join("\n")}`,
    `Source material (transcripts, text, slide text and code chosen by the author):\n\n${material || "(none: work from the topic and instructions, and mark facts TODO(author): check)"}`,
  ].filter(Boolean).join("\n\n");
  return { system, user };
}

/** Reads `=== FILE: path ===` … `=== END FILE ===` blocks. Unsafe paths are dropped with a warning. */
export function parseFiles(text: string): { files: DraftFile[]; warnings: string[] } {
  const files: DraftFile[] = [];
  const warnings: string[] = [];
  const lines = text.split(/\r?\n/);
  let cur: { path: string; lines: string[] } | null = null;
  const close = () => {
    if (!cur) return;
    const rel = safeRelPath(cur.path);
    if (!rel) warnings.push(`Dropped a file with an unsafe path: ${cur.path}`);
    else files.push({ path: rel, content: cur.lines.join("\n").replace(/\n*$/, "\n") });
    cur = null;
  };
  for (const line of lines) {
    const open = /^===\s*FILE:\s*(.+?)\s*===\s*$/.exec(line);
    if (open) { close(); cur = { path: open[1]!.replace(/^`|`$/g, ""), lines: [] }; continue; }
    if (/^===\s*END FILE\s*===\s*$/.test(line)) { close(); continue; }
    cur?.lines.push(line);
  }
  if (cur) { warnings.push(`The last file (${(cur as { path: string }).path}) was not closed; it may be cut off.`); close(); }
  if (!files.length) warnings.push("The model returned no files.");
  return { files, warnings };
}

export function checkHouseStyle(req: GenerateRequest, files: DraftFile[]): string[] {
  const warnings: string[] = [];
  const md = files.filter((f) => /\.(md|markdown)$/i.test(f.path)).map((f) => f.content).join("\n");
  for (const s of REQUIRED[req.kind][req.style]) if (!md.includes(s)) warnings.push(`House style: missing "${s}".`);
  if (req.kind === "exercise") {
    for (const name of ["starter.js", "tests.js", "solution.js"]) if (!files.some((f) => f.path.endsWith(`/${name}`))) warnings.push(`Exercise: missing ${name}.`);
    const tests = files.find((f) => f.path.endsWith("/tests.js"));
    if (tests && !/export\s+const\s+tests\b/.test(tests.content)) warnings.push("Exercise: tests.js does not export `tests`.");
  }
  return warnings;
}

const marker = (req: GenerateRequest) =>
  `<!-- Draft generated by STCourses AI on ${new Date().toISOString().slice(0, 10)} (${req.kind}, ${req.style} style) from author-selected sources. An author must review it before publishing. -->\n\n`;

/** Runs the model once and writes the files it returns through the DraftSink. Never publishes. */
export async function generateDraft(o: {
  provider: AiProvider;
  repo: CourseRepository;
  req: GenerateRequest;
  sink: DraftSink;
  authorId: string;
  dataDir: string;
  signal?: AbortSignal;
  onProgress?: (fraction: number, message: string) => Promise<void> | void;
}): Promise<GenerateResult> {
  const course = o.repo.get(o.req.courseId);
  if (!course || course.status !== "ok") throw new Error(`Course not available: ${o.req.courseId}`);
  await o.onProgress?.(0.05, "Reading sources");
  const material = await gatherMaterial(o.repo, course, o.req.sources, o.dataDir);
  const prompt = buildGeneratorPrompt(o.req, material.text, course);
  await o.onProgress?.(0.15, "Drafting");
  let text = "";
  let chars = 0;
  for await (const ev of o.provider.chat([{ role: "system", content: prompt.system }, { role: "user", content: prompt.user }], { signal: o.signal })) {
    if (ev.text) {
      text += ev.text;
      chars += ev.text.length;
      if (chars > 4000) { chars = 0; await o.onProgress?.(Math.min(0.85, 0.15 + text.length / 60_000), "Drafting"); }
    }
  }
  const parsed = parseFiles(text);
  if (!parsed.files.length) throw new Error(parsed.warnings.join(" ") || "The model returned no files.");
  const files = parsed.files.map((f) => (/\.(md|markdown)$/i.test(f.path) ? { ...f, content: marker(o.req) + f.content } : f));
  await o.onProgress?.(0.9, "Writing drafts");
  const written = await o.sink.writeDraftFiles(course.id, o.authorId, files, `AI draft: ${o.req.kind} "${o.req.topic}" (${o.req.style} style)`);
  return {
    location: draftLocation(written),
    ...(written.branch ? { branch: written.branch } : {}),
    ...(written.commit ? { commit: written.commit } : {}),
    paths: files.map((f) => f.path),
    warnings: [...material.warnings, ...parsed.warnings, ...checkHouseStyle(o.req, parsed.files)],
  };
}