//! Variables captured at patch location
// Copyright (C) 2025  Nguyễn Gia Phong
//
// This file is part of taosc.
//
// Taosc is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Taosc is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with taosc.  If not, see <https://www.gnu.org/licenses/>.

const Allocator = std.mem.Allocator;
const Writer = std.io.Writer;
const assert = std.debug.assert;
const bytesAsSlice = std.mem.bytesAsSlice;
const cwd = std.fs.cwd;
const divCeil = std.math.divCeil;
const fields = std.meta.fields;
const std = @import("std");

pub const RegisterEnum = enum(u4) {
    rax, rbx, rcx, rdx, rsi, rdi, rsp, rbp,
    r8, r9, r10, r11, r12, r13, r14, r15,

    pub fn format(tag: RegisterEnum, writer: *Writer) Writer.Error!void {
        try writer.print("{s}", .{ @tagName(tag) });
    }
};
const Register = i64;
const Registers = [fields(RegisterEnum).len]Register;

pub const signed_integers = .{ i64, i32, i16, i8 };
const alignment = @alignOf(Register);
comptime {
    for (signed_integers) |Int|
        assert(alignment >= @alignOf(Int));
}

fn alignedSize(T: type, count: usize) !usize {
    return try divCeil(usize, @sizeOf(T) * count, @alignOf(T)) * @alignOf(T);
}

fn packedSize(T: type, container_size: usize, count: usize) !usize {
    return @divExact(container_size, @sizeOf(T)) * try alignedSize(T, count);
}

const Variables = @This();
bytes: []align(alignment) const u8,
samples: usize,

pub fn init(allocator: Allocator, stack_size: usize,
            path: []const u8) !Variables {
    var dir = try cwd().openDir(path, .{ .iterate = true });
    defer dir.close();
    var entries = dir.iterate();
    var count: usize = 0;
    while (try entries.next()) |entry| {
        assert(entry.kind == .file);
        count += 1;
    }
    var size = try packedSize(Register, @sizeOf(Registers), count);
    inline for (signed_integers) |Int|
        size += try packedSize(Int, stack_size, count);
    const bytes = try allocator.alignedAlloc(u8,
        .fromByteUnits(alignment), size);
    errdefer allocator.free(bytes);

    const file_size = @sizeOf(Registers) + stack_size;
    const buffer = try allocator.alignedAlloc(u8,
        .fromByteUnits(@alignOf(Registers)), file_size);
    defer allocator.free(buffer);
    const registers: *Registers = @ptrCast(buffer.ptr);
    entries.reset();
    for (0..count) |index| {
        const entry = try entries.next();
        const file = try dir.openFile(entry.?.name, .{});
        defer file.close();
        assert(try file.getEndPos() == file_size);
        assert(try file.read(buffer) == file_size);
        var offset: usize = 0;
        inline for (registers) |reg| {
            const slice = bytesAsSlice(Register, bytes[offset..]);
            slice[index] = reg;
            offset += try alignedSize(Register, count);
        }
        inline for (signed_integers) |Int| {
            const aligned_size = try alignedSize(Int, count);
            for (bytesAsSlice(Int, buffer[@sizeOf(Registers)..])) |value| {
                const slice = bytesAsSlice(Int, bytes[offset..]);
                slice[index] = value;
                offset += aligned_size;
            }
        }
    }
    return .{ .bytes = bytes, .samples = count };
}

pub fn deinit(variables: Variables, allocator: Allocator) void {
    allocator.free(variables.bytes);
}

pub fn register(variables: Variables, name: RegisterEnum) ![]const Register {
    const aligned_size = try alignedSize(Register, variables.samples);
    const offset = aligned_size * @intFromEnum(name);
    const slice = bytesAsSlice(Register, variables.bytes[offset..]);
    return @alignCast(slice[0..variables.samples]);
}