page-edit-service.ts

import crypto from "node:crypto"
import path from "node:path"
import { createBookStorage } from "@adt/storage"
import { createLLMModel, createPromptEngine } from "@adt/llm"
import type { LLMModel } from "@adt/llm"
import { renderPage, buildRenderStrategyResolver, createTemplateEngine, loadBookConfig, createScreenshotRenderer, runVisualReviewLoop, DEFAULT_VISUAL_REVIEW_MODEL_ID, structurePage, buildStructureConfig, renderSectionThumbnail } from "@adt/pipeline"
import type { VisualRefinementDeps, ScreenshotRenderer } from "@adt/pipeline"
import { PageSectioningOutput, WebRenderingOutput, webRenderingLLMSchema, ImageClassificationOutput, type ContentNodeData } from "@adt/types"
import { loadStyleguideContent } from "./styleguide.js"

export interface ReRenderOptions {
  label: string
  pageId: string
  sectionIndex?: number
  /** Optional user prompt/instructions to guide the LLM during re-render */
  prompt?: string
  booksDir: string
  promptsDir: string
  webAssetsDir?: string
  configPath?: string
  apiKey: string
}

export interface ReRenderResult {
  version: number
  rendering: unknown
}

export interface ReStructurePageOptions {
  label: string
  pageId: string
  booksDir: string
  promptsDir: string
  configPath?: string
  apiKey: string
}

export interface ReStructurePageResult {
  version: number
  pageStructuring: unknown
}

export interface AiEditSectionOptions {
  label: string
  pageId: string
  sectionIndex: number
  instruction: string
  /** Optional: current HTML from the frontend (for successive edits on unsaved changes) */
  currentHtml?: string
  booksDir: string
  promptsDir: string
  webAssetsDir?: string
  configPath?: string
  apiKey: string
}

export interface AiEditSectionResult {
  html: string
  reasoning: string
}

export async function reRenderPage(
  options: ReRenderOptions
): Promise<ReRenderResult> {
  const { label, pageId, sectionIndex, prompt, booksDir, promptsDir, webAssetsDir, configPath, apiKey } = options

  // Set API key
  const previousKey = process.env.OPENAI_API_KEY
  process.env.OPENAI_API_KEY = apiKey

  const storage = createBookStorage(label, booksDir)
  let visualRefinement: VisualRefinementDeps | undefined
  let screenshotRenderer: ScreenshotRenderer | undefined

  try {
    // Read latest pipeline data
    const sectionRow = storage.getLatestNodeData("page-sectioning", pageId)

    if (!sectionRow) {
      throw new Error(
        "Page must have page-sectioning data before re-rendering"
      )
    }

    const sectioningParsed = PageSectioningOutput.safeParse(sectionRow.data)
    if (!sectioningParsed.success) {
      throw new Error("Invalid page-sectioning data")
    }
    const sectioning = sectioningParsed.data

    // Build image map: start with all page images, then add any additional images
    // referenced in section parts (e.g. from cross-page merges).
    const allImages = storage.getPageImages(pageId)
    const renderImages = new Map<string, { base64: string; width?: number; height?: number }>()
    for (const img of allImages) {
      renderImages.set(img.imageId, { base64: storage.getImageBase64(img.imageId), width: img.width, height: img.height })
    }
    // Add any images referenced in sections but not found on this page
    for (const section of sectioning.sections) {
      for (const part of section.parts) {
        if (part.type === "image" && !part.isPruned && !renderImages.has(part.imageId)) {
          const dims = storage.getImageDimensions(part.imageId)
          renderImages.set(part.imageId, { base64: storage.getImageBase64(part.imageId), width: dims?.width, height: dims?.height })
        } else if (part.type === "content_node" && !part.isPruned) {
          collectImageIdsFromNode(part.node, renderImages, storage)
        }
      }
    }

    // Load config and build render strategy resolver
    const config = loadBookConfig(label, booksDir, configPath)
    const resolveRenderConfig = buildRenderStrategyResolver(config)

    const styleguideContent = loadStyleguideContent(config.styleguide, configPath)

    // Create LLM model resolver (model-specific, cached)
    const cacheDir = path.join(path.resolve(booksDir), label, ".cache")
    const bookPromptsDir = path.join(path.resolve(booksDir), label, "prompts")
    const promptEngine = createPromptEngine([bookPromptsDir, promptsDir])
    const templatesDir = path.join(path.dirname(promptsDir), "templates")
    const templateEngine = createTemplateEngine(templatesDir)
    const renderModels = new Map<string, LLMModel>()
    const resolveRenderModel = (modelId: string): LLMModel => {
      const existing = renderModels.get(modelId)
      if (existing) return existing
      const model = createLLMModel({
        modelId,
        cacheDir,
        promptEngine,
        onLog: (entry) => storage.appendLlmLog(entry),
      })
      renderModels.set(modelId, model)
      return model
    }

    // Get page image
    const pageImageBase64 = storage.getPageImageBase64(pageId)

    if (sectionIndex !== undefined && (sectionIndex < 0 || sectionIndex >= sectioning.sections.length)) {
      throw new Error(`Section index ${sectionIndex} out of range`)
    }

    // Set up shared screenshot renderer for visual refinement + thumbnails
    if (webAssetsDir) {
      const hasVisualRefinement = Object.values(config.render_strategies ?? {}).some(
        (s) => s.config?.visual_refinement?.enabled
      )
      screenshotRenderer = await createScreenshotRenderer()
      if (hasVisualRefinement) {
        visualRefinement = {
          screenshotRenderer,
          webAssetsDir,
          llmModel: resolveRenderModel(DEFAULT_VISUAL_REVIEW_MODEL_ID),
          storeScreenshot: (base64: string) => {
            const hash = crypto.createHash("sha256").update(base64).digest("hex").slice(0, 16)
            storage.putDebugImage(hash, Buffer.from(base64, "base64"))
          },
        }
      }
    }

    const captureThumbnails = async (rendering: WebRenderingOutput): Promise<void> => {
      if (!screenshotRenderer || !webAssetsDir) return
      const thumbImages = new Map<string, { base64: string }>()
      for (const [id, img] of renderImages) thumbImages.set(id, { base64: img.base64 })
      for (const section of rendering.sections) {
        try {
          const buffer = await renderSectionThumbnail({
            section,
            label,
            images: thumbImages,
            webAssetsDir,
            screenshotRenderer,
          })
          storage.putSectionThumbnail(pageId, section.sectionIndex, buffer)
        } catch (err) {
          console.error(
            `[page-edit] ${label}: thumbnail failed for ${pageId} sec${section.sectionIndex}: ${err instanceof Error ? err.message : String(err)}`
          )
        }
      }
    }

    // Render either a single section (preferred) or the full page.
    // For section re-render we force all other sections to pruned in-memory so
    // renderPage preserves the original sectionIndex while skipping extra LLM calls.
    const sectioningForRender = sectionIndex === undefined
      ? sectioning
      : {
          ...sectioning,
          sections: sectioning.sections.map((section, idx) =>
            idx === sectionIndex ? section : { ...section, isPruned: true }
          ),
        }

    const renderResult = await renderPage(
      {
        label,
        pageId,
        pageImageBase64,
        sectioning: sectioningForRender,
        images: renderImages,
        styleguide: styleguideContent,
        userPrompt: prompt,
      },
      resolveRenderConfig,
      resolveRenderModel,
      templateEngine,
      visualRefinement,
    )

    if (sectionIndex === undefined) {
      const version = storage.putNodeData("web-rendering", pageId, renderResult)
      await captureThumbnails(renderResult)
      return { version, rendering: renderResult }
    }

    // Merge the newly rendered section back into existing rendering, preserving
    // other sections as-is.
    const existingRenderingRow = storage.getLatestNodeData("web-rendering", pageId)
    const existingRenderingParsed = existingRenderingRow
      ? WebRenderingOutput.safeParse(existingRenderingRow.data)
      : null
    if (existingRenderingRow && !existingRenderingParsed?.success) {
      throw new Error("Invalid web-rendering data")
    }
    const existingSections = existingRenderingParsed?.success
      ? existingRenderingParsed.data.sections
      : []
    const withoutTarget = existingSections.filter((s) => s.sectionIndex !== sectionIndex)
    const newTarget = renderResult.sections.find((s) => s.sectionIndex === sectionIndex)
    const mergedSections = newTarget
      ? [...withoutTarget, newTarget].sort((a, b) => a.sectionIndex - b.sectionIndex)
      : withoutTarget.sort((a, b) => a.sectionIndex - b.sectionIndex)
    const mergedRendering = { sections: mergedSections }

    const version = storage.putNodeData("web-rendering", pageId, mergedRendering)
    await captureThumbnails(mergedRendering)
    return { version, rendering: mergedRendering }
  } finally {
    if (screenshotRenderer) {
      await screenshotRenderer.close()
    }
    storage.clearNodesByType(["image-captioning", "text-catalog", "text-catalog-translation", "tts", "tts-timestamps"])
    storage.clearStepRuns(["image-captioning", "text-catalog", "catalog-translation", "tts"])
    storage.close()
    // Restore previous key
    if (previousKey !== undefined) {
      process.env.OPENAI_API_KEY = previousKey
    } else {
      delete process.env.OPENAI_API_KEY
    }
  }
}

/**
 * Use LLM to edit a single section's HTML based on a natural language instruction.
 * Returns the edited HTML and reasoning without saving — the frontend previews first.
 */
export async function aiEditSection(
  options: AiEditSectionOptions
): Promise<AiEditSectionResult> {
  const { label, pageId, sectionIndex, instruction, currentHtml: providedHtml, booksDir, promptsDir, webAssetsDir, configPath, apiKey } = options

  const previousKey = process.env.OPENAI_API_KEY
  process.env.OPENAI_API_KEY = apiKey

  const storage = createBookStorage(label, booksDir)

  try {
    // Use provided HTML (from frontend pending state) or read from DB
    let currentHtml: string
    if (providedHtml) {
      currentHtml = providedHtml
    } else {
      const renderingRow = storage.getLatestNodeData("web-rendering", pageId)
      if (!renderingRow) {
        throw new Error("Page must have web-rendering data before AI editing")
      }
      const renderingParsed = WebRenderingOutput.safeParse(renderingRow.data)
      if (!renderingParsed.success) {
        throw new Error("Invalid web-rendering data")
      }
      const section = renderingParsed.data.sections.find((s) => s.sectionIndex === sectionIndex)
      if (!section) {
        throw new Error(`Section ${sectionIndex} not found in rendering`)
      }
      currentHtml = section.html
    }

    // Load config to get model ID for editing
    const config = loadBookConfig(label, booksDir, configPath)
    const modelId = (config as Record<string, unknown>).page_sectioning
      ? ((config as Record<string, unknown>).page_sectioning as Record<string, unknown>).model as string
      : "openai:gpt-4o"

    // Build LLM model
    const cacheDir = path.join(path.resolve(booksDir), label, ".cache")
    const bookPromptsDir = path.join(path.resolve(booksDir), label, "prompts")
    const promptEngine = createPromptEngine([bookPromptsDir, promptsDir])
    const model = createLLMModel({
      modelId,
      cacheDir,
      promptEngine,
      onLog: (entry) => storage.appendLlmLog(entry),
    })

    // Extract existing data-ids and img tags for validation
    const dataIdRegex = /data-id="([^"]+)"/g
    const existingIds = new Set<string>()
    let match
    while ((match = dataIdRegex.exec(currentHtml)) !== null) {
      existingIds.add(match[1])
    }

    // Extract img tags with their data-ids and srcs
    const imgTagRegex = /<img\s[^>]*data-id="([^"]+)"[^>]*src="([^"]+)"[^>]*>/g
    const existingImgs = new Map<string, string>() // data-id → src
    while ((match = imgTagRegex.exec(currentHtml)) !== null) {
      existingImgs.set(match[1], match[2])
    }
    // Also catch imgs where src comes before data-id
    const imgTagRegex2 = /<img\s[^>]*src="([^"]+)"[^>]*data-id="([^"]+)"[^>]*>/g
    while ((match = imgTagRegex2.exec(currentHtml)) !== null) {
      if (!existingImgs.has(match[2])) {
        existingImgs.set(match[2], match[1])
      }
    }

    // Load the original page image so the LLM can see the intended layout
    let pageImageBase64: string | undefined
    try {
      pageImageBase64 = storage.getPageImageBase64(pageId)
    } catch {
      // Page image not available — proceed without it
    }

    const validateEditedHtml = (rawHtml: string) => {
      const cleanedHtml = rawHtml
        .replace(/^```(?:html)?\s*\n?/i, "")
        .replace(/\n?```\s*$/, "")
      const errors: string[] = []

      if (!cleanedHtml.includes("<section")) {
        errors.push("Result must contain a <section> element")
      }

      for (const id of existingIds) {
        if (!cleanedHtml.includes(`data-id="${id}"`)) {
          errors.push(`Missing data-id="${id}" in result`)
        }
      }

      for (const [imgDataId, imgSrc] of existingImgs) {
        const escaped = imgDataId.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")
        const imgCheck = new RegExp(`<img\\s[^>]*data-id="${escaped}"[^>]*>`)
        if (!imgCheck.test(cleanedHtml)) {
          const imgCheck2 = new RegExp(`<img\\s[^>]*src="[^"]*"[^>]*data-id="${escaped}"[^>]*>`)
          if (!imgCheck2.test(cleanedHtml)) {
            errors.push(`Image data-id="${imgDataId}" must remain an <img> tag`)
          }
        }
        if (!cleanedHtml.includes(imgSrc)) {
          errors.push(`Image src="${imgSrc}" was removed or changed`)
        }
      }

      return { valid: errors.length === 0, errors, cleanedHtml }
    }

    const result = await model.generateObject<{ reasoning: string; content: string }>({
      schema: webRenderingLLMSchema,
      prompt: "html_edit",
      context: { current_html: currentHtml, instruction, page_image_base64: pageImageBase64 },
      validate: (obj) => {
        const r = obj as { content: string }
        const check = validateEditedHtml(r.content)
        return { valid: check.valid, errors: check.errors }
      },
      maxRetries: 3,
      log: { taskType: "web-rendering", pageId, promptName: "html_edit" },
    })

    let html = validateEditedHtml(result.object.content).cleanedHtml

    // Visual refinement loop — screenshot the edited HTML and verify
    if (webAssetsDir) {
      // Find first render strategy with visual refinement enabled
      const vrStrategyConfig = Object.values(config.render_strategies ?? {})
        .find((s) => s.config?.visual_refinement?.enabled)?.config?.visual_refinement
      if (vrStrategyConfig?.enabled) {
        const maxIterations = vrStrategyConfig.max_iterations ?? 3
        const vrTimeout = vrStrategyConfig.timeout ?? 120
        const vrTemperature = vrStrategyConfig.temperature

        const reviewModel = createLLMModel({
          modelId: DEFAULT_VISUAL_REVIEW_MODEL_ID,
          cacheDir,
          promptEngine,
          onLog: (entry) => storage.appendLlmLog(entry),
        })

        // Build image map from data-ids in the HTML for screenshot rendering
        const imagesForScreenshot = new Map<string, { base64: string }>()
        for (const [imgDataId] of existingImgs) {
          try {
            imagesForScreenshot.set(imgDataId, { base64: storage.getImageBase64(imgDataId) })
          } catch {
            // Image not found in storage — skip (will show broken in screenshot)
          }
        }

        const screenshotRenderer = await createScreenshotRenderer()
        try {
          const review = await runVisualReviewLoop({
            initialHtml: html,
            label,
            pageId,
            images: imagesForScreenshot,
            deps: {
              llmModel: reviewModel,
              screenshotRenderer,
              webAssetsDir,
              storeScreenshot: (base64) => {
                const hash = crypto.createHash("sha256").update(base64).digest("hex").slice(0, 16)
                storage.putDebugImage(hash, Buffer.from(base64, "base64"))
              },
            },
            promptName: "visual_review_edit",
            maxIterations,
            timeoutMs: vrTimeout * 1000,
            temperature: vrTemperature,
            pageImageBase64,
            promptContext: { instruction },
            firstIterationScreenshotsText: "\nHere are screenshots of the edited HTML at three viewport sizes:\n",
            nextIterationScreenshotsText: "Here are the updated screenshots after your revision:\n",
            trailingContextText: `Edit instruction: ${instruction}`,
            validateHtml: (candidateHtml) => {
              const check = validateEditedHtml(candidateHtml)
              return { valid: check.valid, errors: check.errors, cleanedHtml: check.cleanedHtml }
            },
          })
          html = review.html
        } finally {
          await screenshotRenderer.close()
        }
      }
    }

    return { html, reasoning: result.object.reasoning }
  } finally {
    storage.close()
    if (previousKey !== undefined) {
      process.env.OPENAI_API_KEY = previousKey
    } else {
      delete process.env.OPENAI_API_KEY
    }
  }
}

export async function reStructurePage(
  options: ReStructurePageOptions
): Promise<ReStructurePageResult> {
  const { label, pageId, booksDir, promptsDir, configPath, apiKey } = options

  const previousKey = process.env.OPENAI_API_KEY
  process.env.OPENAI_API_KEY = apiKey

  const storage = createBookStorage(label, booksDir)
  try {
    const pages = storage.getPages()
    const page = pages.find((p) => p.pageId === pageId)
    if (!page) throw new Error(`Page not found: ${pageId}`)

    // Load config and build structure config
    const config = loadBookConfig(label, booksDir, configPath)
    const structureConfig = buildStructureConfig(config)

    // Create LLM model
    const cacheDir = path.join(path.resolve(booksDir), label, ".cache")
    const bookPromptsDir = path.join(path.resolve(booksDir), label, "prompts")
    const promptEngine = createPromptEngine([bookPromptsDir, promptsDir])
    const llmModel = createLLMModel({
      modelId: structureConfig.modelId,
      cacheDir,
      promptEngine,
      onLog: (entry) => storage.appendLlmLog(entry),
    })

    // Get page image
    const imageBase64 = storage.getPageImageBase64(pageId)

    // Get unpruned images from current image classification
    const imageClassRow = storage.getLatestNodeData("image-filtering", pageId)
    const allPageImages = storage.getPageImages(pageId)
    let structureImages: Array<{ imageId: string; imageBase64: string }> = []

    if (imageClassRow) {
      const parsed = ImageClassificationOutput.safeParse(imageClassRow.data)
      if (parsed.success) {
        const unprunedIds = new Set(
          parsed.data.images.filter((img) => !img.isPruned).map((img) => img.imageId)
        )
        structureImages = allPageImages
          .filter((img) => unprunedIds.has(img.imageId))
          .map((img) => ({
            imageId: img.imageId,
            imageBase64: storage.getImageBase64(img.imageId),
          }))
      }
    }

    // Run page structuring
    const result = await structurePage(
      {
        pageId: page.pageId,
        pageNumber: page.pageNumber,
        text: page.text,
        imageBase64,
        images: structureImages,
      },
      structureConfig,
      llmModel
    )

    const version = storage.putNodeData("page-structuring", pageId, result)
    return { version, pageStructuring: result }
  } finally {
    storage.close()
    if (previousKey !== undefined) {
      process.env.OPENAI_API_KEY = previousKey
    } else {
      delete process.env.OPENAI_API_KEY
    }
  }
}

/**
 * Recursively collect image IDs from a content node tree and add them to the render images map.
 */
function collectImageIdsFromNode(
  node: ContentNodeData,
  renderImages: Map<string, { base64: string; width?: number; height?: number }>,
  storage: { getImageBase64: (id: string) => string; getImageDimensions: (id: string) => { width: number; height: number } | null }
): void {
  if (node.isPruned) return
  const ensureLoaded = (imageId: string) => {
    if (renderImages.has(imageId)) return
    try {
      const dims = storage.getImageDimensions(imageId)
      renderImages.set(imageId, { base64: storage.getImageBase64(imageId), width: dims?.width, height: dims?.height })
    } catch {
      // Image not found in storage — skip
    }
  }
  if (node.imageId) ensureLoaded(node.imageId)
  if (node.backgroundImageId) ensureLoaded(node.backgroundImageId)
  if (node.children) {
    for (const child of node.children) {
      collectImageIdsFromNode(child, renderImages, storage)
    }
  }
}