#!/usr/bin/env swift
// Layout test harness — compiles standalone, no Xcode needed.
// Usage: swift test_layout.swift

import Foundation
import CoreGraphics

// ─── Minimal model types ───────────────────────────────────────────────

struct WindowInfo {
    let windowID: CGWindowID
    let title: String
    let appName: String
    let appBundleID: String
    let frame: CGRect
    let pid: pid_t
    let isOnScreen: Bool
    var thumbnail: CGImage? = nil
    var appIcon: Any? = nil  // NSImage not available without AppKit in script mode

    var id: CGWindowID { windowID }
    var fingerprintKey: String { "\(appBundleID)|\(title)" }
}

struct WindowGroup: Identifiable {
    let appBundleID: String
    let appName: String
    var windows: [WindowInfo]
    var frontmostIndex: Int = Int.max
    var id: String { appBundleID }
    var appIcon: Any? { nil }
}

struct WindowSnapshot {
    let groups: [WindowGroup]
    let fingerprint: Set<String>
    let totalWindowCount: Int

    init(groups: [WindowGroup]) {
        self.groups = groups
        self.totalWindowCount = groups.reduce(0) { $0 + $1.windows.count }
        var fp = Set<String>()
        for group in groups {
            for window in group.windows {
                fp.insert(window.fingerprintKey)
            }
        }
        self.fingerprint = fp
    }
}

// ─── Copy of LayoutGeometry ────────────────────────────────────────────

enum LayoutGeometry {
    static let margin: CGFloat = 24
    static let groupHeaderHeight: CGFloat = 28
    static let compactGroupHeaderHeight: CGFloat = 4
    static let groupSpacing: CGFloat = 8
    static let itemSpacing: CGFloat = 8
    static let titleHeight: CGFloat = 22
    static let minThumbnailWidth: CGFloat = 80
    static let maxThumbnailWidth: CGFloat = 480
    static let compactGroupThreshold = 6

    static func effectiveHeaderHeight(groupCount: Int) -> CGFloat {
        groupCount > compactGroupThreshold ? compactGroupHeaderHeight : groupHeaderHeight
    }

    static func uniformThumbnailSize(
        groupWindowCounts: [Int],
        availableWidth: CGFloat,
        availableHeight: CGFloat
    ) -> (size: CGSize, columns: Int) {
        let groupCount = groupWindowCounts.count
        guard groupCount > 0 else { return (CGSize(width: 200, height: 150), 1) }

        let headerH = effectiveHeaderHeight(groupCount: groupCount)
        let headerOverhead = CGFloat(groupCount) * (headerH + groupSpacing)
        let contentHeight = availableHeight - headerOverhead

        guard contentHeight > 0 else {
            return (CGSize(width: minThumbnailWidth, height: minThumbnailWidth * 0.625 + titleHeight), 1)
        }

        func totalRows(cols: Int) -> Int {
            groupWindowCounts.reduce(0) { $0 + Int(ceil(Double($1) / Double(cols))) }
        }

        var bestWidth: CGFloat = 0
        var bestHeight: CGFloat = 0
        var bestCols = 1

        let maxCols = max(1, Int((availableWidth + itemSpacing) / (minThumbnailWidth + itemSpacing)))

        for cols in 1...maxCols {
            let thumbWidth = min(maxThumbnailWidth, (availableWidth - CGFloat(cols - 1) * itemSpacing) / CGFloat(cols))
            guard thumbWidth >= minThumbnailWidth else { continue }
            let thumbHeight = thumbWidth * 0.625 + titleHeight
            let rows = totalRows(cols: cols)
            let neededHeight = CGFloat(rows) * thumbHeight + CGFloat(max(0, rows - 1)) * itemSpacing

            if neededHeight <= contentHeight && thumbWidth > bestWidth {
                bestWidth = thumbWidth
                bestHeight = thumbHeight
                bestCols = cols
            }
        }

        if bestWidth == 0 {
            var lo: CGFloat = 20
            var hi = min(maxThumbnailWidth, availableWidth)

            for _ in 0..<40 {
                let mid = (lo + hi) / 2
                let cols = max(1, Int((availableWidth + itemSpacing) / (mid + itemSpacing)))
                let thumbHeight = mid * 0.625 + titleHeight
                let rows = totalRows(cols: cols)
                let neededHeight = CGFloat(rows) * thumbHeight + CGFloat(max(0, rows - 1)) * itemSpacing

                if neededHeight <= contentHeight {
                    bestWidth = mid
                    bestHeight = thumbHeight
                    bestCols = cols
                    lo = mid
                } else {
                    hi = mid
                }
            }

            if bestWidth == 0 {
                bestCols = maxCols
                bestWidth = max(20, (availableWidth - CGFloat(bestCols - 1) * itemSpacing) / CGFloat(bestCols))
                bestHeight = bestWidth * 0.625 + titleHeight
            }
        }

        return (CGSize(width: bestWidth, height: bestHeight), bestCols)
    }
}

// ─── Simplified layout engine (no stability/reuse) ─────────────────────

struct GroupLayout {
    let columns: Int
    let thumbnailSize: CGSize
    let itemRange: Range<Int>
}

struct LayoutItem {
    let title: String
    let appName: String
    let frame: CGRect
    let groupIndex: Int
}

struct LayoutGroupHeader {
    let appName: String
    let frame: CGRect
}

struct LayoutResult {
    let items: [LayoutItem]
    let headers: [LayoutGroupHeader]
    let groupLayouts: [GroupLayout]
    let totalHeight: CGFloat
    let isCompactHeaders: Bool
}

func computeLayout(snapshot: WindowSnapshot, screenBounds: CGRect) -> LayoutResult {
    let availableWidth = screenBounds.width - LayoutGeometry.margin * 2
    let availableHeight = screenBounds.height - LayoutGeometry.margin * 2
    let groupCount = snapshot.groups.count
    let isCompact = groupCount > LayoutGeometry.compactGroupThreshold
    let headerH = LayoutGeometry.effectiveHeaderHeight(groupCount: groupCount)

    let groupWindowCounts = snapshot.groups.map { $0.windows.count }
    let (thumbSize, cols) = LayoutGeometry.uniformThumbnailSize(
        groupWindowCounts: groupWindowCounts,
        availableWidth: availableWidth,
        availableHeight: availableHeight
    )

    var items: [LayoutItem] = []
    var headers: [LayoutGroupHeader] = []
    var groupLayouts: [GroupLayout] = []

    var y = screenBounds.origin.y + LayoutGeometry.margin

    for (groupIndex, group) in snapshot.groups.enumerated() {
        let windowsInRow = min(cols, group.windows.count)
        let gridWidth = CGFloat(windowsInRow) * thumbSize.width + CGFloat(max(0, windowsInRow - 1)) * LayoutGeometry.itemSpacing
        let xOffset = screenBounds.origin.x + LayoutGeometry.margin + (availableWidth - gridWidth) / 2

        let headerFrame = CGRect(
            x: screenBounds.origin.x + LayoutGeometry.margin,
            y: y,
            width: availableWidth,
            height: headerH
        )
        headers.append(LayoutGroupHeader(appName: group.appName, frame: headerFrame))
        y += headerH

        let itemStart = items.count

        for (windowIndex, window) in group.windows.enumerated() {
            let col = windowIndex % cols
            let row = windowIndex / cols

            let x = xOffset + CGFloat(col) * (thumbSize.width + LayoutGeometry.itemSpacing)
            let itemY = y + CGFloat(row) * (thumbSize.height + LayoutGeometry.itemSpacing)

            let frame = CGRect(x: x, y: itemY, width: thumbSize.width, height: thumbSize.height)
            items.append(LayoutItem(title: window.title, appName: window.appName, frame: frame, groupIndex: groupIndex))
        }

        let rowsInGroup = Int(ceil(Double(group.windows.count) / Double(cols)))
        let groupContentHeight = CGFloat(rowsInGroup) * thumbSize.height + CGFloat(max(0, rowsInGroup - 1)) * LayoutGeometry.itemSpacing

        groupLayouts.append(GroupLayout(
            columns: cols,
            thumbnailSize: thumbSize,
            itemRange: itemStart..<items.count
        ))

        y += groupContentHeight + LayoutGeometry.groupSpacing
    }

    let totalHeight = y - screenBounds.origin.y
    return LayoutResult(items: items, headers: headers, groupLayouts: groupLayouts, totalHeight: totalHeight, isCompactHeaders: isCompact)
}

// ─── Test scenarios ────────────────────────────────────────────────────

func makeWindow(id: Int, title: String, appName: String, bundleID: String) -> WindowInfo {
    WindowInfo(
        windowID: CGWindowID(id),
        title: title,
        appName: appName,
        appBundleID: bundleID,
        frame: CGRect(x: 0, y: 0, width: 800, height: 600),
        pid: pid_t(id),
        isOnScreen: true
    )
}

func groupWindows(_ windows: [WindowInfo]) -> [WindowGroup] {
    var dict: [String: WindowGroup] = [:]
    for (i, w) in windows.enumerated() {
        if dict[w.appBundleID] == nil {
            dict[w.appBundleID] = WindowGroup(appBundleID: w.appBundleID, appName: w.appName, windows: [])
        }
        dict[w.appBundleID]!.windows.append(w)
        dict[w.appBundleID]!.frontmostIndex = min(dict[w.appBundleID]!.frontmostIndex, i)
    }
    return dict.values.sorted { $0.frontmostIndex < $1.frontmostIndex }
}

func printResult(_ result: LayoutResult, screenBounds: CGRect, label: String) {
    print("═══════════════════════════════════════════════════════════════")
    print("TEST: \(label)")
    print("Screen: \(Int(screenBounds.width))×\(Int(screenBounds.height))")
    print("Total height used: \(Int(result.totalHeight)) / \(Int(screenBounds.height))")
    print("Compact headers: \(result.isCompactHeaders)")
    print("")

    var allFit = true

    for (gi, gl) in result.groupLayouts.enumerated() {
        let header = result.headers[gi]
        print("  Group \(gi): \(header.appName)")
        print("    Header at y=\(Int(header.frame.minY)), h=\(Int(header.frame.height))")
        print("    Thumbnail: \(Int(gl.thumbnailSize.width))×\(Int(gl.thumbnailSize.height)), cols=\(gl.columns)")

        for idx in gl.itemRange {
            let item = result.items[idx]
            let f = item.frame
            let withinScreen = f.minX >= screenBounds.minX && f.minY >= screenBounds.minY &&
                              f.maxX <= screenBounds.maxX && f.maxY <= screenBounds.maxY
            let marker = withinScreen ? "✓" : "✗ OFF-SCREEN"
            if !withinScreen { allFit = false }
            print("    [\(idx)] \"\(item.title)\" @ (\(Int(f.minX)),\(Int(f.minY))) \(Int(f.width))×\(Int(f.height)) \(marker)")
        }
    }

    // Check uniform sizing within each group
    for (gi, gl) in result.groupLayouts.enumerated() {
        let sizes = Set(gl.itemRange.map { idx -> String in
            let f = result.items[idx].frame
            return "\(Int(f.width))×\(Int(f.height))"
        })
        if sizes.count > 1 {
            print("  ⚠️  Group \(gi) has NON-UNIFORM sizes: \(sizes)")
        }
    }

    // Check sizes are consistent across groups (ideally similar or justified)
    let allSizes = result.groupLayouts.map { "\(Int($0.thumbnailSize.width))×\(Int($0.thumbnailSize.height))" }
    let uniqueSizes = Set(allSizes)
    if uniqueSizes.count > 1 {
        print("\n  ℹ️  Different thumbnail sizes across groups: \(uniqueSizes)")
        let widths = result.groupLayouts.map { $0.thumbnailSize.width }
        let ratio = (widths.max() ?? 1) / max(widths.min() ?? 1, 1)
        if ratio > 3.0 {
            print("  ⚠️  Size ratio \(String(format: "%.1f", ratio))x between largest and smallest — may look unbalanced")
        }
    }

    print("\n  All items on screen: \(allFit ? "✓ YES" : "✗ NO")")
    print("")
}

// ─── Run tests ─────────────────────────────────────────────────────────

let screen1080 = CGRect(x: 0, y: 0, width: 1920, height: 1080)
let screen1440 = CGRect(x: 0, y: 0, width: 2560, height: 1440)
let screenLaptop = CGRect(x: 0, y: 0, width: 1470, height: 956)  // 14" MacBook Pro effective

// Test 1: Few windows, few apps
do {
    let windows = [
        makeWindow(id: 1, title: "Document", appName: "Safari", bundleID: "com.apple.Safari"),
        makeWindow(id: 2, title: "Gmail", appName: "Safari", bundleID: "com.apple.Safari"),
        makeWindow(id: 3, title: "Main.swift", appName: "Xcode", bundleID: "com.apple.dt.Xcode"),
    ]
    let groups = groupWindows(windows)
    let snapshot = WindowSnapshot(groups: groups)
    let result = computeLayout(snapshot: snapshot, screenBounds: screen1080)
    printResult(result, screenBounds: screen1080, label: "3 windows, 2 apps — 1080p")
}

// Test 2: Many windows, many apps
do {
    var windows: [WindowInfo] = []
    var id = 1
    for app in ["Safari", "Xcode", "Terminal", "Finder", "Slack", "Discord", "Notes", "Mail"] {
        let bundleID = "com.test.\(app.lowercased())"
        let count = app == "Safari" ? 5 : (app == "Xcode" ? 3 : (app == "Terminal" ? 4 : 1))
        for w in 0..<count {
            windows.append(makeWindow(id: id, title: "\(app) Window \(w+1)", appName: app, bundleID: bundleID))
            id += 1
        }
    }
    let groups = groupWindows(windows)
    let snapshot = WindowSnapshot(groups: groups)
    let result = computeLayout(snapshot: snapshot, screenBounds: screen1080)
    printResult(result, screenBounds: screen1080, label: "\(windows.count) windows, 8 apps — 1080p")
}

// Test 3: Lots of windows on laptop screen
do {
    var windows: [WindowInfo] = []
    var id = 1
    for app in ["Safari", "Chrome", "Xcode", "Terminal", "Finder", "Slack", "Discord", "Notes", "Mail", "Calendar"] {
        let bundleID = "com.test.\(app.lowercased())"
        let count = app == "Safari" ? 6 : (app == "Chrome" ? 4 : (app == "Xcode" ? 3 : 1))
        for w in 0..<count {
            windows.append(makeWindow(id: id, title: "\(app) Window \(w+1)", appName: app, bundleID: bundleID))
            id += 1
        }
    }
    let groups = groupWindows(windows)
    let snapshot = WindowSnapshot(groups: groups)
    let result = computeLayout(snapshot: snapshot, screenBounds: screenLaptop)
    printResult(result, screenBounds: screenLaptop, label: "\(windows.count) windows, 10 apps — Laptop (1470×956)")
}

// Test 4: Single app, many windows
do {
    let windows = (1...12).map { i in
        makeWindow(id: i, title: "Tab \(i)", appName: "Safari", bundleID: "com.apple.Safari")
    }
    let groups = groupWindows(windows)
    let snapshot = WindowSnapshot(groups: groups)
    let result = computeLayout(snapshot: snapshot, screenBounds: screen1080)
    printResult(result, screenBounds: screen1080, label: "12 windows, 1 app — 1080p")
}

// Test 5: Each app has exactly 1 window
do {
    let apps = ["Safari", "Xcode", "Terminal", "Finder", "Slack", "Discord", "Notes", "Mail", "Calendar", "Reminders", "Music", "Photos"]
    let windows = apps.enumerated().map { (i, app) in
        makeWindow(id: i + 1, title: "Main", appName: app, bundleID: "com.test.\(app.lowercased())")
    }
    let groups = groupWindows(windows)
    let snapshot = WindowSnapshot(groups: groups)
    let result = computeLayout(snapshot: snapshot, screenBounds: screenLaptop)
    printResult(result, screenBounds: screenLaptop, label: "12 windows, 12 apps (1 each) — Laptop")
}

print("═══════════════════════════════════════════════════════════════")
print("DONE")