import Foundation

struct SquareClassification {
    let pieceType: PieceType?
    let pieceColor: PieceColor?
    let isHighlighted: Bool
    
    var isEmpty: Bool {
        return pieceType == nil
    }
    
    static func empty(highlighted: Bool = false) -> SquareClassification {
        return SquareClassification(pieceType: nil, pieceColor: nil, isHighlighted: highlighted)
    }
    
    init(pieceType: PieceType?, pieceColor: PieceColor?, isHighlighted: Bool = false) {
        self.pieceType = pieceType
        self.pieceColor = pieceColor
        self.isHighlighted = isHighlighted
    }
    
    init?(label: String) {
        // Match exact categories from trained model
        switch label {
        case "white_pawn": 
            self.init(pieceType: .pawn, pieceColor: .white)
        case "white_knight":
            self.init(pieceType: .knight, pieceColor: .white)
        case "white_bishop":
            self.init(pieceType: .bishop, pieceColor: .white)
        case "white_rook":
            self.init(pieceType: .rook, pieceColor: .white)
        case "white_queen":
            self.init(pieceType: .queen, pieceColor: .white)
        case "white_king":
            self.init(pieceType: .king, pieceColor: .white)
        case "black_pawn":
            self.init(pieceType: .pawn, pieceColor: .black)
        case "black_knight":
            self.init(pieceType: .knight, pieceColor: .black)
        case "black_bishop":
            self.init(pieceType: .bishop, pieceColor: .black)
        case "black_rook":
            self.init(pieceType: .rook, pieceColor: .black)
        case "black_queen":
            self.init(pieceType: .queen, pieceColor: .black)
        case "black_king":
            self.init(pieceType: .king, pieceColor: .black)
        default:
            // Any unrecognized label returns an empty square
            self.init(pieceType: nil, pieceColor: nil, isHighlighted: false)
        }
    }
}
