import AppKit
import CoreGraphics
import ImageIO
import ScreenCaptureKit
import UniformTypeIdentifiers

struct WindowInfo {
    let id: CGWindowID
    let pid: pid_t
    let app: String
    let title: String
    let frame: CGRect
}

func windowInfos(_ list: [[String: Any]]) -> [WindowInfo] {
    return list.compactMap { w in
        guard (w[kCGWindowLayer as String] as? Int) == 0,
              let id = w[kCGWindowNumber as String] as? Int,
              let pid = w[kCGWindowOwnerPID as String] as? Int,
              let b = w[kCGWindowBounds as String] as? NSDictionary,
              let frame = CGRect(dictionaryRepresentation: b),
              frame.width > 1, frame.height > 1
        else {
            return nil
        }
        return WindowInfo(id: CGWindowID(id), pid: pid_t(pid), app: w[kCGWindowOwnerName as String] as? String ?? "", title: w[kCGWindowName as String] as? String ?? "", frame: frame)
    }
}

func onScreenWindows() -> [WindowInfo] {
    let list = CGWindowListCopyWindowInfo([.optionOnScreenOnly, .excludeDesktopElements], kCGNullWindowID) as? [[String: Any]] ?? []
    return windowInfos(list)
}

func window(_ id: Int) throws -> WindowInfo {
    guard let wid = CGWindowID(exactly: id) else {
        throw ToolError("no window \(id); screen_list shows the current ones")
    }
    let list = CGWindowListCopyWindowInfo(.optionIncludingWindow, wid) as? [[String: Any]] ?? []
    guard let w = windowInfos(list).first else {
        throw ToolError("no window \(id); screen_list shows the current ones")
    }
    return w
}

func describe(_ w: WindowInfo) -> String {
    return "\(w.id) \(w.app) \(w.title.debugDescription) \(Int(w.frame.width))x\(Int(w.frame.height)) at \(Int(w.frame.minX)),\(Int(w.frame.minY))"
}

func screenList(_ args: [String: Any]) async throws -> [Content] {
    var out = "apps (pid name bundle):\n"
    for app in NSWorkspace.shared.runningApplications where app.activationPolicy == .regular {
        out += "\(app.processIdentifier) \(app.localizedName ?? "") \(app.bundleIdentifier ?? "")\(app.isActive ? " active" : "")\(app.isHidden ? " hidden" : "")\n"
    }
    out += "\nwindows, front to back (id app title size at x,y):\n"
    for w in onScreenWindows() {
        out += describe(w) + "\n"
    }
    out += "\ndisplays (index size at x,y):\n"
    for (i, s) in NSScreen.screens.enumerated() {
        let id = s.deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? CGDirectDisplayID ?? 0
        let b = CGDisplayBounds(id)
        out += "\(i) \(Int(b.width))x\(Int(b.height)) at \(Int(b.minX)),\(Int(b.minY))\(CGDisplayIsMain(id) != 0 ? " main" : "")\n"
    }
    return [.text(out)]
}

struct Shot {
    let png: Data
    let note: String
}

func capture(window id: Int?, display: Int?, crop: [String: Any]?, maxEdge: Double) async throws -> Shot {
    if !maxEdge.isFinite || maxEdge < 16 || maxEdge > 8192 {
        throw ToolError("max_edge must be between 16 and 8192 pixels")
    }
    let content = try await SCShareableContent.excludingDesktopWindows(false, onScreenWindowsOnly: false)
    let filter: SCContentFilter
    let size: CGSize
    let what: String
    var base = CGPoint.zero
    if let id {
        guard let w = content.windows.first(where: { Int($0.windowID) == id }), let app = w.owningApplication else {
            throw ToolError("no window \(id) to capture; screen_list shows the current ones")
        }
        let center = CGPoint(x: w.frame.midX, y: w.frame.midY)
        guard w.isOnScreen, let d = content.displays.first(where: { $0.frame.contains(center) }) else {
            throw ToolError("window \(id) is not on screen")
        }
        filter = SCContentFilter(display: d, including: [app], exceptingWindows: [])
        size = w.frame.size
        base = CGPoint(x: w.frame.minX - d.frame.minX, y: w.frame.minY - d.frame.minY)
        what = "window \(id)"
    } else {
        let screens = NSScreen.screens
        let i = display ?? 0
        guard i >= 0, i < screens.count,
              let did = screens[i].deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? CGDirectDisplayID,
              let d = content.displays.first(where: { $0.displayID == did })
        else {
            throw ToolError("no display \(i)")
        }
        filter = SCContentFilter(display: d, excludingWindows: [])
        size = CGSize(width: d.width, height: d.height)
        what = "display \(i)"
    }
    var src = CGRect(origin: .zero, size: size)
    if let crop {
        src = CGRect(x: crop.double("x") ?? 0, y: crop.double("y") ?? 0, width: crop.double("width") ?? size.width, height: crop.double("height") ?? size.height).intersection(src)
        if src.isEmpty {
            throw ToolError("crop is outside the \(Int(size.width))x\(Int(size.height)) \(what)")
        }
    }
    let k = min(1, maxEdge / max(src.width, src.height))
    let cfg = SCStreamConfiguration()
    cfg.sourceRect = src.offsetBy(dx: base.x, dy: base.y)
    cfg.width = max(1, Int((src.width * k).rounded()))
    cfg.height = max(1, Int((src.height * k).rounded()))
    cfg.showsCursor = false
    let img = try await SCScreenshotManager.captureImage(contentFilter: filter, configuration: cfg)
    let data = NSMutableData()
    let dest = CGImageDestinationCreateWithData(data, UTType.png.identifier as CFString, 1, nil)!
    CGImageDestinationAddImage(dest, img, nil)
    CGImageDestinationFinalize(dest)
    var note = "\(what), \(Int(size.width))x\(Int(size.height)) points; image \(img.width)x\(img.height)"
    if k < 1 || src.origin != .zero {
        note += String(format: "; \(what) point = image pixel / %.4f", k)
        if src.origin != .zero {
            note += " + (\(Int(src.minX)), \(Int(src.minY)))"
        }
    } else {
        note += ", 1 pixel = 1 point"
    }
    return Shot(png: data as Data, note: note)
}

func captureTool(_ args: [String: Any]) async throws -> [Content] {
    let shot = try await capture(window: args.int("window"), display: args.int("display"), crop: args.dict("crop"), maxEdge: args.double("max_edge") ?? 1568)
    return [.text(shot.note), .png(shot.png)]
}