const std = @import("std");

pub const Color = enum(u1) {
    black = 0,
    white = 1,
};

pub const PieceType = enum(u3) {
    none = 0,
    pawn = 1,
    knight = 2,
    bishop = 3,
    rook = 4,
    queen = 5,
    king = 6,
};

pub fn encode(color: Color, piece_type: PieceType) u4 {
    if (piece_type == .none) return 0;
    return (@as(u4, @intFromEnum(color)) << 3) | @as(u4, @intFromEnum(piece_type));
}

pub fn typeOf(encoded: u4) PieceType {
    const type_bits: u3 = @intCast(encoded & 0b111);
    return @enumFromInt(type_bits);
}

pub fn colorOf(encoded: u4) ?Color {
    if (typeOf(encoded) == .none) return null;
    const color_bit: u1 = @intCast(encoded >> 3);
    return @enumFromInt(color_bit);
}

pub fn fromFENChar(ch: u8) !u4 {
    switch (ch) {
        'P' => return encode(.white, .pawn),
        'N' => return encode(.white, .knight),
        'B' => return encode(.white, .bishop),
        'R' => return encode(.white, .rook),
        'Q' => return encode(.white, .queen),
        'K' => return encode(.white, .king),
        'p' => return encode(.black, .pawn),
        'n' => return encode(.black, .knight),
        'b' => return encode(.black, .bishop),
        'r' => return encode(.black, .rook),
        'q' => return encode(.black, .queen),
        'k' => return encode(.black, .king),
        else => return error.InvalidPieceType,
    }
}

test "encode pieces uses low bits for type and bit 3 for color" {
    try std.testing.expectEqual(@as(u4, 1), encode(.black, .pawn));
    try std.testing.expectEqual(@as(u4, 6), encode(.black, .king));
    try std.testing.expectEqual(@as(u4, 9), encode(.white, .pawn));
    try std.testing.expectEqual(@as(u4, 14), encode(.white, .king));
    try std.testing.expectEqual(@as(u4, 0), encode(.white, .none));
    try std.testing.expectEqual(@as(u4, 10), encode(.white, .knight));
    try std.testing.expectEqual(@as(u4, 3), encode(.black, .bishop));
}

test "fromFENChar encodes white and black pieces" {
    try std.testing.expectEqual(encode(.white, .pawn), fromFENChar('P'));
    try std.testing.expectEqual(encode(.black, .king), fromFENChar('k'));
}

test "fromFENChar rejects invalid character" {
    try std.testing.expectError(error.InvalidPieceType, fromFENChar('x'));
}

test "typeOf returns correct type" {
    try std.testing.expectEqual(.pawn, typeOf(9));
    try std.testing.expectEqual(.king, typeOf(6));
    try std.testing.expectEqual(.none, typeOf(0));
}

test "colorOf returns correct color" {
    try std.testing.expectEqual(null, colorOf(0));
    try std.testing.expectEqual(.white, colorOf(14));
    try std.testing.expectEqual(.black, colorOf(3));
}