const std = @import("std");
const mem = std.mem;
const stdout = std.io.getStdOut().writer();
const stderr = std.io.getStdErr().writer();

const utils = @import("utils.zig");
const print = utils.print;
const exit = utils.exit;
const pErr = utils.pErr;

const walkers = @import("walker.zig");
const BitWalker = walkers.BitWalker;
const BitWalkerUint = walkers.BitWalkerUint;

fn usage() void {
    pErr("reader <path>", .{});
    exit("", .{}, 1);
}

fn HuffmanGraph(comptime valueType: type) type {
    return struct {
        const Node = struct {
            const NodeSelf = @This();

            left: ?*NodeSelf,
            right: ?*NodeSelf,
            value: ?valueType,

            allocator: mem.Allocator,

            fn init(allocator: mem.Allocator) !*NodeSelf {
                var node = try allocator.create(NodeSelf);
                node.left = null;
                node.right = null;
                node.value = null;
                node.allocator = allocator;

                return node;
            }

            fn deinit(self: *NodeSelf) void {
                if (self.left) |left| {
                    left.deinit();
                }
                if (self.right) |right| {
                    right.deinit();
                }
                self.allocator.destroy(self);
            }

            fn print(self: *NodeSelf, curDepth: usize, targetDepth: usize) void {
                if (curDepth != targetDepth) {
                    if (self.left) |l| {
                        l.print(curDepth + 1, targetDepth);
                    } else {
                        utils.printf(" , ", .{});
                    }
                    if (self.right) |r| {
                        r.print(curDepth + 1, targetDepth);
                    } else {
                        utils.printf(" . ", .{});
                    }
                    return;
                }

                if (self.value) |v| {
                    utils.printf(" {any} ", .{v});
                } else {
                    utils.printf(" _ ", .{});
                }
            }

            fn depth(self: *NodeSelf) usize {
                var d: usize = 0;
                if (self.right) |r| {
                    d = r.depth();
                }
                if (self.left) |l| {
                    d = std.math.max(d, l.depth());
                }
                return d + 1;
            }
        };

        root: *Node,
        allocator: mem.Allocator,
        stored_depth: ?usize = null,

        iterNode: *Node,

        const Self = @This();

        fn init(allocator: mem.Allocator) !Self {
            var root = try Node.init(allocator);
            return Self{
                .root = root,
                .iterNode = root,
                .allocator = allocator,
            };
        }

        fn print(self: *Self) void {
            for (0..(self.depth() + 1)) |i| {
                self.root.print(0, i);
                utils.printf("\n", .{});
            }
        }

        fn addValue(self: *Self, code: u64, size: u8, value: valueType) !void {
            var walker = try BitWalkerUint(u64, true).init(code, size);

            var curNode: *Node = self.root;

            while (walker.walkBit()) |bit| {
                if (bit == 1) {
                    if (curNode.left) |nextNode| {
                        curNode = nextNode;
                    } else {
                        curNode.left = try Node.init(self.allocator);
                        curNode = curNode.left.?;
                    }
                } else {
                    if (curNode.right) |nextNode| {
                        curNode = nextNode;
                    } else {
                        curNode.right = try Node.init(self.allocator);
                        curNode = curNode.right.?;
                    }
                }
            }

            curNode.value = value;
            self.stored_depth = null;
        }

        fn depth(self: *Self) usize {
            if (self.stored_depth) |d| {
                return d;
            }

            self.stored_depth = self.root.depth() - 1;
            return self.stored_depth.?;
        }

        fn iter(self: *Self, bit: u1) anyerror!?valueType {
            var node = self.iterNode;

            var nextNode: ?*Node = null;

            if (bit == 1) {
                nextNode = node.left;
            } else {
                nextNode = node.right;
            }

            if (nextNode) |new_node| {
                if (new_node.value) |value| {
                    self.iterNode = self.root;
                    return value;
                } else {
                    self.iterNode = new_node;
                    return null;
                }
            }

            return error.InvalidBitSequence;
        }

        fn deinit(self: *Self) void {
            self.root.deinit();
        }
    };
}

const DynamicDecoder = struct {
    const Self = @This();

    len_to_read: usize,
    codes: [19]u64,
    walker: *BitWalker,
    allocator: mem.Allocator,
    graph: HuffmanGraph(u64),

    fn init(walker: *BitWalker, len_to_read: usize, allocator: mem.Allocator) !Self {
        const list: [19]u8 = .{ 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 };
        var lenList: [19]u3 = .{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

        var bl_count: [7]u3 = .{ 0, 0, 0, 0, 0, 0, 0 };
        var max: u8 = 0;

        for (0..len_to_read) |i| {
            var data: u3 = @intCast(u3, try walker.walk(3));
            lenList[i] = data;
            if (data == 0) {
                continue;
            }
            bl_count[data] += 1;
            if (data > max) {
                max = data;
            }
        }

        var next_code: [19]u64 = .{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
        var code: u64 = 0;
        for (1..(max + 1)) |bits| {
            code = @shlExact(code + bl_count[bits - 1], 1);
            next_code[bits] = code;
        }

        var codes: [19]u64 = .{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

        for (0..19) |n| {
            var len = lenList[n];
            if (len != 0) {
                codes[n] = next_code[len];
                next_code[len] += 1;
            }
        }

        var graph = try HuffmanGraph(u64).init(allocator);

        for (0..19) |i| {
            if (lenList[i] == 0)
                continue;
            try graph.addValue(codes[i], lenList[i], list[i]);
        }

        return .{
            .len_to_read = len_to_read,
            .codes = codes,
            .walker = walker,
            .allocator = allocator,
            .graph = graph,
        };
    }

    fn deinit(self: *Self) void {
        self.graph.deinit();
    }
};

const BlockData = struct {
    const Self = @This();

    bitw: *BitWalker,
    allocator: mem.Allocator,

    last: bool,
    blockType: u8,

    fn init(allocator: mem.Allocator, bitw: *BitWalker) !Self {
        return Self{
            .bitw = bitw,
            .allocator = allocator,
            .last = try bitw.walk(1) == 1,
            .blockType = try bitw.walk(2),
        };
    }

    fn uncompress(self: *Self) !void {
        if (self.blockType != 2) {
            return error.unsuported_block_type;
        }
        
        try self.dynamic_huffman;

    }

    fn dynamic_huffman() !void {
        var bitw = self.bitw;

        var number_of_literal_codes: u32 = @as(u32, try bitw.walk(5)) + 257;
        var number_of_dist_codes = try bitw.walk(5) + 1;
        var number_of_length_codes = try bitw.walk(4) + 4;

        print("number of literal codes: {}", .{number_of_literal_codes});
        print("number of dist codes: {}", .{number_of_dist_codes});
        print("number_of_length_codes: {}", .{number_of_length_codes});

        var dynamic_decoder = try DynamicDecoder.init(bitw, number_of_length_codes, self.allocator);
        defer dynamic_decoder.deinit();

        dynamic_decoder.graph.print();

        var code_len: usize = 0;
        while (code_len < number_of_literal_codes) {
            var decode_value = try dynamic_decoder.graph.iter(try bitw.bitWalk());
            while (decode_value == null) {
                decode_value = try dynamic_decoder.graph.iter(try bitw.bitWalk());
            }

            print("Test {any}", .{decode_value});

            if (decode_value.? == 16) {
                return error.not_implemented;
            } else if (decode_value.? == 17) {
                code_len += try bitw.walk(3);
            } else if (decode_value.? == 18) {
                code_len += try bitw.walk(7);
            } else {
                code_len += 1;
            }
        }

        return error.todo;
    }

    fn read_len_code() void {}
};

const LOCAL_FILE_HEADER_SIGNATURE = 0x04034b50;
const ZipFileHeader = struct {
    version: u16,
    general: u16,
    compression_method: u16,
    last_mod_time: u16,
    last_mod_date: u16,
    crc_32: u32,
    compressed_size: u32,
    uncompressed_size: u32,
    file_name_length: u16,
    extra_field_length: u16,

    file_name: []u8,
    extra_field: []u8,
    compressed_content: []u8,
    uncompressed_content: []u8,
    decompressed: bool,

    allocator: std.mem.Allocator,

    const Self = @This();

    fn init(allocator: std.mem.Allocator, reader: std.fs.File.Reader) !Self {
        if (try reader.readInt(u32, .Big) == LOCAL_FILE_HEADER_SIGNATURE) {
            return error.InvalidError;
        }

        var self = Self{
            .allocator = allocator,
            .version = try reader.readInt(u16, .Little),
            .general = try reader.readInt(u16, .Little),
            .compression_method = try reader.readInt(u16, .Little),
            .last_mod_time = try reader.readInt(u16, .Little),
            .last_mod_date = try reader.readInt(u16, .Little),
            .crc_32 = try reader.readInt(u32, .Little),
            .compressed_size = try reader.readInt(u32, .Little),
            .uncompressed_size = try reader.readInt(u32, .Little),
            .file_name_length = try reader.readInt(u16, .Little),
            .extra_field_length = try reader.readInt(u16, .Little),
            .file_name = undefined,
            .extra_field = undefined,
            .compressed_content = undefined,
            .uncompressed_content = undefined,
            .decompressed = false,
        };

        self.file_name = try allocator.alloc(u8, self.file_name_length);
        self.extra_field = try allocator.alloc(u8, self.extra_field_length);
        self.compressed_content = try allocator.alloc(u8, self.compressed_size);

        _ = try reader.read(self.file_name);
        _ = try reader.read(self.extra_field);
        _ = try reader.read(self.compressed_content);

        return self;
    }

    fn extract(self: *Self) !void {
        if (self.decompressed) {
            return error.AlreadyDecompressed;
        }

        if (self.compression_method == 0) {
            return error.uncompressed_file;
        }

        if (self.compression_method != 8) {
            return error.unsuported_compression_method;
        }

        self.uncompressed_content = try self.allocator.alloc(u8, self.uncompressed_size);
        errdefer self.allocator.free(self.uncompressed_content);

        var bitw = BitWalker.init(&self.compressed_content, false);

        var block = try BlockData.init(self.allocator, &bitw);

        if (block.last) {
            print("last block", .{});
        } else {
            print("not last block", .{});
        }

        try block.uncompress();

        self.decompressed = true;
    }

    fn deinit(self: *Self) void {
        self.allocator.free(self.file_name);
        self.allocator.free(self.extra_field);
        self.allocator.free(self.compressed_content);
        if (self.decompressed) {
            self.allocator.free(self.uncompressed_content);
        }
    }
};

pub fn main() !void {
    var args = std.process.args();

    var filePath: ?[]const u8 = null;

    // Skip the current path
    _ = args.next();
    while (args.next()) |arg| {
        if (filePath == null) {
            filePath = arg;
        } else {
            pErr("Invalid argument: {s}", .{arg});
            usage();
        }
    }

    if (filePath == null) {
        pErr("File path not provided. Please provide a path", .{});
        usage();
    }

    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();

    var allocator = gpa.allocator();

    var book_path = try std.fs.realpathAlloc(allocator, filePath.?);
    defer allocator.free(book_path);

    var file = std.fs.openFileAbsolute(book_path, .{}) catch |err| {
        exit("Please provide a file path! Error: {?}", .{err}, 1);
        return err;
    };
    defer file.close();

    var stat = try file.stat();

    if (stat.kind != .File) {
        exit("Please provide a valid file", .{}, 1);
    }

    var reader = file.reader();

    var first_file = try ZipFileHeader.init(allocator, reader);
    defer first_file.deinit();

    if (!std.mem.eql(u8, first_file.file_name, "mimetype")) {
        exit("Invalid file provided", .{}, 1);
    }

    if (!std.mem.startsWith(u8, first_file.compressed_content, "application/epub")) {
        exit("Invalid file provided", .{}, 1);
    }

    print("H: {}", .{first_file.compression_method});

    var second_file = try ZipFileHeader.init(allocator, reader);
    defer second_file.deinit();

    try second_file.extract();

    print("G: {s}", .{second_file.file_name});
    print("GI: {}", .{second_file.compression_method});
    print("xml stuff:\n{s}", .{second_file.compressed_content});
}