const std = @import("std");
const mecha = @import("mecha");
const print = std.debug.print;

pub fn solve(part: []u8, buffer: []u8, allocator: std.mem.Allocator) !void {
    if (std.mem.eql(u8, part, "1")) {
        try part1(buffer, allocator);
    } else {
        try part2(buffer, allocator);
    }
}

fn amount(line: []const u8, allocator: std.mem.Allocator) usize {
    const match = mecha.oneOf(.{ mecha.string("XMAS"), mecha.string("SAMX") });

    var rest = line;
    var output: usize = 0;
    while (rest.len > 0) {
        _ = match.parse(allocator, rest) catch {
            rest = rest[1..];
            continue;
        };
        rest = rest[3..];
        output += 1;
    }
    return output;
}

fn diagonal(list: []const []const u8, startr: usize, startc: usize, dirr: i32, dirc: i32, allocator: std.mem.Allocator) !usize {
    var line = std.ArrayList(u8).init(allocator);
    defer line.deinit();

    const rows: i32 = @intCast(list.len);
    const width: i32 = @intCast(list[0].len);

    var r: i32 = @intCast(startr);
    var c: i32 = @intCast(startc);
    while (r >= 0 and c >= 0 and r < rows and c < width) {
        try line.append(list[@intCast(r)][@intCast(c)]);

        r += dirr;
        c += dirc;
    }
    return amount(line.items, allocator);
}

fn part1(buffer: []const u8, allocator: std.mem.Allocator) !void {
    var lines = std.mem.splitScalar(u8, buffer, '\n');

    var list = std.ArrayList([]const u8).init(allocator);
    defer list.deinit();

    var sum: usize = 0;

    while (lines.next()) |line| {
        if (line.len > 0) {
            try list.append(line);
            sum += amount(line, allocator); // HORIZONTAL
        }
    }

    for (0..list.items.len) |r| {
        sum += try diagonal(list.items, r, 0, -1, 1, allocator); // UP-RIGHT
        sum += try diagonal(list.items, r, 0, 1, 1, allocator); // DOWN-RIGHT
    }

    for (1..list.items[0].len) |c| {
        sum += try diagonal(list.items, 0, c, 1, 1, allocator); // DOWN-RIGHT
        sum += try diagonal(list.items, list.items.len - 1, c, -1, 1, allocator); // UP-RIGHT
    }

    for (0..list.items[0].len) |c| { // VERTICAL
        sum += try diagonal(list.items, 0, c, 1, 0, allocator);
    }

    print("{d}\n", .{sum});
}

//------------------------------PART 2---------------------------------------

fn diagonal2(list: []const []const u8, startr: usize, startc: usize, dirr: i32, dirc: i32, allocator: std.mem.Allocator, hashmap: anytype) !usize {
    var line = std.ArrayList(u8).init(allocator);
    defer line.deinit();

    const rows: i32 = @intCast(list.len);
    const width: i32 = @intCast(list[0].len);
    var output: usize = 0;

    var r: i32 = @intCast(startr);
    var c: i32 = @intCast(startc);
    while (r >= 0 and c >= 0 and r < rows and c < width) {
        try line.append(list[@intCast(r)][@intCast(c)]);

        if (line.items.len >= 3 and (std.mem.eql(u8, line.items[line.items.len - 3 ..], "MAS") or std.mem.eql(u8, line.items[line.items.len - 3 ..], "SAM"))) {
            if (hashmap.contains((r - dirr) * rows + (c - dirc))) {
                output += 1;
            } else {
                try hashmap.put((r - dirr) * rows + (c - dirc), 1);
            }
        }

        r += dirr;
        c += dirc;
    }
    return output;
}

fn part2(buffer: []const u8, allocator: std.mem.Allocator) !void {
    var map = std.AutoHashMap(i32, usize).init(allocator);
    defer map.deinit();

    var lines = std.mem.splitScalar(u8, buffer, '\n');

    var list = std.ArrayList([]const u8).init(allocator);
    defer list.deinit();

    var sum: usize = 0;

    while (lines.next()) |line| {
        if (line.len > 0) {
            try list.append(line);
        }
    }

    for (0..list.items.len) |r| {
        sum += try diagonal2(list.items, r, 0, -1, 1, allocator, &map); // UP-RIGHT
        sum += try diagonal2(list.items, r, 0, 1, 1, allocator, &map); // DOWN-RIGHT
    }

    for (1..list.items[0].len) |c| {
        sum += try diagonal2(list.items, 0, c, 1, 1, allocator, &map); // DOWN-RIGHT
        sum += try diagonal2(list.items, list.items.len - 1, c, -1, 1, allocator, &map); // UP-RIGHT
    }

    print("{d}\n", .{sum});
}