import AppKit
import Foundation

/// Captures the contents of an NSWindow as image data.
@MainActor
public enum Snapshot {

    /// Image output format.
    public enum ImageFormat: Sendable {
        case jpeg(quality: Double)
        case png
    }

    /// Capture the current contents of a window's content view.
    /// Uses `cacheDisplay` to render directly from the view hierarchy, which works
    /// regardless of window position (including off-screen/headless) and captures
    /// only the SwiftUI content without the title bar.
    /// - Parameters:
    ///   - window: The window to capture.
    ///   - format: Output format (default: JPEG at 0.85 quality).
    public static func capture(window: NSWindow, format: ImageFormat = .jpeg(quality: 0.85)) throws -> Data {
        guard let contentView = window.contentView else {
            throw SnapshotError.captureFailed
        }
        let bounds = contentView.bounds
        guard bounds.width > 0, bounds.height > 0 else {
            throw SnapshotError.captureFailed
        }
        // Pin the raster to a deterministic 1x (one pixel per point). bitmapImageRepForCachingDisplay
        // inherits the host window's backingScaleFactor — 2x on a Retina display — which made the
        // snapshot's pixel dimensions depend on which machine the daemon ran on. Building the rep
        // manually with pixel dimensions equal to the point bounds keeps output reproducible.
        guard
            let bitmapRep = NSBitmapImageRep(
                bitmapDataPlanes: nil,
                pixelsWide: Int(bounds.width.rounded()),
                pixelsHigh: Int(bounds.height.rounded()),
                bitsPerSample: 8,
                samplesPerPixel: 4,
                hasAlpha: true,
                isPlanar: false,
                colorSpaceName: .deviceRGB,
                bytesPerRow: 0,
                bitsPerPixel: 0
            )
        else {
            throw SnapshotError.captureFailed
        }
        bitmapRep.size = bounds.size
        contentView.cacheDisplay(in: bounds, to: bitmapRep)

        switch format {
        case .jpeg(let quality):
            guard
                let data = bitmapRep.representation(
                    using: .jpeg,
                    properties: [.compressionFactor: NSNumber(value: quality)]
                )
            else {
                throw SnapshotError.encodingFailed
            }
            return data
        case .png:
            guard let data = bitmapRep.representation(using: .png, properties: [:]) else {
                throw SnapshotError.encodingFailed
            }
            return data
        }
    }

    /// Capture and write to a file.
    public static func capture(window: NSWindow, format: ImageFormat = .jpeg(quality: 0.85), to path: URL) throws {
        let data = try capture(window: window, format: format)
        try data.write(to: path)
    }

    /// Re-encode an agent-rendered image file (PNG) to the requested output format.
    /// PNG output returns the bytes unchanged; JPEG transcodes. Used by the JIT
    /// structural path, where the snapshot is rendered in the agent, not the window.
    public static func encode(imageAt path: URL, format: ImageFormat) throws -> Data {
        let data = try Data(contentsOf: path)
        switch format {
        case .png:
            return data
        case .jpeg(let quality):
            guard
                let rep = NSBitmapImageRep(data: data),
                let out = rep.representation(
                    using: .jpeg,
                    properties: [.compressionFactor: NSNumber(value: quality)]
                )
            else {
                throw SnapshotError.encodingFailed
            }
            return out
        }
    }
}

public enum SnapshotError: Error, LocalizedError, CustomStringConvertible {
    case captureFailed
    case encodingFailed

    public var description: String {
        switch self {
        case .captureFailed: return "Failed to capture window contents"
        case .encodingFailed: return "Failed to encode screenshot"
        }
    }

    public var errorDescription: String? { description }
}
