users/lib/documentation/enhanced-screenshot-service.ts

import { screenshotService, type ScreenshotOptions, type ScreenshotResult } from "./screenshot-service"

export interface EnhancedScreenshotOptions extends ScreenshotOptions {
  waitForImages?: boolean
  waitForFonts?: boolean
  removeElements?: string[]
  addWatermark?: boolean
  watermarkText?: string
}

class EnhancedScreenshotService {
  async captureViewport(options: EnhancedScreenshotOptions = {}): Promise<ScreenshotResult> {
    try {
      // 等待页面完全加载
      await this.waitForPageLoad(options)

      // 移除不需要的元素
      const removedElements = this.removeElements(options.removeElements || [])

      try {
        // 调用基础截图服务
        const result = await screenshotService.captureViewport(options)

        // 添加水印
        if (options.addWatermark && result.success && result.dataUrl) {
          result.dataUrl = await this.addWatermark(result.dataUrl, options.watermarkText)
        }

        return result
      } finally {
        // 恢复移除的元素
        this.restoreElements(removedElements)
      }
    } catch (error) {
      console.error("增强截图失败:", error)
      return {
        success: false,
        error: error instanceof Error ? error.message : "未知错误",
      }
    }
  }

  async captureElement(element: HTMLElement, options: EnhancedScreenshotOptions = {}): Promise<ScreenshotResult> {
    try {
      // 等待页面完全加载
      await this.waitForPageLoad(options)

      // 调用基础截图服务
      const result = await screenshotService.captureElement(element, options)

      // 添加水印
      if (options.addWatermark && result.success && result.dataUrl) {
        result.dataUrl = await this.addWatermark(result.dataUrl, options.watermarkText)
      }

      return result
    } catch (error) {
      console.error("增强元素截图失败:", error)
      return {
        success: false,
        error: error instanceof Error ? error.message : "未知错误",
      }
    }
  }

  private async waitForPageLoad(options: EnhancedScreenshotOptions): Promise<void> {
    // 等待图片加载
    if (options.waitForImages) {
      await this.waitForImages()
    }

    // 等待字体加载
    if (options.waitForFonts) {
      await this.waitForFonts()
    }

    // 额外等待时间确保渲染完成
    await new Promise((resolve) => setTimeout(resolve, 1000))
  }

  private async waitForImages(): Promise<void> {
    const images = Array.from(document.images)
    const promises = images.map((img) => {
      if (img.complete) return Promise.resolve()

      return new Promise((resolve, reject) => {
        img.onload = resolve
        img.onerror = resolve // 即使图片加载失败也继续
        setTimeout(resolve, 5000) // 5秒超时
      })
    })

    await Promise.all(promises)
  }

  private async waitForFonts(): Promise<void> {
    if ("fonts" in document) {
      try {
        await document.fonts.ready
      } catch (error) {
        console.warn("字体加载等待失败:", error)
      }
    }
  }

  private removeElements(selectors: string[]): Array<{ element: Element; parent: Node; nextSibling: Node | null }> {
    const removedElements: Array<{ element: Element; parent: Node; nextSibling: Node | null }> = []

    selectors.forEach((selector) => {
      const elements = document.querySelectorAll(selector)
      elements.forEach((element) => {
        if (element.parentNode) {
          removedElements.push({
            element,
            parent: element.parentNode,
            nextSibling: element.nextSibling,
          })
          element.parentNode.removeChild(element)
        }
      })
    })

    return removedElements
  }

  private restoreElements(removedElements: Array<{ element: Element; parent: Node; nextSibling: Node | null }>): void {
    removedElements.forEach(({ element, parent, nextSibling }) => {
      if (nextSibling) {
        parent.insertBefore(element, nextSibling)
      } else {
        parent.appendChild(element)
      }
    })
  }

  private async addWatermark(dataUrl: string, watermarkText?: string): Promise<string> {
    return new Promise((resolve) => {
      const canvas = document.createElement("canvas")
      const ctx = canvas.getContext("2d")
      const img = new Image()

      img.onload = () => {
        canvas.width = img.width
        canvas.height = img.height

        // 绘制原图
        ctx?.drawImage(img, 0, 0)

        if (ctx && watermarkText) {
          // 添加水印
          ctx.font = "16px Arial"
          ctx.fillStyle = "rgba(0, 0, 0, 0.5)"
          ctx.textAlign = "right"
          ctx.fillText(
            watermarkText || `截图时间: ${new Date().toLocaleString()}`,
            canvas.width - 20,
            canvas.height - 20,
          )
        }

        resolve(canvas.toDataURL("image/png"))
      }

      img.src = dataUrl
    })
  }
}

export const enhancedScreenshotService = new EnhancedScreenshotService()
export type { EnhancedScreenshotOptions }