src/lib/streaming-artifact.ts

import type { ArtifactDocument } from './agent-tools';
import { prepareArtifactDeployment } from './artifact-runtime';
import { inferArtifactMime, normalizeArtifactPath, type ArtifactFileInput } from './artifact-workspace';
import type { ArtifactDraft, ArtifactDraftFile } from './contracts';

const MAX_STREAMED_FILE_SLOTS = 12;
const SCRIPT_OR_DATA_FILE = /\.(?:js|mjs|cjs|json)$/iu;

export interface StreamingJsonStrings {
  values: Record<string, string>;
  partial?: {
    key: string;
    value: string;
  };
}

function jsonStringEnd(source: string, start: number): number {
  for (let index = start + 1; index < source.length; index += 1) {
    if (source[index] === '\\') {
      index += 1;
      continue;
    }
    if (source[index] === '"') return index;
  }
  return -1;
}

function decodePartialJsonString(raw: string): string {
  let safeLength = raw.length;
  for (let index = 0; index < raw.length; index += 1) {
    const character = raw[index];
    if (character !== '\\') continue;
    const escape = raw[index + 1];
    if (escape === undefined) {
      safeLength = index;
      break;
    }
    if (escape === 'u') {
      const codePoint = raw.slice(index + 2, index + 6);
      if (codePoint.length < 4 || !/^[0-9a-f]{4}$/iu.test(codePoint)) {
        safeLength = index;
        break;
      }
      index += 5;
      continue;
    }
    if (!'"\\/bfnrt'.includes(escape)) {
      safeLength = index;
      break;
    }
    index += 1;
  }
  try {
    return JSON.parse(`"${raw.slice(0, safeLength)}"`) as string;
  } catch {
    return '';
  }
}

export function extractStreamingJsonStrings(source: string): StreamingJsonStrings {
  const values: Record<string, string> = {};
  let cursor = 0;
  while (cursor < source.length) {
    while (cursor < source.length && /[\s,{]/u.test(source[cursor] ?? '')) cursor += 1;
    if (source[cursor] !== '"') break;
    const keyEnd = jsonStringEnd(source, cursor);
    if (keyEnd < 0) break;
    let key: string;
    try {
      key = JSON.parse(source.slice(cursor, keyEnd + 1)) as string;
    } catch {
      break;
    }
    cursor = keyEnd + 1;
    while (cursor < source.length && /\s/u.test(source[cursor] ?? '')) cursor += 1;
    if (source[cursor] !== ':') break;
    cursor += 1;
    while (cursor < source.length && /\s/u.test(source[cursor] ?? '')) cursor += 1;
    if (source[cursor] !== '"') break;
    const valueStart = cursor;
    const valueEnd = jsonStringEnd(source, valueStart);
    if (valueEnd < 0) {
      return {
        values,
        partial: {
          key,
          value: decodePartialJsonString(source.slice(valueStart + 1)),
        },
      };
    }
    try {
      values[key] = JSON.parse(source.slice(valueStart, valueEnd + 1)) as string;
    } catch {
      break;
    }
    cursor = valueEnd + 1;
  }
  return { values };
}

function artifactFiles(artifact: ArtifactDocument): ArtifactFileInput[] {
  return artifact.files?.length
    ? artifact.files.map(({ path, content }) => ({ path, content }))
    : [{
      path: artifact.entryPath ?? 'index.html',
      content: artifact.source,
    }];
}

function deletePaths(value: string | undefined): Set<string> {
  if (!value) return new Set();
  return new Set(value.split(/\r?\n/u).map((path) => {
    try {
      return normalizeArtifactPath(path);
    } catch {
      return '';
    }
  }).filter(Boolean));
}

function draftFile(path: string, content: string, partial: boolean): ArtifactDraftFile | null {
  try {
    const normalizedPath = normalizeArtifactPath(path);
    return {
      path: normalizedPath,
      content,
      mime: inferArtifactMime(normalizedPath),
      bytes: new TextEncoder().encode(content).byteLength,
      partial,
    };
  } catch {
    return null;
  }
}

function streamedFilePairs(parsed: StreamingJsonStrings): ArtifactDraftFile[] {
  const files: ArtifactDraftFile[] = [];
  for (let slot = 1; slot <= MAX_STREAMED_FILE_SLOTS; slot += 1) {
    const path = parsed.values[`file_${slot}_path`];
    if (!path) continue;
    const contentKey = `file_${slot}_content`;
    const completeContent = parsed.values[contentKey];
    const partialContent = parsed.partial?.key === contentKey ? parsed.partial.value : undefined;
    const file = completeContent !== undefined
      ? draftFile(path, completeContent, false)
      : partialContent !== undefined
        ? draftFile(path, partialContent, true)
        : null;
    if (file) files.push(file);
  }
  return files;
}

function mergeFiles(
  existing: ArtifactFileInput[],
  streamed: ArtifactDraftFile[],
  removed: Set<string>,
): ArtifactDraftFile[] {
  const merged = new Map<string, ArtifactDraftFile>();
  for (const file of existing) {
    const next = draftFile(file.path, file.content, false);
    if (next && !removed.has(next.path)) merged.set(next.path, next);
  }
  for (const file of streamed) {
    if (!removed.has(file.path)) merged.set(file.path, file);
  }
  return [...merged.values()];
}

function previewFiles(
  files: ArtifactDraftFile[],
  existing: ArtifactFileInput[],
): ArtifactFileInput[] {
  const stable = new Map(existing.map((file) => [file.path, file]));
  for (const file of files) {
    if (file.partial && SCRIPT_OR_DATA_FILE.test(file.path)) continue;
    stable.set(file.path, { path: file.path, content: file.content });
  }
  return [...stable.values()];
}

export function buildStreamingArtifactDraft(
  toolName: string,
  streamedArguments: string,
  draftId: string,
  artifacts: ReadonlyMap<string, ArtifactDocument> = new Map(),
): ArtifactDraft | null {
  if (!['artifact_deploy', 'artifact_write', 'html_artifact'].includes(toolName)) return null;
  const parsed = extractStreamingJsonStrings(streamedArguments);
  const streamedFiles = toolName === 'html_artifact'
    ? (() => {
      const html = parsed.values.html ?? (parsed.partial?.key === 'html' ? parsed.partial.value : undefined);
      const file = html === undefined ? null : draftFile('index.html', html, parsed.values.html === undefined);
      return file ? [file] : [];
    })()
    : streamedFilePairs(parsed);
  const sourceArtifact = toolName === 'artifact_write' && parsed.values.artifact_id
    ? artifacts.get(parsed.values.artifact_id)
    : undefined;
  const existingFiles = sourceArtifact ? artifactFiles(sourceArtifact) : [];
  const removed = deletePaths(parsed.values.delete_paths);
  const files = mergeFiles(existingFiles, streamedFiles, removed);
  if (files.length === 0) return null;

  const title = parsed.values.title
    ?? sourceArtifact?.title
    ?? (toolName === 'artifact_write' ? 'Updating artifact' : 'Untitled artifact');
  const entryPath = parsed.values.entry_path ?? sourceArtifact?.entryPath ?? 'index.html';
  let preview: ArtifactDocument | undefined;
  try {
    preview = prepareArtifactDeployment(
      previewFiles(files, existingFiles).filter((file) => !removed.has(file.path)),
      title,
      entryPath,
      draftId,
    );
  } catch {
    // A partial file tree is still useful before its HTML entry becomes renderable.
  }
  return {
    id: draftId,
    title,
    entryPath,
    files,
    ...(preview ? { preview } : {}),
    updatedAt: Date.now(),
  };
}