const std = @import("std");

const board = @import("chess/board.zig");
const fen = @import("chess/fen.zig");
const bitboard = @import("chess/bitboard.zig");
const geometry = @import("geometry.zig");
const piece = @import("chess/piece.zig");

pub const PieceGroup = enum(u5) {
    white_pawn,
    white_knight,
    white_bishop,
    white_rook,
    white_queen,
    white_king,
    black_pawn,
    black_knight,
    black_bishop,
    black_rook,
    black_queen,
    black_king,
    dragged_white_pawn,
    dragged_white_knight,
    dragged_white_bishop,
    dragged_white_rook,
    dragged_white_queen,
    dragged_white_king,
    dragged_black_pawn,
    dragged_black_knight,
    dragged_black_bishop,
    dragged_black_rook,
    dragged_black_queen,
    dragged_black_king,
};

pub const PieceGroupCount = 24;

pub const PaletteEntry = struct {
    group: PieceGroup,
    encoded: u4,
};

pub const promotion_piece_types = [_]piece.PieceType{ .queen, .rook, .bishop, .knight };

pub const palette_entries = [_]PaletteEntry{
    .{ .group = .white_king, .encoded = piece.encode(.white, .king) },
    .{ .group = .white_queen, .encoded = piece.encode(.white, .queen) },
    .{ .group = .white_rook, .encoded = piece.encode(.white, .rook) },
    .{ .group = .white_bishop, .encoded = piece.encode(.white, .bishop) },
    .{ .group = .white_knight, .encoded = piece.encode(.white, .knight) },
    .{ .group = .white_pawn, .encoded = piece.encode(.white, .pawn) },
    .{ .group = .black_king, .encoded = piece.encode(.black, .king) },
    .{ .group = .black_queen, .encoded = piece.encode(.black, .queen) },
    .{ .group = .black_rook, .encoded = piece.encode(.black, .rook) },
    .{ .group = .black_bishop, .encoded = piece.encode(.black, .bishop) },
    .{ .group = .black_knight, .encoded = piece.encode(.black, .knight) },
    .{ .group = .black_pawn, .encoded = piece.encode(.black, .pawn) },
};

pub fn paletteRectForBoard(board_rect: geometry.BoardRect) geometry.BoardRect {
    const gap = board_rect.width / 48.0;
    return .{
        .left = board_rect.left + board_rect.width + gap,
        .bottom = board_rect.bottom,
        .width = board_rect.width / @as(f32, @floatFromInt(palette_entries.len)),
        .height = board_rect.height,
    };
}

pub fn promotionPopupRectForSquare(board_rect: geometry.BoardRect, square: bitboard.Square) geometry.BoardRect {
    const cell_width = (board_rect.width / 8.0) * 0.65;
    const cell_height = (board_rect.height / 8.0) * 0.65;
    const width = cell_width * @as(f32, @floatFromInt(promotion_piece_types.len));
    const height = cell_height;
    const center = geometry.boardToNdc(
        board_rect,
        @as(f32, @floatFromInt(square % 8)) + 0.5,
        @as(f32, @floatFromInt(square / 8)) + 0.5,
    );
    const board_right = board_rect.left + board_rect.width;
    const board_top = board_rect.bottom + board_rect.height;
    const unclamped_left = center[0] - (width / 2.0);
    const unclamped_bottom = center[1] - (height / 2.0);

    return .{
        .left = @max(board_rect.left, @min(unclamped_left, board_right - width)),
        .bottom = @max(board_rect.bottom, @min(unclamped_bottom, board_top - height)),
        .width = width,
        .height = height,
    };
}

pub fn promotionChoiceRectForIndex(board_rect: geometry.BoardRect, square: bitboard.Square, index: usize) geometry.BoardRect {
    const popup_rect = promotionPopupRectForSquare(board_rect, square);
    const cell_width = popup_rect.width / @as(f32, @floatFromInt(promotion_piece_types.len));
    return .{
        .left = popup_rect.left + (@as(f32, @floatFromInt(index)) * cell_width),
        .bottom = popup_rect.bottom,
        .width = cell_width,
        .height = popup_rect.height,
    };
}

pub const PieceVertexGroups = struct {
    groups: [PieceGroupCount]std.ArrayList(geometry.Vertex),

    pub fn init() PieceVertexGroups {
        return .{
            .groups = [_]std.ArrayList(geometry.Vertex){.empty} ** PieceGroupCount,
        };
    }

    pub fn deinit(self: *PieceVertexGroups, allocator: std.mem.Allocator) void {
        for (&self.groups) |*vertex_group| {
            vertex_group.deinit(allocator);
        }
    }

    pub fn group(self: *PieceVertexGroups, piece_group: PieceGroup) *std.ArrayList(geometry.Vertex) {
        return &self.groups[@intFromEnum(piece_group)];
    }

    pub fn constGroup(self: *const PieceVertexGroups, piece_group: PieceGroup) *const std.ArrayList(geometry.Vertex) {
        return &self.groups[@intFromEnum(piece_group)];
    }
};

pub fn pieceGroupFromEncoded(encoded: u4) ?PieceGroup {
    return pieceGroupFromEncodedWithOffset(encoded, false);
}

fn draggedPieceGroupFromEncoded(encoded: u4) ?PieceGroup {
    return pieceGroupFromEncodedWithOffset(encoded, true);
}

fn pieceGroupFromEncodedWithOffset(encoded: u4, dragged: bool) ?PieceGroup {
    if (encoded == 0) return null;

    const color = piece.colorOf(encoded) orelse return null;
    const piece_type = piece.typeOf(encoded);

    return switch (color) {
        .white => switch (piece_type) {
            .pawn => if (dragged) .dragged_white_pawn else .white_pawn,
            .knight => if (dragged) .dragged_white_knight else .white_knight,
            .bishop => if (dragged) .dragged_white_bishop else .white_bishop,
            .rook => if (dragged) .dragged_white_rook else .white_rook,
            .queen => if (dragged) .dragged_white_queen else .white_queen,
            .king => if (dragged) .dragged_white_king else .white_king,
            .none => null,
        },
        .black => switch (piece_type) {
            .pawn => if (dragged) .dragged_black_pawn else .black_pawn,
            .knight => if (dragged) .dragged_black_knight else .black_knight,
            .bishop => if (dragged) .dragged_black_bishop else .black_bishop,
            .rook => if (dragged) .dragged_black_rook else .black_rook,
            .queen => if (dragged) .dragged_black_queen else .black_queen,
            .king => if (dragged) .dragged_black_king else .black_king,
            .none => null,
        },
    };
}

fn appendSidewaysPieceQuad(
    vertices: *std.ArrayList(geometry.Vertex),
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    file: f32,
    rank: f32,
) !void {
    const inset: f32 = 0.08;
    const x0 = file + inset;
    const x1 = file + 1.0 - inset;
    const y0 = rank + inset;
    const y1 = rank + 1.0 - inset;

    try geometry.appendTexturedQuad(
        vertices,
        allocator,
        geometry.boardToNdc(board_rect, x1, y0),
        geometry.boardToNdc(board_rect, x1, y1),
        geometry.boardToNdc(board_rect, x0, y1),
        geometry.boardToNdc(board_rect, x0, y0),
        geometry.White,
        .{ 0.0, 1.0 },
        .{ 1.0, 1.0 },
        .{ 1.0, 0.0 },
        .{ 0.0, 0.0 },
    );
}

fn appendPieceQuadInRect(
    vertices: *std.ArrayList(geometry.Vertex),
    allocator: std.mem.Allocator,
    rect: geometry.BoardRect,
) !void {
    const inset_x = rect.width * 0.08;
    const inset_y = rect.height * 0.08;
    const x0 = rect.left + inset_x;
    const x1 = rect.left + rect.width - inset_x;
    const y0 = rect.bottom + inset_y;
    const y1 = rect.bottom + rect.height - inset_y;

    try geometry.appendTexturedQuad(
        vertices,
        allocator,
        .{ x0, y0 },
        .{ x1, y0 },
        .{ x1, y1 },
        .{ x0, y1 },
        geometry.White,
        .{ 0.0, 1.0 },
        .{ 1.0, 1.0 },
        .{ 1.0, 0.0 },
        .{ 0.0, 0.0 },
    );
}

pub fn appendPiecesGroupedFromBoard(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    state: board.BoardState,
) !void {
    try appendPiecesGroupedFromBoardExcept(groups, board_rect, allocator, state, null);
}

pub fn appendPiecesGroupedFromBoardExcept(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    state: board.BoardState,
    except_square: ?@import("board_input.zig").SquareCoord,
) !void {
    try appendPiecesGroupedFromBoardExceptWithGameOver(groups, board_rect, allocator, state, except_square, null);
}

pub fn appendPiecesGroupedFromBoardExceptWithGameOver(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    state: board.BoardState,
    except_square: ?@import("board_input.zig").SquareCoord,
    losing_king_square: ?@import("board_input.zig").SquareCoord,
) !void {
    for (0..8) |rank| {
        for (0..8) |file| {
            if (except_square) |except| {
                if (except.file == file and except.rank == rank) continue;
            }

            const encoded = state.getSquare(@intCast((rank * 8) + file));
            const piece_group = pieceGroupFromEncoded(encoded) orelse continue;

            if (losing_king_square) |losing| {
                if (losing.file == file and losing.rank == rank) {
                    try appendSidewaysPieceQuad(
                        groups.group(piece_group),
                        board_rect,
                        allocator,
                        @as(f32, @floatFromInt(file)),
                        @as(f32, @floatFromInt(rank)),
                    );
                    continue;
                }
            }

            try geometry.appendPieceQuad(
                groups.group(piece_group),
                board_rect,
                allocator,
                @as(f32, @floatFromInt(file)),
                @as(f32, @floatFromInt(rank)),
            );

        }
    }
}

pub fn appendPieceGhostAtSquare(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    encoded: u4,
    square: @import("board_input.zig").SquareCoord,
) !void {
    const piece_group = pieceGroupFromEncoded(encoded) orelse return;
    try geometry.appendPieceQuadWithColor(
        groups.group(piece_group),
        board_rect,
        allocator,
        @floatFromInt(square.file),
        @floatFromInt(square.rank),
        .{ 1.0, 1.0, 1.0, 0.35 },
    );
}

pub fn appendDraggedPiece(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    encoded: u4,
    center_ndc: [2]f32,
) !void {
    const piece_group = draggedPieceGroupFromEncoded(encoded) orelse return;
    try geometry.appendPieceQuadCenteredAtNdc(
        groups.group(piece_group),
        allocator,
        center_ndc,
        .{ board_rect.width / 8.0 * 0.84, board_rect.height / 8.0 * 0.84 },
        geometry.White,
    );
}

pub fn appendPalettePiecesGrouped(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
) !void {
    const palette_rect = paletteRectForBoard(board_rect);
    const cell_height = palette_rect.height / @as(f32, @floatFromInt(palette_entries.len));

    for (palette_entries, 0..) |entry, i| {
        const y = palette_rect.bottom + (@as(f32, @floatFromInt(palette_entries.len - 1 - i)) * cell_height);
        const cell_rect = geometry.BoardRect{
            .left = palette_rect.left,
            .bottom = y,
            .width = palette_rect.width,
            .height = cell_height,
        };
        try appendPieceQuadInRect(groups.group(entry.group), allocator, cell_rect);
    }
}

pub fn appendPromotionPiecesGrouped(
    groups: *PieceVertexGroups,
    board_rect: geometry.BoardRect,
    allocator: std.mem.Allocator,
    color: piece.Color,
    square: bitboard.Square,
) !void {
    for (promotion_piece_types, 0..) |piece_type, i| {
        const encoded = piece.encode(color, piece_type);
        const piece_group = pieceGroupFromEncoded(encoded) orelse continue;
        const cell_rect = promotionChoiceRectForIndex(board_rect, square, i);
        const piece_rect = geometry.BoardRect{
            .left = cell_rect.left + (cell_rect.width * 0.25),
            .bottom = cell_rect.bottom + (cell_rect.height * 0.25),
            .width = cell_rect.width * 0.5,
            .height = cell_rect.height * 0.5,
        };
        try appendPieceQuadInRect(groups.group(piece_group), allocator, piece_rect);
    }
}

pub fn paletteIndexForEncoded(encoded: u4) ?usize {
    for (palette_entries, 0..) |entry, i| {
        if (entry.encoded == encoded) return i;
    }
    return null;
}

fn expectGroupVertexCount(groups: *const PieceVertexGroups, piece_group: PieceGroup, expected: usize) !void {
    try std.testing.expectEqual(expected, groups.constGroup(piece_group).items.len);
}

test "pieceGroupFromEncoded maps encoded pieces to render groups" {
    try std.testing.expectEqual(PieceGroup.white_pawn, pieceGroupFromEncoded(piece.encode(.white, .pawn)).?);
    try std.testing.expectEqual(PieceGroup.white_knight, pieceGroupFromEncoded(piece.encode(.white, .knight)).?);
    try std.testing.expectEqual(PieceGroup.white_bishop, pieceGroupFromEncoded(piece.encode(.white, .bishop)).?);
    try std.testing.expectEqual(PieceGroup.white_rook, pieceGroupFromEncoded(piece.encode(.white, .rook)).?);
    try std.testing.expectEqual(PieceGroup.white_queen, pieceGroupFromEncoded(piece.encode(.white, .queen)).?);
    try std.testing.expectEqual(PieceGroup.white_king, pieceGroupFromEncoded(piece.encode(.white, .king)).?);

    try std.testing.expectEqual(PieceGroup.black_pawn, pieceGroupFromEncoded(piece.encode(.black, .pawn)).?);
    try std.testing.expectEqual(PieceGroup.black_knight, pieceGroupFromEncoded(piece.encode(.black, .knight)).?);
    try std.testing.expectEqual(PieceGroup.black_bishop, pieceGroupFromEncoded(piece.encode(.black, .bishop)).?);
    try std.testing.expectEqual(PieceGroup.black_rook, pieceGroupFromEncoded(piece.encode(.black, .rook)).?);
    try std.testing.expectEqual(PieceGroup.black_queen, pieceGroupFromEncoded(piece.encode(.black, .queen)).?);
    try std.testing.expectEqual(PieceGroup.black_king, pieceGroupFromEncoded(piece.encode(.black, .king)).?);

    try std.testing.expectEqual(null, pieceGroupFromEncoded(0));
}

test "appendPiecesGroupedFromBoard leaves all groups empty for empty board" {
    const state = board.BoardState.empty();
    const board_rect = geometry.boardRectForExtent(800, 600);

    var groups = PieceVertexGroups.init();
    defer groups.deinit(std.testing.allocator);

    try appendPiecesGroupedFromBoard(&groups, board_rect, std.testing.allocator, state);

    for (groups.groups) |vertex_group| {
        try std.testing.expectEqual(@as(usize, 0), vertex_group.items.len);
    }
}

test "appendPalettePiecesGrouped adds one piece to each piece group" {
    const board_rect = geometry.boardRectForExtent(800, 600);

    var groups = PieceVertexGroups.init();
    defer groups.deinit(std.testing.allocator);

    try appendPalettePiecesGrouped(&groups, board_rect, std.testing.allocator);

    for (groups.groups, 0..) |vertex_group, i| {
        const group: PieceGroup = @enumFromInt(@as(u5, @intCast(i)));
        const is_dragged_group = @intFromEnum(group) >= @intFromEnum(PieceGroup.dragged_white_pawn);
        const expected: usize = if (is_dragged_group) 0 else 6;
        try std.testing.expectEqual(expected, vertex_group.items.len);
    }
}

test "appendDraggedPiece uses dragged overlay group" {
    const board_rect = geometry.boardRectForExtent(800, 600);

    var groups = PieceVertexGroups.init();
    defer groups.deinit(std.testing.allocator);

    try appendDraggedPiece(&groups, board_rect, std.testing.allocator, piece.encode(.white, .queen), .{ 0.0, 0.0 });

    try expectGroupVertexCount(&groups, .white_queen, 0);
    try expectGroupVertexCount(&groups, .dragged_white_queen, 6);
}

test "appendPiecesGroupedFromBoardExceptWithGameOver renders sideways losing king" {
    var state = board.BoardState.empty();
    state.setSquare(63, piece.encode(.black, .king));
    state.setSquare(45, piece.encode(.white, .king));
    const board_rect = geometry.boardRectForExtent(800, 600);

    var groups = PieceVertexGroups.init();
    defer groups.deinit(std.testing.allocator);

    try appendPiecesGroupedFromBoardExceptWithGameOver(
        &groups,
        board_rect,
        std.testing.allocator,
        state,
        null,
        .{ .file = 7, .rank = 7 },
    );

    try expectGroupVertexCount(&groups, .white_king, 6);
    try expectGroupVertexCount(&groups, .black_king, 6);
}

test "appendPiecesGroupedFromBoard groups starting position vertices by piece" {
    const state = try fen.parseFen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
    const board_rect = geometry.boardRectForExtent(800, 600);

    var groups = PieceVertexGroups.init();
    defer groups.deinit(std.testing.allocator);

    try appendPiecesGroupedFromBoard(&groups, board_rect, std.testing.allocator, state);

    try expectGroupVertexCount(&groups, .white_pawn, 8 * 6);
    try expectGroupVertexCount(&groups, .white_knight, 2 * 6);
    try expectGroupVertexCount(&groups, .white_bishop, 2 * 6);
    try expectGroupVertexCount(&groups, .white_rook, 2 * 6);
    try expectGroupVertexCount(&groups, .white_queen, 1 * 6);
    try expectGroupVertexCount(&groups, .white_king, 1 * 6);

    try expectGroupVertexCount(&groups, .black_pawn, 8 * 6);
    try expectGroupVertexCount(&groups, .black_knight, 2 * 6);
    try expectGroupVertexCount(&groups, .black_bishop, 2 * 6);
    try expectGroupVertexCount(&groups, .black_rook, 2 * 6);
    try expectGroupVertexCount(&groups, .black_queen, 1 * 6);
    try expectGroupVertexCount(&groups, .black_king, 1 * 6);

    var total_vertices: usize = 0;
    for (groups.groups) |vertex_group| {
        total_vertices += vertex_group.items.len;
    }
    try std.testing.expectEqual(@as(usize, 32 * 6), total_vertices);
}