const std = @import("std");
const rand = std.crypto.random;

const markov_data = @embedFile("markov.bin");

pub const DataPoint = struct {
    char: u8,
    prob: f32,

    pub fn desc(context: void, a: DataPoint, b: DataPoint) bool {
        _ = context;
        return a.prob > b.prob;
    }
};

pub const MarkovChain = struct {
    allocator: std.mem.Allocator,
    map: std.AutoHashMap(u8, []DataPoint),

    pub fn fromFile(path: []const u8, allocator: std.mem.Allocator) !MarkovChain {
        var markovBinFile = try std.fs.cwd().openFile(path, .{ .mode = .read_only });
        const reader = markovBinFile.reader();

        return try fromReader(&reader, allocator);
    }

    pub fn fromMemory(data: []const u8, allocator: std.mem.Allocator) !MarkovChain {
        var stream = std.io.fixedBufferStream(data);
        const reader = stream.reader();

        return try fromReader(reader, allocator);
    }

    fn fromReader(reader: anytype, allocator: std.mem.Allocator) !MarkovChain {
        var self = MarkovChain{
            .allocator = allocator,
            .map = std.AutoHashMap(u8, []DataPoint).init(allocator),
        };

        for (0..256) |prevChar| {
            const cnt = try reader.readInt(u8, .little);
            var nextChars: []DataPoint = try self.allocator.alloc(DataPoint, cnt);
            for (0..cnt) |i| {
                const nextByte = try reader.readByte();
                const prob: f32 = @bitCast(try reader.readInt(u32, .little));
                nextChars[i] = DataPoint{
                    .char = nextByte,
                    .prob = prob,
                };
            }
            try self.map.put(@as(u8, @intCast(prevChar)), nextChars);
        }

        return self;
    }

    pub fn deinit(self: *MarkovChain) void {
        var iter = self.map.iterator();
        while (iter.next()) |entry| {
            self.allocator.free(entry.value_ptr.*);
        }
        self.map.deinit();
    }

    pub fn generate(self: *MarkovChain, size: u8, allocator: std.mem.Allocator) ![]u8 {
        var result = try allocator.alloc(u8, size);
        var previous: u8 = 0;
        for (0..size) |i| {
            const choices = self.map.get(previous).?;
            const randFloat = rand.float(f32);
            var cumul: f32 = 0;
            for (0..choices.len) |j| {
                cumul += choices[j].prob;
                if (randFloat < cumul) {
                    result[i] = choices[j].char;
                    previous = choices[j].char;
                    break;
                }
            }
        }
        return result;
    }
};

test "basic test from file" {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();
    defer {
        _ = gpa.deinit();
    }

    var markov = try MarkovChain.fromFile("markov.bin", allocator);
    defer markov.deinit();

    for (0..6) |_| {
        const randName = try markov.generate(8, allocator);
        defer allocator.free(randName);
        std.debug.print("generated from file: {s}\n", .{randName});
    }
}

test "basic test from embedded file" {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();
    defer {
        _ = gpa.deinit();
    }

    var markov = try MarkovChain.fromMemory(markov_data, allocator);
    defer markov.deinit();

    for (0..6) |_| {
        const randName = try markov.generate(8, allocator);
        defer allocator.free(randName);
        std.debug.print("generated from memory: {s}\n", .{randName});
    }
}