const std = @import("std");

const shapeModule = @import("shape.zig");
const inputModule = @import("input.zig");
const areaModule = @import("area.zig");

const Shape = shapeModule.Shape;
const Area = areaModule.Area;
const kShapeSize = shapeModule.kShapeSize;
const readInputIntoString = inputModule.readInputIntoString;
const parseInput = inputModule.parseInput;

const Coord = struct {
    row: u8,
    col: u8,
};

pub fn getEligibleCoords(allocator: std.mem.Allocator,
                         area: Area,
                         eligibleCoords: *std.ArrayList(Coord)) !void {
    eligibleCoords.clearRetainingCapacity();

    for (0 .. area.rows) |row| {
        for (0 .. area.cols) |col|  {
            if (area.getCell(@intCast(row), @intCast(col)) == '.') {
                var shouldAdd = false;
                if (row > 0 and area.getCell(@intCast(row-1), @intCast(col)) != '.') {
                    shouldAdd = true;
                } else if (col > 0 and area.getCell(@intCast(row), @intCast(col-1)) != '.') {
                    shouldAdd = true;
                } else if (row < area.rows - 1 and area.getCell(@intCast(row+1), @intCast(col)) != '.') {
                    shouldAdd = true;
                } else if (col < area.cols - 1 and area.getCell(@intCast(row), @intCast(col+1)) != '.') {
                    shouldAdd = true;
                }

                if (shouldAdd) {
                    var newEligibleCoord = try eligibleCoords.addOne(allocator);
                    newEligibleCoord.row = @intCast(row);
                    newEligibleCoord.col = @intCast(col);
                }
            }
        }
    }

    if (eligibleCoords.items.len == 0 and area.getCell(0, 0) == '.') {
        var newEligibleCoord = try eligibleCoords.addOne(allocator);
        newEligibleCoord.row = 0;
        newEligibleCoord.col = 0;
    }
}

pub fn calculateShapeSizes(allocator: std.mem.Allocator, shapeTransforms: []const []const Shape) ![]usize{
    var shapeSizes = try allocator.alloc(usize, shapeTransforms.len);
    for (0..shapeTransforms.len) |i| {
        shapeSizes[i] = 0;
        for (shapeTransforms[i][0].data) |d| {
            if (d != '.') {
                shapeSizes[i] += 1;
            }
        }
    }
    return shapeSizes;
}

pub fn calcShapesTotalSize(shapeCounts: []u8, shapeSizes: []usize) usize {
    var totalSize: usize = 0;

    for (0..shapeCounts.len) |i| {
        totalSize += shapeSizes[i] * shapeCounts[i];
    }
    return totalSize;
}


pub fn canContainShapes(allocator: std.mem.Allocator, area: Area,
                        shapeCounts: []u8, shapeTransforms: []const []const Shape,
                        shapeSizes: []usize) !bool {
    // try std.fs.File.stdout().writeAll("canContainShapes: [");
    // for (shapeCounts) |shapeCount| {
    //     try std.fs.File.stdout().writeAll(
    //       try std.fmt.allocPrint(allocator, " {d} ", .{shapeCount}));
    // }
    // try std.fs.File.stdout().writeAll("]\n");

    if (calcShapesTotalSize(shapeCounts, shapeSizes) > area.totalFreeSize()) {
        return false;
    }

    var eligibleCoords = std.ArrayList(Coord).empty;
    defer eligibleCoords.deinit(allocator);
    
    var allEmpty = true;
    for (shapeCounts) |shapeCount| {
        if (shapeCount != 0) {
            allEmpty = false;
            break;
        }
    }
    if (allEmpty) {
        try area.print(allocator);
        return true;
    }

    for (0 .. shapeCounts.len) |shapeIndex| {
        if (shapeCounts[shapeIndex] == 0) {
            continue;
        }

        for (0..shapeTransforms[shapeIndex].len) |shapeVariant| {
            try getEligibleCoords(allocator, area, &eligibleCoords);

            for (eligibleCoords.items) |coord| {
                const row = coord.row;
                const col = coord.col;

                if (try area.appendShape(
                      allocator,
                      &shapeTransforms[shapeIndex][shapeVariant],
                      @intCast(row), @intCast(col), '#')) |newArea| {
                    defer newArea.deinit();

                    // try std.fs.File.stdout().writeAll("Step: \n");
                    // try newArea.print(allocator);
                    // try std.fs.File.stdout().writeAll("\n");

                    var newShapeCounts = try allocator.alloc(u8, shapeCounts.len);
                    defer allocator.free(newShapeCounts);
                    std.mem.copyForwards(u8, newShapeCounts, shapeCounts);
                    newShapeCounts[shapeIndex] -= 1;

                    if (try canContainShapes(allocator, newArea, newShapeCounts,
                                             shapeTransforms, shapeSizes)) {
                        return true;
                    }
                }
            }
        }
    }

    return false;
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    const inputBuffer = try readInputIntoString(allocator, "input.txt");
    const input = try parseInput(allocator, inputBuffer);
    
    for (input.areas) |area| {
         try std.fs.File.stdout().writeAll(
           try std.fmt.allocPrint(allocator, "area: rows={d}, cols={d}, shapes=[ ", .{area.rows, area.cols}));
         for (area.shapesCount) |shapeCount| {
             try std.fs.File.stdout().writeAll(
               try std.fmt.allocPrint(allocator, "{d}, ", .{shapeCount}));
         }
         try std.fs.File.stdout().writeAll("]\n");
    }

    var shapeTransforms: [][]Shape = try allocator.alloc([]Shape, input.shapes.len);
    for (0..input.shapes.len) |i| {
        var allTransforms = try input.shapes[i].getAllTransforms(allocator);
        defer allTransforms.deinit();

        shapeTransforms[i] = try allocator.alloc(Shape, allTransforms.count());

        var j: usize = 0;
        var it = allTransforms.keyIterator();
        while (it.next()) |shapeTransform| {
            shapeTransforms[i][j] = shapeTransform.*;
            j += 1;
        }
    }

    const shapeSizes = try calculateShapeSizes(allocator, shapeTransforms);
    defer allocator.free(shapeSizes);

    var counter: usize = 0;
    for (0..input.areas.len) |i| {
        const area = try Area.init(allocator, input.areas[i].rows, input.areas[i].cols);
        defer area.deinit();
        
        if(try canContainShapes(allocator, area, input.areas[i].shapesCount, shapeTransforms, shapeSizes)) {
            try std.fs.File.stdout().writeAll("Can contain shapes!\n");
            counter += 1;
        }
        else {
            try std.fs.File.stdout().writeAll("Can NOT contain shapes!\n");
        }
    }

    try std.fs.File.stdout().writeAll(
      try std.fmt.allocPrint(allocator, "Answer: {d}\n", .{counter}));
}