const std = @import("std");
const builtin = @import("builtin");

pub fn readMessage(a: std.mem.Allocator, reader: anytype, comptime msgType: type) !msgType {
    var nka: u32 = 0;
    while (nka < 1000) {
        var len = try reader.readInt(u32, .big);
        // All messages except Keepalive start with a single byte message type.
        // Skip keepalive messages, we don't care (unless you're spamming them)
        if (len == 0) {
            nka += 1; 
            continue;
        }
        var mt = try reader.readByte();
        if (mt != msgType.Tag) return error.ProtocolError;
        return try msgType.read(a, len-1, reader);
    } else {
        return error.ProtocolError;
    }
}

// When you're expecting several possible messages.
// 'Expected' should be a tagged union of message types you are expecting.
pub fn readAnyMessage(a: std.mem.Allocator, reader: anytype, comptime Expected: type) !Expected {
    var nka: u32 = 0;
    while (nka < 1000) {
        var len = try reader.readInt(u32, .big);
        // All messages except Keepalive start with a single byte message type.
        // Skip keepalive messages, we don't care (unless you're spamming them)
        if (len == 0) {
            nka += 1; 
            continue;
        }
        var mt = try reader.readByte();
        inline for (@typeInfo(Expected).Union.fields) |field| {
            const msgType = field.type;
            if (msgType.Tag == mt) {
                return @unionInit(Expected, field.name, try msgType.read(a, len-1, reader));
            }
        } else {
            return error.ProtocolError;
        }
    } else {
        return error.ProtocolError;
    }
}

// Handshake message has a different structure to the rest. And it's only read once per connection.
pub const Handshake = struct {
    info_hash: [20]u8,
    peer_id: [20]u8,

    pub fn read(reader: anytype) !Handshake {
        var msg = [_]u8{0} ** 68;
        try reader.readNoEof(&msg);
        if (msg[0] != 19) return error.ProtocolError;
        if (!std.mem.eql(u8, msg[1..20], "BitTorrent protocol")) return error.ProtocolError;
        //if (!std.mem.allEqual(u8, msg[20..28], 0)) return error.ProtocolError;
        var res: Handshake = undefined;
        @memcpy(&res.info_hash, msg[28..48]);
        @memcpy(&res.peer_id, msg[48..68]);
        return res;
    }

    pub fn write(self: Handshake, writer: anytype) !void {
        try writer.writeByte(19);
        try writer.writeAll("BitTorrent protocol");
        try writer.writeByteNTimes(0, 8);
        try writer.writeAll(&self.info_hash);
        try writer.writeAll(&self.peer_id);
    }
};


pub const Unchoke = struct {
    pub const Tag: u8 = 1;
    pub fn read(a: std.mem.Allocator, len: usize, reader: anytype) !Unchoke {
        _ = a;
        _ = reader;
        if (len != 0) return error.ProtocolError;
        return .{};
    }
};

pub const Interested = struct {
    pub const Tag: u8 = 2;
    pub fn read(a: std.mem.Allocator, len: usize, reader: anytype) !Interested {
        _ = a;
        _ = reader;
        if (len != 0) return error.ProtocolError;
        return .{};   
    }
    pub fn write(writer: anytype) !void {
        try writer.writeInt(u32, 1, .big);
        try writer.writeInt(u8, Tag, .big);
    }
};

pub const Bitfield = struct {
    pub const Tag: u8 = 5;
    pub fn read(a: std.mem.Allocator, len: usize, reader: anytype) !Bitfield {
        // TODO actually read this message and do something useful.
        _ = a;
        try reader.skipBytes(len, .{}); // 
        return .{};
    }
};

pub const Request = struct {
    pub const Tag: u8 = 6;
    index: u32, 
    begin: u32, 
    length: u32,
    pub fn write(self: Request, writer: anytype) !void {
        try writer.writeInt(u32, 13, .big);
        try writer.writeInt(u8, Tag, .big);
        try writer.writeInt(u32,self.index, .big);
        try writer.writeInt(u32,self.begin, .big);
        try writer.writeInt(u32,self.length, .big);
    }
};

pub const Piece = struct {
    pub const Tag: u8 = 7;
    index: u32,
    begin: u32,
    block: []const u8,
    pub fn read(a: std.mem.Allocator, len: usize, reader: anytype) !Piece {
        if (len <= 8) {
            std.log.err("Piece#read len {}", .{len});
            return error.ProtocolError;
        }
        var ix = try reader.readInt(u32, .big);
        var be = try reader.readInt(u32, .big);
        var bl = try a.alloc(u8, len-8);
        errdefer a.free(bl);
        try reader.readNoEof(bl);
        return .{
            .index = ix, 
            .begin = be,
            .block = bl,
        };
    }
    pub fn deinit(self: *Piece, a: std.mem.Allocator) void {
        a.free(self.block);
    }
};

test "read any" {
    const a = std.testing.allocator;
    var fbs = std.io.fixedBufferStream(&[_]u8{0, 0, 0, 1, Unchoke.Tag});
    var r = fbs.reader();
    const T = union(enum) {u: Unchoke, i: Interested, b: Bitfield};
    var msg: T = try readAnyMessage(a, r, T);
    try std.testing.expect(msg == .u);
}