import SwiftUI
import UIKit

extension UIImage {
    /// Extracts dominant colors from the image using pixel sampling.
    /// - Parameter count: Number of colors to extract (default 3)
    /// - Returns: Array of SwiftUI Colors, always includes white and black as fallbacks
    func dominantColors(count: Int = 3) -> [Color] {
        guard let cgImage = self.cgImage else {
            return defaultColors(count: count)
        }
        
        let width = cgImage.width
        let height = cgImage.height
        
        guard width > 0, height > 0 else {
            return defaultColors(count: count)
        }
        
        // Create a smaller sample size for performance
        let sampleSize = 50
        let scaleX = max(1, width / sampleSize)
        let scaleY = max(1, height / sampleSize)
        
        guard let context = CGContext(
            data: nil,
            width: sampleSize,
            height: sampleSize,
            bitsPerComponent: 8,
            bytesPerRow: sampleSize * 4,
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
        ) else {
            return defaultColors(count: count)
        }
        
        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: sampleSize, height: sampleSize))
        
        guard let pixelData = context.data else {
            return defaultColors(count: count)
        }
        
        let data = pixelData.bindMemory(to: UInt8.self, capacity: sampleSize * sampleSize * 4)
        
        // Collect colors with their frequencies
        var colorCounts: [ColorBucket: Int] = [:]
        
        for y in 0..<sampleSize {
            for x in 0..<sampleSize {
                let offset = (y * sampleSize + x) * 4
                let r = data[offset]
                let g = data[offset + 1]
                let b = data[offset + 2]
                let a = data[offset + 3]
                
                // Skip transparent or near-transparent pixels
                guard a > 128 else { continue }
                
                // Skip near-white and near-black (we'll add those as defaults)
                let brightness = (Int(r) + Int(g) + Int(b)) / 3
                guard brightness > 30 && brightness < 225 else { continue }
                
                // Bucket colors to reduce noise (group similar colors)
                let bucket = ColorBucket(
                    r: UInt8((Int(r) / 32) * 32),
                    g: UInt8((Int(g) / 32) * 32),
                    b: UInt8((Int(b) / 32) * 32)
                )
                
                colorCounts[bucket, default: 0] += 1
            }
        }
        
        // Sort by frequency and take top colors
        let sortedColors = colorCounts.sorted { $0.value > $1.value }
        
        var extractedColors: [Color] = []
        
        // Add top colors, ensuring they're distinct
        for (bucket, _) in sortedColors {
            let color = Color(
                red: Double(bucket.r) / 255.0,
                green: Double(bucket.g) / 255.0,
                blue: Double(bucket.b) / 255.0
            )
            
            // Check if this color is distinct enough from existing ones
            if !extractedColors.contains(where: { $0.isClose(to: color) }) {
                extractedColors.append(color)
            }
            
            if extractedColors.count >= count - 1 {
                break
            }
        }
        
        // Always include white as first option
        var result: [Color] = [.white]
        result.append(contentsOf: extractedColors)
        
        // Add black if we have room
        if result.count < count {
            result.append(.black)
        }
        
        return Array(result.prefix(count))
    }
    
    private func defaultColors(count: Int) -> [Color] {
        let defaults: [Color] = [.white, .black, Color(red: 0.2, green: 0.2, blue: 0.2)]
        return Array(defaults.prefix(count))
    }
}

// MARK: - Color Bucket for grouping similar colors

private struct ColorBucket: Hashable {
    let r: UInt8
    let g: UInt8
    let b: UInt8
}

// MARK: - Color Comparison Extension

private extension Color {
    /// Checks if two colors are visually similar.
    func isClose(to other: Color, threshold: Double = 0.15) -> Bool {
        guard let selfComponents = self.cgColor?.components,
              let otherComponents = other.cgColor?.components,
              selfComponents.count >= 3,
              otherComponents.count >= 3 else {
            return false
        }
        
        let rDiff = abs(selfComponents[0] - otherComponents[0])
        let gDiff = abs(selfComponents[1] - otherComponents[1])
        let bDiff = abs(selfComponents[2] - otherComponents[2])
        
        return rDiff < threshold && gDiff < threshold && bDiff < threshold
    }
}
