AITuner.swift

import Foundation
import CoreGraphics
import AppKit
import WebKit

/// AI parameter tuning. The model does NOT draw curves (vtracer is better at that).
/// Its job is to look at the original image — and, when available, the current trace —
/// and suggest *parameter* adjustments that best capture the image's shapes.
///
/// One-shot flow (see `Tracer.autoTune`):
///   1. trace with current settings (already done; we have the SVG)
///   2. best-effort rasterize that SVG back to a bitmap
///   3. send original (+ trace) + current params to a vision model
///   4. apply the suggested settings and re-trace
enum AITuner {

    struct Suggestion: Decodable {
        var maxColors: Int?
        var filterSpeckle: Int?
        var cornerThreshold: Int?
        var colorMode: String?      // "color" | "bw"
        var rationale: String?
        // token usage, filled from the response (not the content JSON)
        var promptTokens: Int?
        var completionTokens: Int?
    }

    enum AITunerError: LocalizedError {
        case badResponse(String)
        case noContent
        case decode(String)

        var errorDescription: String? {
            switch self {
            case .badResponse(let m): return "AI request failed: \(m)"
            case .noContent:          return "AI returned an empty response."
            case .decode(let m):      return "Couldn't read the AI's reply: \(m)"
            }
        }
    }

    // MARK: - Encoding helpers

    static func pngData(from cg: CGImage) -> Data? {
        NSBitmapImageRep(cgImage: cg).representation(using: .png, properties: [:])
    }

    /// Clamp a size so the longest edge is <= maxDim (keeps the API payload small/cheap).
    static func clampSize(_ s: CGSize, maxDim: CGFloat = 768) -> CGSize {
        let m = max(s.width, s.height)
        guard m > maxDim, m > 0 else { return s }
        let scale = maxDim / m
        return CGSize(width: max(1, s.width * scale), height: max(1, s.height * scale))
    }

    // MARK: - The model call

    /// Ask the model for parameter suggestions. Returns .failure on any HTTP/parse error
    /// (caller keeps current settings).
    static func suggest(apiKey: String,
                        endpoint: String,
                        model: String,
                        originalPNG: Data,
                        tracePNG: Data?,
                        currentParams: [String: Any]) async -> Result<Suggestion, Error> {

        guard let url = URL(string: endpoint) else {
            return .failure(AITunerError.badResponse("invalid endpoint URL"))
        }

        let paramsJSON = (try? JSONSerialization.data(withJSONObject: currentParams))
            .flatMap { String(data: $0, encoding: .utf8) } ?? "{}"

        let system = """
        You are an expert in raster-to-vector tracing with vtracer. Given a source image \
        (and optionally its current trace), recommend settings that best reproduce the \
        image's SHAPES as clean vectors. Reply with ONLY a JSON object, no prose, with keys:
        maxColors (int 2-64; fewer for flat logos/line art, more for photos),
        filterSpeckle (int 0-20; higher removes noise but drops fine detail),
        cornerThreshold (int 0-180; higher = rounder, fewer sharp corners),
        colorMode ("color" or "bw"; use "bw" for line art / monochrome),
        rationale (one short sentence).
        """

        var userContent: [[String: Any]] = [
            ["type": "text",
             "text": "Current settings: \(paramsJSON). The FIRST image is the original."
                   + (tracePNG != nil ? " The SECOND image is the current trace — improve on it." : "")],
            ["type": "image_url",
             "image_url": ["url": "data:image/png;base64,\(originalPNG.base64EncodedString())"]],
        ]
        if let tracePNG {
            userContent.append([
                "type": "image_url",
                "image_url": ["url": "data:image/png;base64,\(tracePNG.base64EncodedString())"],
            ])
        }

        let body: [String: Any] = [
            "model": model,
            "messages": [
                ["role": "system", "content": system],
                ["role": "user", "content": userContent],
            ],
            "response_format": ["type": "json_object"],
            "max_completion_tokens": 600,
        ]

        var req = URLRequest(url: url)
        req.httpMethod = "POST"
        req.setValue("application/json", forHTTPHeaderField: "Content-Type")
        req.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        req.httpBody = try? JSONSerialization.data(withJSONObject: body)
        req.timeoutInterval = 60

        do {
            let (data, response) = try await URLSession.shared.data(for: req)
            guard let http = response as? HTTPURLResponse else {
                return .failure(AITunerError.badResponse("no HTTP response"))
            }
            guard (200..<300).contains(http.statusCode) else {
                let msg = String(data: data, encoding: .utf8) ?? "status \(http.statusCode)"
                return .failure(AITunerError.badResponse("HTTP \(http.statusCode): \(msg.prefix(300))"))
            }
            guard let content = extractContent(from: data) else {
                return .failure(AITunerError.noContent)
            }
            guard var suggestion = try? JSONDecoder().decode(Suggestion.self, from: Data(content.utf8)) else {
                return .failure(AITunerError.decode(String(content.prefix(200))))
            }
            if let usage = extractUsage(from: data) {
                suggestion.promptTokens = usage.prompt
                suggestion.completionTokens = usage.completion
            }
            return .success(suggestion)
        } catch {
            return .failure(error)
        }
    }

    /// Pull `usage.prompt_tokens` / `usage.completion_tokens` from the response.
    private static func extractUsage(from data: Data) -> (prompt: Int, completion: Int)? {
        guard let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let usage = obj["usage"] as? [String: Any] else { return nil }
        return (usage["prompt_tokens"] as? Int ?? 0, usage["completion_tokens"] as? Int ?? 0)
    }

    /// Pull `choices[0].message.content` out of an OpenAI-style chat completion response.
    private static func extractContent(from data: Data) -> String? {
        guard let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let choices = obj["choices"] as? [[String: Any]],
              let message = choices.first?["message"] as? [String: Any],
              let content = message["content"] as? String else {
            return nil
        }
        return content
    }
}

/// Renders an SVG string to PNG data via an offscreen WebKit view. Best-effort:
/// returns nil if rendering fails. Must be used on the main actor (WebKit requirement).
@MainActor
final class SVGRasterizer: NSObject, WKNavigationDelegate {
    private var webView: WKWebView?
    private var continuation: CheckedContinuation<Data?, Never>?
    private let size: CGSize

    init(size: CGSize) {
        self.size = size
        super.init()
    }

    func render(svg: String) async -> Data? {
        await withCheckedContinuation { cont in
            self.continuation = cont
            let wv = WKWebView(frame: CGRect(origin: .zero, size: size))
            wv.navigationDelegate = self
            self.webView = wv
            let w = Int(size.width), h = Int(size.height)
            let html = """
            <!doctype html><html><head><meta charset="utf-8">
            <style>html,body{margin:0;padding:0;background:#fff;}
            svg{width:\(w)px;height:\(h)px;}</style></head>
            <body>\(svg)</body></html>
            """
            wv.loadHTMLString(html, baseURL: nil)
        }
    }

    private func finish(_ data: Data?) {
        continuation?.resume(returning: data)
        continuation = nil
    }

    func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
        let config = WKSnapshotConfiguration()
        config.rect = CGRect(origin: .zero, size: size)
        webView.takeSnapshot(with: config) { [weak self] image, _ in
            guard let self else { return }
            guard let image,
                  let tiff = image.tiffRepresentation,
                  let rep = NSBitmapImageRep(data: tiff),
                  let png = rep.representation(using: .png, properties: [:]) else {
                self.finish(nil); return
            }
            self.finish(png)
        }
    }

    func webView(_ webView: WKWebView, didFail navigation: WKNavigation!, withError error: Error) {
        finish(nil)
    }

    func webView(_ webView: WKWebView, didFailProvisionalNavigation navigation: WKNavigation!, withError error: Error) {
        finish(nil)
    }
}