#!/usr/bin/env swift

import AppKit
import Foundation
import ImageIO
import UniformTypeIdentifiers
import WebKit

enum ExportError: Error, CustomStringConvertible {
  case invalidArguments
  case snapshotFailed
  case pngEncodingFailed(String)
  case clockRegionMissing

  var description: String {
    switch self {
    case .invalidArguments:
      return "用法：export-kindle-background.swift <url> <png-output> <region-output>"
    case .snapshotFailed:
      return "网页截图失败"
    case let .pngEncodingFailed(path):
      return "PNG 编码失败：\(path)"
    case .clockRegionMissing:
      return "页面里没有找到 data-clock-region 标记"
    }
  }
}

final class SnapshotExporter: NSObject, WKNavigationDelegate {
  private enum ExportOrientation {
    case portrait
    case landscape
  }

  private struct ExportLayout {
    let orientation: ExportOrientation
    let viewportSize: CGSize
    let outputSize: CGSize
  }

  private let url: URL
  private let pngOutputURL: URL
  private let regionOutputURL: URL
  private let completion: (Result<Void, Error>) -> Void
  // 网页渲染尺寸和 Kindle framebuffer 输出尺寸不是一回事。
  // landscape 需要先按 1448x1072 渲染，再旋转成 1072x1448 输出给设备。
  private let layout: ExportLayout

  private lazy var window: NSWindow = {
    let window = NSWindow(
      contentRect: CGRect(origin: .zero, size: layout.viewportSize),
      styleMask: [.borderless],
      backing: .buffered,
      defer: false
    )
    window.isReleasedWhenClosed = false
    window.backgroundColor = .white
    return window
  }()

  private lazy var webView: WKWebView = {
    let config = WKWebViewConfiguration()
    let view = WKWebView(frame: CGRect(origin: .zero, size: layout.viewportSize), configuration: config)
    view.navigationDelegate = self
    return view
  }()

  init(url: URL, pngOutputURL: URL, regionOutputURL: URL, completion: @escaping (Result<Void, Error>) -> Void) {
    self.url = url
    self.pngOutputURL = pngOutputURL
    self.regionOutputURL = regionOutputURL
    self.completion = completion
    self.layout = SnapshotExporter.resolveLayout(url: url)
    super.init()
  }

  private static func resolveLayout(url: URL) -> ExportLayout {
    let components = URLComponents(url: url, resolvingAgainstBaseURL: false)
    let orientation = components?.queryItems?.first(where: { $0.name == "orientation" })?.value

    if orientation == "landscape" {
      return ExportLayout(
        orientation: .landscape,
        viewportSize: CGSize(width: 1448, height: 1072),
        outputSize: CGSize(width: 1072, height: 1448)
      )
    }

    return ExportLayout(
      orientation: .portrait,
      viewportSize: CGSize(width: 1072, height: 1448),
      outputSize: CGSize(width: 1072, height: 1448)
    )
  }

  func start() {
    window.contentView = webView
    window.orderBack(nil)
    webView.load(URLRequest(url: url))
  }

  func webView(_ webView: WKWebView, didFinish _: WKNavigation!) {
    // 等页面数据与图片稳定下来，避免只截到骨架。
    DispatchQueue.main.asyncAfter(deadline: .now() + 4) {
      self.capture()
    }
  }

  private func capture() {
    let regionScript = """
      (() => {
        const node = document.querySelector('[data-clock-region="true"]');
        if (!node) return null;
        const dashboard = document.querySelector('.dashboard-frame');
        const rect = node.getBoundingClientRect();
        const frameRect = dashboard ? dashboard.getBoundingClientRect() : { left: 0, top: 0, width: window.innerWidth, height: window.innerHeight };
        return {
          clock: {
            x: Math.round(rect.left),
            y: Math.round(rect.top),
            width: Math.round(rect.width),
            height: Math.round(rect.height)
          },
          frame: {
            x: Math.round(frameRect.left),
            y: Math.round(frameRect.top),
            width: Math.round(frameRect.width),
            height: Math.round(frameRect.height)
          }
        };
      })();
      """

    webView.evaluateJavaScript(regionScript) { value, error in
      if let error {
        self.finish(.failure(error))
        return
      }

      guard
        let result = value as? [String: Any],
        let region = result["clock"] as? [String: NSNumber],
        let frame = result["frame"] as? [String: NSNumber]
      else {
        self.finish(.failure(ExportError.clockRegionMissing))
        return
      }

      let snapshotConfig = WKSnapshotConfiguration()
      snapshotConfig.rect = CGRect(origin: .zero, size: self.layout.viewportSize)
      snapshotConfig.snapshotWidth = NSNumber(value: Float(self.layout.viewportSize.width))

      self.webView.takeSnapshot(with: snapshotConfig) { image, snapshotError in
        if let snapshotError {
          self.finish(.failure(snapshotError))
          return
        }

        guard let image else {
          self.finish(.failure(ExportError.snapshotFailed))
          return
        }

        do {
          try self.savePNG(image: image, to: self.pngOutputURL, frameOffset: frame)
          try self.saveRegion(region: region, frameOffset: frame)
          self.finish(.success(()))
        } catch {
          self.finish(.failure(error))
        }
      }
    }
  }

  private func savePNG(image: NSImage, to url: URL, frameOffset: [String: NSNumber]) throws {
    let normalizedImage = NSImage(size: NSSize(width: layout.viewportSize.width, height: layout.viewportSize.height))
    normalizedImage.lockFocus()
    image.draw(in: NSRect(origin: .zero, size: layout.viewportSize))
    normalizedImage.unlockFocus()

    guard let sourceCGImage = normalizedImage.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
      throw ExportError.pngEncodingFailed(url.path)
    }

    let width = Int(layout.outputSize.width)
    let height = Int(layout.outputSize.height)
    let colorSpace = CGColorSpaceCreateDeviceGray()

    // 输出 8-bit 灰度 PNG，但页面本身仍按纯白底和纯黑字设计，避免额外灰阶装饰。
    guard let context = CGContext(
      data: nil,
      width: width,
      height: height,
      bitsPerComponent: 8,
      bytesPerRow: 0,
      space: colorSpace,
      bitmapInfo: CGImageAlphaInfo.none.rawValue
    ) else {
      throw ExportError.pngEncodingFailed(url.path)
    }

    context.setFillColor(gray: 1, alpha: 1)
    context.fill(CGRect(x: 0, y: 0, width: width, height: height))
    context.interpolationQuality = .high
    let offsetX = CGFloat(frameOffset["x"]?.doubleValue ?? 0)
    let offsetY = CGFloat(frameOffset["y"]?.doubleValue ?? 0)

    // Kindle framebuffer 固定是纵向；landscape 导出时这里把横向截图旋转成设备图。
    switch layout.orientation {
    case .portrait:
      context.draw(
        sourceCGImage,
        in: CGRect(
          x: -offsetX,
          y: -offsetY,
          width: layout.viewportSize.width,
          height: layout.viewportSize.height
        )
      )
    case .landscape:
      context.saveGState()
      // 横向网页需要先顺时针旋转 90 度写入纵向 framebuffer，
      // 这样 Kindle 在 logo_right 摆放时，最终看到的页面方向才是正的。
      context.translateBy(x: 0, y: layout.outputSize.height)
      context.rotate(by: -.pi / 2)
      context.draw(
        sourceCGImage,
        in: CGRect(
          x: -offsetX,
          y: -offsetY,
          width: layout.viewportSize.width,
          height: layout.viewportSize.height
        )
      )
      context.restoreGState()
    }

    guard let grayscaleImage = context.makeImage() else {
      throw ExportError.pngEncodingFailed(url.path)
    }

    try FileManager.default.createDirectory(at: url.deletingLastPathComponent(), withIntermediateDirectories: true)

    guard let destination = CGImageDestinationCreateWithURL(
      url as CFURL,
      UTType.png.identifier as CFString,
      1,
      nil
    ) else {
      throw ExportError.pngEncodingFailed(url.path)
    }

    CGImageDestinationAddImage(destination, grayscaleImage, nil)

    guard CGImageDestinationFinalize(destination) else {
      throw ExportError.pngEncodingFailed(url.path)
    }
  }

  private func saveRegion(region: [String: NSNumber], frameOffset: [String: NSNumber]) throws {
    let transformedRegion = transformRegion(region: region, frameOffset: frameOffset)
    let jsonObject: [String: Int] = [
      "x": transformedRegion["x"]?.intValue ?? 0,
      "y": transformedRegion["y"]?.intValue ?? 0,
      "width": transformedRegion["width"]?.intValue ?? 0,
      "height": transformedRegion["height"]?.intValue ?? 0,
    ]
    let data = try JSONSerialization.data(withJSONObject: jsonObject, options: [.prettyPrinted, .sortedKeys])
    try FileManager.default.createDirectory(at: regionOutputURL.deletingLastPathComponent(), withIntermediateDirectories: true)
    try data.write(to: regionOutputURL)
  }

  private func transformRegion(region: [String: NSNumber], frameOffset: [String: NSNumber]) -> [String: NSNumber] {
    let normalizedX = CGFloat(region["x"]?.doubleValue ?? 0) - CGFloat(frameOffset["x"]?.doubleValue ?? 0)
    let normalizedY = CGFloat(region["y"]?.doubleValue ?? 0) - CGFloat(frameOffset["y"]?.doubleValue ?? 0)
    let normalizedRegion: [String: NSNumber] = [
      "x": NSNumber(value: Int(round(normalizedX))),
      "y": NSNumber(value: Int(round(normalizedY))),
      "width": region["width"] ?? 0,
      "height": region["height"] ?? 0,
    ]

    guard layout.orientation == .landscape else {
      return normalizedRegion
    }

    let x = CGFloat(normalizedRegion["x"]?.doubleValue ?? 0)
    let y = CGFloat(normalizedRegion["y"]?.doubleValue ?? 0)
    let width = CGFloat(normalizedRegion["width"]?.doubleValue ?? 0)
    let height = CGFloat(normalizedRegion["height"]?.doubleValue ?? 0)

    // “logo 在右侧”对应设备图坐标系下，横向网页需要顺时针旋转 90 度写入 framebuffer。
    let transformedX = layout.viewportSize.height - (y + height)
    let transformedY = x

    return [
      "x": NSNumber(value: Int(round(transformedX))),
      "y": NSNumber(value: Int(round(transformedY))),
      "width": NSNumber(value: Int(round(height))),
      "height": NSNumber(value: Int(round(width))),
    ]
  }

  private func finish(_ result: Result<Void, Error>) {
    completion(result)
    NSApplication.shared.terminate(nil)
  }
}

guard CommandLine.arguments.count >= 4 else {
  fputs("\(ExportError.invalidArguments)\n", stderr)
  exit(1)
}

let url = URL(string: CommandLine.arguments[1])!
let pngOutputURL = URL(fileURLWithPath: CommandLine.arguments[2])
let regionOutputURL = URL(fileURLWithPath: CommandLine.arguments[3])

let app = NSApplication.shared
app.setActivationPolicy(.prohibited)

let exporter = SnapshotExporter(url: url, pngOutputURL: pngOutputURL, regionOutputURL: regionOutputURL) { result in
  switch result {
  case .success:
    print("Exported \(pngOutputURL.path)")
    print("Region saved to \(regionOutputURL.path)")
  case let .failure(error):
    fputs("\(error)\n", stderr)
    exit(1)
  }
}

exporter.start()
app.run()