import Foundation
import CoreImage
import Vision

/// Errors that can occur during board detection and processing
enum BoardDetectionError: Error {
    case boardNotFound
    case invalidBoardDimensions
    case squareExtractionFailed
    case imageProcessingFailed
}

final class BoardDetector {
    // Share CIContext to avoid creating too many Metal command queues
    private static let shared = CIContext()
    private var context: CIContext { BoardDetector.shared }
    
    func detectBoard(in image: CIImage) -> CGRect? {
        let request = VNDetectRectanglesRequest()
        request.minimumAspectRatio = 0.8
        request.maximumAspectRatio = 1.2
        request.minimumSize = 0.4
        request.maximumObservations = 1
        request.quadratureTolerance = 30
        request.minimumConfidence = 0.9
        
        let requestHandler = VNImageRequestHandler(ciImage: image, options: [:])
        do {
            try requestHandler.perform([request])
        } catch {
            print("ERROR: Rectangle detection failed - \(error)")
            return nil
        }
        
        guard let observations = request.results,
              !observations.isEmpty else {
            return nil
        }
        
        let bestObservation = observations[0]
        let imageSize = image.extent.size
        let transform = CGAffineTransform(scaleX: imageSize.width, y: imageSize.height)
        let detectedRect = bestObservation.boundingBox.applying(transform)
        
        guard validateDetectedRect(detectedRect, in: image) else {
            return nil
        }
        
        return detectedRect
    }
    
    /// Piece recognizer for analyzing extracted squares
    private let pieceRecognizer: PieceRecognizer
    
    /// Initialize with a piece recognizer
    init(pieceRecognizer: PieceRecognizer) {
        self.pieceRecognizer = pieceRecognizer
    }
    
    /// Extract a specific square from the board image
    private func extractSquare(from image: CIImage, in rect: CGRect, at position: BoardPosition) throws -> CGImage {
        let squareSize = rect.width / 8
        let x = rect.minX + CGFloat(position.file) * squareSize
        let y = rect.minY + CGFloat(position.rank) * squareSize
        let squareRect = CGRect(x: x, y: y, width: squareSize, height: squareSize)
        
        guard image.extent.contains(squareRect) else {
            print("ERROR: Square bounds outside image extent at \(position.file),\(position.rank)")
            throw BoardDetectionError.squareExtractionFailed
        }
        
        let croppedImage = image.cropped(to: squareRect)
        
        guard let cgImage = context.createCGImage(croppedImage, from: croppedImage.extent) else {
            print("ERROR: Failed to create square image at \(position.file),\(position.rank)")
            throw BoardDetectionError.squareExtractionFailed
        }
        
        return cgImage
    }
    
    /// Analyze the board and return the chess position
    func analyzeBoard(in image: CIImage) async throws -> ChessPosition {
        print("=== ANALYZING BOARD ===")
        
        guard let boardRect = detectBoard(in: image) else {
            print("ERROR: Board detection failed")
            throw BoardDetectionError.boardNotFound
        }
        
        let aspectRatio = boardRect.width / boardRect.height
        guard boardRect.width > 0, boardRect.height > 0,
              abs(1 - aspectRatio) < 0.1 else {
            print("ERROR: Invalid board dimensions")
            throw BoardDetectionError.invalidBoardDimensions
        }
        
        var position = ChessPosition()
        var recognizedPieces = 0
        
        // Reset piece counts before starting new scan
        pieceRecognizer.resetPieceCounts()
        
        // Process each square
        for rank in 0..<8 {
            for file in 0..<8 {
                let boardPosition = BoardPosition(file: file, rank: rank)
                let squareImage = try extractSquare(from: image, in: boardRect, at: boardPosition)
                
                // Set the current position being analyzed
                pieceRecognizer.currentPosition = boardPosition
                
                if let piece = try await pieceRecognizer.recognizePiece(from: squareImage) {
                    position[boardPosition] = piece
                    recognizedPieces += 1
                }
            }
        }
        
        // Validate piece counts after all squares are processed
        try pieceRecognizer.validatePieceCounts()
        
        print("\nPieces recognized: \(recognizedPieces)")
        print("\nRecognized position:")
        for rank in (0...7).reversed() {
            var rankStr = "\(rank + 1) "
            for file in 0...7 {
                if let piece = position[BoardPosition(file: file, rank: rank)] {
                    rankStr += "\(piece.color == .white ? "w" : "b")\(piece.type.fenSymbol.uppercased()) "
                } else {
                    rankStr += ".. "
                }
            }
            print(rankStr)
        }
        print("  a  b  c  d  e  f  g  h")
        
        guard position.isValid else {
            print("\nERROR: Invalid chess position")
            throw BoardDetectionError.imageProcessingFailed
        }
        
        print("=== BOARD ANALYSIS COMPLETE ===")
        return position
    }
    
    private func validateDetectedRect(_ rect: CGRect, in image: CIImage) -> Bool {
        let imageSize = image.extent.size
        
        guard image.extent.contains(rect) else {
            print("ERROR: Detected rectangle outside image bounds")
            return false
        }
        
        let aspectRatio = rect.width / rect.height
        guard aspectRatio >= 0.9 && aspectRatio <= 1.1 else {
            print("ERROR: Invalid board aspect ratio: \(aspectRatio)")
            return false
        }
        
        let minDimension = min(imageSize.width, imageSize.height)
        let boardSize = max(rect.width, rect.height)
        let sizeRatio = boardSize / minDimension
        guard sizeRatio >= 0.4 else {
            print("ERROR: Board too small relative to image")
            return false
        }
        
        return true
    }
}
