import Foundation
import Vision
import CoreML
import CoreImage

/// Errors that can occur during piece recognition
enum PieceRecognitionError: Error {
    case invalidImageDimensions
    case modelLoadError
    case recognitionFailed(String)
    case lowConfidence
    case invalidInput
}

/// A class responsible for recognizing chess pieces from images
final class PieceRecognizer {
    // MARK: - Properties
    
    /// Vision model for piece classification
    private let vnModel: VNCoreMLModel
    
    /// Shared CIContext for image processing
    private static let ciContext = CIContext()
    
    /// Latest classification results
    private var classificationResults: [VNClassificationObservation]?
    
    /// Confidence threshold
    private let confidenceThreshold: Float = 0.75
    
    // MARK: - Initialization
    
    init() throws {
        print("=== INITIALIZING PIECE RECOGNIZER ===")
        
        let bundle = Bundle.main
        
        // Load model
        guard let modelURL = bundle.url(forResource: "ChessPieceClassifier", withExtension: "mlmodelc") else {
            print("ERROR: Model not found in bundle at \(bundle.bundlePath)")
            throw PieceRecognitionError.modelLoadError
        }
        
        do {
            let config = MLModelConfiguration()
            config.computeUnits = .all
            let model = try MLModel(contentsOf: modelURL, configuration: config)
            self.vnModel = try VNCoreMLModel(for: model)
            print("Model loaded successfully from: \(modelURL.path)")
        } catch {
            print("ERROR: Failed to load model - \(error)")
            throw PieceRecognitionError.modelLoadError
        }
    }
    
    // MARK: - Recognition Methods
    
    func recognizeSquare(from image: CGImage, row: Int = 0, col: Int = 0) async throws -> SquareClassification {
        // Validate image dimensions
        guard image.width > 0, image.height > 0,
              abs(1 - Float(image.width) / Float(image.height)) < 0.1 else {
            print("ERROR: Invalid square dimensions \(image.width)x\(image.height)")
            throw PieceRecognitionError.invalidImageDimensions
        }
        
        let handler = VNImageRequestHandler(cgImage: image)
        var classificationError: Error?
        
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
            let request = VNCoreMLRequest(model: vnModel) { request, error in
                if let error = error {
                    classificationError = error
                    continuation.resume(throwing: error)
                    return
                }
                self.classificationResults = request.results as? [VNClassificationObservation]
                continuation.resume()
            }
            request.imageCropAndScaleOption = .centerCrop
            
            do {
                try handler.perform([request])
            } catch {
                continuation.resume(throwing: error)
            }
        }
        
        if let error = classificationError {
            print("ERROR: Classification failed - \(error)")
            throw PieceRecognitionError.recognitionFailed(error.localizedDescription)
        }
        
        guard let results = self.classificationResults,
              let topResult = results.first else {
            print("ERROR: No classification results")
            throw PieceRecognitionError.recognitionFailed("No results")
        }
        
        // Print results
        print("\nClassification results:")
        for result in results.prefix(3) {
            print("- \(result.identifier): \(result.confidence)")
        }
        
        // Check confidence threshold
        if topResult.confidence < confidenceThreshold {
            print("Low confidence (\(topResult.confidence)) for \(topResult.identifier)")
            return .empty()
        }
        
        // Try to create a SquareClassification from the label
        if let classification = SquareClassification(label: topResult.identifier) {
            print("Classified as \(topResult.identifier) with confidence \(topResult.confidence)")
            return classification
        }
        
        // Return empty square if classification fails
        print("Classification failed for \(topResult.identifier) (\(topResult.confidence))")
        return .empty()
    }
    
    /// Preprocess an image for recognition
    /// - Parameter image: Input CGImage
    /// - Returns: Preprocessed CGImage
    func preprocessImage(_ image: CGImage) throws -> CGImage {
        let ciImage = CIImage(cgImage: image)
        
        // Apply preprocessing filters
        // First pass: Enhance contrast and edges
        var processed = ciImage
            .applyingFilter("CIColorControls", parameters: [
                "inputContrast": 1.3,
                "inputBrightness": 0.0,
                "inputSaturation": 1.0
            ])
            .applyingFilter("CIUnsharpMask", parameters: [
                "inputRadius": 2.0,
                "inputIntensity": 0.8
            ])
        
        // Second pass: Reduce noise and enhance details
        processed = processed
            .applyingFilter("CINoiseReduction", parameters: [
                "inputNoiseLevel": 0.2,
                "inputSharpness": 0.6
            ])
            .applyingFilter("CIColorControls", parameters: [
                "inputContrast": 1.2,
                "inputBrightness": 0.0,
                "inputSaturation": 1.0
            ])
        
        // Convert back to CGImage
        guard let outputImage = Self.ciContext.createCGImage(processed, from: processed.extent) else {
            throw PieceRecognitionError.invalidInput
        }
        
        return outputImage
    }
}
