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
    case invalidPosition(String)
}

/// 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()
    
    /// Piece counts for validation
    private var whitePieceCount: [PieceType: Int] = [:]
    private var blackPieceCount: [PieceType: Int] = [:]
    
    // MARK: - Initialization
    
    init() throws {
        print("=== INITIALIZING PIECE RECOGNIZER ===")
        
        let bundle = Bundle.main
        
        // Load model from bundle
        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
        }
        
        resetPieceCounts()
    }
    
    // MARK: - Recognition Methods
    
    /// Reset piece counts for new position
    func resetPieceCounts() {
        whitePieceCount = [
            .king: 0,
            .queen: 0,
            .rook: 0,
            .bishop: 0,
            .knight: 0,
            .pawn: 0
        ]
        blackPieceCount = [
            .king: 0,
            .queen: 0,
            .rook: 0,
            .bishop: 0,
            .knight: 0,
            .pawn: 0
        ]
    }
    
    /// Validate piece counts and ensure the position is legal
    func validatePieceCounts() throws {
        var errors: [String] = []
        
        // Check white pieces
        if whitePieceCount[.king] != 1 {
            let count = whitePieceCount[.king] ?? 0
            errors.append("Invalid white king count: \(count)")
        }
        if let count = whitePieceCount[.queen], count > 1 {
            errors.append("Too many white queens: \(count)")
        }
        if let count = whitePieceCount[.rook], count > 2 {
            errors.append("Too many white rooks: \(count)")
        }
        if let count = whitePieceCount[.bishop], count > 2 {
            errors.append("Too many white bishops: \(count)")
        }
        if let count = whitePieceCount[.knight], count > 2 {
            errors.append("Too many white knights: \(count)")
        }
        if let count = whitePieceCount[.pawn], count > 8 {
            errors.append("Too many white pawns: \(count)")
        }
        
        // Check black pieces
        if blackPieceCount[.king] ?? 0 != 1 {
            let count = blackPieceCount[.king] ?? 0
            errors.append("Invalid black king count: \(count)")
        }
        if let count = blackPieceCount[.queen], count > 1 {
            errors.append("Too many black queens: \(count)")
        }
        if let count = blackPieceCount[.rook], count > 2 {
            errors.append("Too many black rooks: \(count)")
        }
        if let count = blackPieceCount[.bishop], count > 2 {
            errors.append("Too many black bishops: \(count)")
        }
        if let count = blackPieceCount[.knight], count > 2 {
            errors.append("Too many black knights: \(count)")
        }
        if let count = blackPieceCount[.pawn], count > 8 {
            errors.append("Too many black pawns: \(count)")
        }
        
        if !errors.isEmpty {
            throw PieceRecognitionError.invalidPosition(errors.joined(separator: ", "))
        }
    }
    
    /// Update piece count
    private func updatePieceCount(piece: ChessPiece) {
        if piece.color == .white {
            whitePieceCount[piece.type] = (whitePieceCount[piece.type] ?? 0) + 1
        } else {
            blackPieceCount[piece.type] = (blackPieceCount[piece.type] ?? 0) + 1
        }
    }
    
    /// Check if adding this piece would exceed limits
    private func wouldExceedLimits(_ piece: ChessPiece) -> Bool {
        let count = piece.color == .white ? whitePieceCount[piece.type] ?? 0 : blackPieceCount[piece.type] ?? 0
        switch piece.type {
        case .king: return count >= 1
        case .queen: return count >= 1
        case .rook, .bishop, .knight: return count >= 2
        case .pawn: return count >= 8
        }
    }
    
    /// Recognize a chess piece from a square image
    /// - Parameters:
    ///   - image: CGImage of the chess square
    ///   - completion: Callback with result (ChessPiece if recognized, nil if empty)
    /// - Throws: PieceRecognitionError
    func recognizePiece(from image: CGImage) async throws -> ChessPiece? {
        // 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 classificationResults: [VNClassificationObservation]?
        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
                }
                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 = classificationResults,
              let topResult = results.first else {
            print("ERROR: No classification results")
            throw PieceRecognitionError.recognitionFailed("No results")
        }
        
        // Print all results to help diagnose recognition issues
        if let pos = currentPosition {
            print("\nClassification results for \(String(describing: pos)):")
        } else {
            print("\nClassification results:")
        }
        for result in results.prefix(3) {
            print("- \(result.identifier): \(result.confidence)")
        }
        
        // Check for empty squares
        if topResult.identifier == "empty_dark" || topResult.identifier == "empty_light" {
            if topResult.confidence > 0.9 {
                if let pos = currentPosition {
                    print("\nEmpty square confirmed at \(String(describing: pos))")
                }
                return nil
            }
        }
        
        // Get confidence ratio between top predictions
        let secondBestConfidence = results.count > 1 ? results[1].confidence : 0
        let confidenceRatio = topResult.confidence / (secondBestConfidence + Float.ulpOfOne)
        
        // Adjust confidence based on position-specific knowledge
        let adjustedConfidence = adjustConfidence(topResult.confidence, 
                                                for: topResult.identifier,
                                                at: currentPosition)
        
        // Much stricter piece recognition:
        // 1. Must have very high confidence (>0.98)
        // 2. Must have strong separation from second best (>5.0 ratio)
        // 3. Must make sense for position
        // 4. Must not exceed piece limits
        if adjustedConfidence > 0.98 && confidenceRatio > 5.0 {
            if let piece = try? createPiece(from: topResult.identifier) {
                if isValidPieceForPosition(piece, at: currentPosition) && !wouldExceedLimits(piece) {
                    updatePieceCount(piece: piece)
                    return piece
                }
            }
        }
        
        // Log classification details
        if let pos = currentPosition {
            print("\nClassification rejected at \(String(describing: pos)):")
        }
        print("Top result: \(topResult.identifier) (\(topResult.confidence))")
        if results.count > 1 {
            print("Second best: \(results[1].identifier) (\(results[1].confidence))")
            print("Confidence ratio: \(confidenceRatio)")
        }
        print("Adjusted confidence: \(adjustedConfidence)")
        
        return nil // Default to empty for unclear cases
    }
    
    /// Position where this piece is being recognized
    var currentPosition: BoardPosition?
    
    /// Adjust confidence based on position-specific knowledge
    private func adjustConfidence(_ confidence: Float, for identifier: String, at position: BoardPosition?) -> Float {
        guard let position = currentPosition else { return confidence }
        
        // Parse the piece info
        let components = identifier.split(separator: "_")
        guard components.count == 2,
              let color = PieceColor(rawValue: String(components[0])),
              let type = PieceType(rawValue: String(components[1])) else {
            return confidence
        }
        
        var adjustment: Float = 0.0
        
        // Back rank pieces are more likely to be correct
        if position.isBackRank(for: color) {
            // Corners should be rooks
            if type == .rook && position.isEdgeFile {
                adjustment += 0.1
            }
            // Next to corners should be knights
            if type == .knight && (position.file == 1 || position.file == 6) {
                adjustment += 0.1
            }
            // Next to knights should be bishops
            if type == .bishop && (position.file == 2 || position.file == 5) {
                adjustment += 0.1
            }
            // Center should be king/queen
            if (type == .king || type == .queen) && position.isCenterFile {
                adjustment += 0.1
            }
        }
        
        // Pawns are more likely on their starting ranks
        if type == .pawn {
            if (color == .white && position.rank == 1) ||
               (color == .black && position.rank == 6) {
                adjustment += 0.1
            }
        }
        
        return min(1.0, confidence + adjustment)
    }
    
    /// Validate if a piece makes sense for its position
    private func isValidPieceForPosition(_ piece: ChessPiece, at position: BoardPosition?) -> Bool {
        guard let pos = position else { return true }
        
        // Basic position validation
        switch piece.type {
        case .king:
            // Kings can't be on the first or last rank of opponent's side
            if piece.color == .white && pos.rank == 7 { return false }
            if piece.color == .black && pos.rank == 0 { return false }
            
        case .pawn:
            // Pawns can't be on first or last rank
            if pos.rank == 0 || pos.rank == 7 { return false }
            // White pawns can't be behind their starting rank
            if piece.color == .white && pos.rank > 6 { return false }
            // Black pawns can't be behind their starting rank
            if piece.color == .black && pos.rank < 1 { return false }
            
        default:
            // Other pieces can move freely
            break
        }
        
        return true
    }
    
    /// Helper to create a chess piece from a classification label
    private func createPiece(from identifier: String) throws -> ChessPiece {
        // Validate it's not an empty square
        guard !identifier.starts(with: "empty_") else {
            print("ERROR: Cannot create piece from empty square label: \(identifier)")
            throw PieceRecognitionError.invalidInput
        }
        
        let components = identifier.split(separator: "_")
        guard components.count == 2,
              let color = PieceColor(rawValue: String(components[0])),
              let type = PieceType(rawValue: String(components[1])) else {
            print("ERROR: Invalid piece label format: \(identifier)")
            throw PieceRecognitionError.invalidInput
        }
        return ChessPiece(type: type, color: color)
    }
    
    /// 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
        let processed = ciImage
            .applyingFilter("CIColorControls", parameters: [
                kCIInputContrastKey: 1.1,
                kCIInputBrightnessKey: 0.0,
                kCIInputSaturationKey: 1.1
            ])
            .applyingFilter("CIUnsharpMask", parameters: [
                kCIInputRadiusKey: 1.0,
                kCIInputIntensityKey: 0.5
            ])
        
        // Convert back to CGImage
        guard let outputImage = Self.ciContext.createCGImage(processed, from: processed.extent) else {
            throw PieceRecognitionError.invalidInput
        }
        
        return outputImage
    }
}

// MARK: - BoardPosition Extensions

extension BoardPosition {
    /// Initialize from file and rank indices
    init(file: Int, rank: Int) {
        self.file = file
        self.rank = rank
    }
    
    /// Whether this position is on the edge of the board (files a or h)
    var isEdgeFile: Bool {
        return file == 0 || file == 7
    }
    
    /// Whether this position is on the back rank for the given color
    func isBackRank(for color: PieceColor) -> Bool {
        return (color == .white && rank == 0) || (color == .black && rank == 7)
    }
    
    /// Whether this position is in the center files (d or e)
    var isCenterFile: Bool {
        return file == 3 || file == 4
    }
}
