const std = @import("std");
const MetaInfo = @import("metainfo.zig");
const bencode = @import("bencode.zig");
const AnyWriter = @import("anywriter.zig");
const peerproto = @import("peer_protocol.zig");
const trackproto = @import("tracker_protocol.zig");

// TODO figure this out. It's not that important, I think, unless
// other clients have special handling for different patterns.
// Spec looks like a bit of a free-for-all here.
const peer_id = [20]u8{ 0x30, 0x30, 0x31, 0x31, 0x32, 0x32, 0x33, 0x33, 0x34, 0x34, 0x35, 0x35, 0x36, 0x36, 0x37, 0x37, 0x38, 0x38, 0x39, 0x39 };
// To make this a practical torrent app...
// Add some UI? or keep it a console app? I might keep it as a console app for now
const Args = struct {
    torrent_file: []const u8,
};

fn read_args() !Args {
    var args = std.process.args();
    if (!args.skip()) return error.NotEnoughArgs;
    const nxt = args.next() orelse return error.NotEnoughArgs;
    if (std.mem.eql(u8, nxt, "-h") or std.mem.eql(u8, nxt, "--help")) {
        try printUsage(false);
        return error.Help;
    }
    return .{
        .torrent_file = nxt,
    };
}

fn do_main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const a = gpa.allocator();

    const args = read_args() catch |e| switch (e) {
        error.Help => return,
        else => |err| return err,
    };

    const f = try std.fs.cwd().openFile(args.torrent_file, .{});
    defer f.close();
    var fr = f.reader();
    var mib = try bencode.bdecode(a, fr);
    defer mib.deinit(a);
    var mi = try MetaInfo.parse(a, mib);
    defer mi.deinit(a);
    const info_hash = try mi.info.hash(a);

    var c = std.http.Client{
        .allocator = a,
    };
    defer c.deinit();

    const url = try trackproto.trackerRequestUrl(a, info_hash, peer_id, mi.info.files[0].length, mi.announce);
    defer a.free(url);
    var res = try c.fetch(a, .{ .location = .{ .url = url } });
    defer res.deinit();
    if (res.status != .ok) {
        return error.TrackerHttpError;
    }

    var trb = try bencode.bdecodeBuf(a, res.body.?);
    defer trb.deinit(a);
    var tr = try trackproto.TrackerResp.parse(a, trb);
    defer tr.deinit(a);

    if (tr.peers.len == 0) {
        std.log.info("no peers", .{});
        return;
    }

    for (tr.peers) |peer| {
        std.log.info("peer: {}", .{peer});
    }

    // Handle peers, PoC we're just going to handle 1 peer and download everything from them very simplistically.
    const p = tr.peers[0];
    const file = mi.info.files[0];
    var ps = try std.net.tcpConnectToAddress(p);
    defer ps.close();
    var pw = ps.writer();
    var pr = ps.reader();
    var hs: peerproto.Handshake = .{
        .info_hash = info_hash,
        .peer_id = peer_id,
    };

    try hs.write(pw);
    var phs = try peerproto.Handshake.read(pr);
    std.log.info("peer at {} peer_id {s}", .{ p, std.fmt.fmtSliceHexLower(&phs.peer_id) });

    var bf = try peerproto.readMessage(a, pr, peerproto.Bitfield);
    _ = bf; // ignore it for now.
    try peerproto.Interested.write(pw);
    _ = try peerproto.readMessage(a, pr, peerproto.Unchoke);

    var of = try std.fs.cwd().createFile(file.name, .{});
    defer of.close();
    errdefer {
        // try to truncate the now-bad file...
        of.setEndPos(0) catch {};
    }

    // Read the piece into memory, we'll check the hash before it goes to disk...
    var piece_buf = try a.alloc(u8, mi.info.piece_length);
    defer a.free(piece_buf);

    for (0..mi.info.pieceCount()) |pi| {
        const piece_length = @min(mi.info.piece_length, file.length - (pi * mi.info.piece_length));
        var s1 = std.crypto.hash.Sha1.init(.{});

        // Send a request message for each 16KiB block of the first piece
        const blklen: u32 = 16 * 1024;
        var blkcount = try std.math.divCeil(u32, piece_length, blklen);
        for (0..blkcount) |i| {
            const begin = std.math.cast(u32, i * blklen).?;
            const len = @min(blklen, piece_length - begin);
            const req = peerproto.Request{
                .index = @intCast(pi),
                .begin = begin,
                .length = len,
            };
            std.log.info("Request {any}", .{req});
            try req.write(pw);
            var piece = try peerproto.readMessage(a, pr, peerproto.Piece);
            defer piece.deinit(a);
            if (piece.index != req.index) return error.ProtocolError;
            if (piece.begin != req.begin) return error.ProtocolError;
            if (piece.block.len != req.length) return error.ProtocolError;
            s1.update(piece.block);
            @memcpy(piece_buf[piece.begin .. piece.begin + piece.block.len], piece.block);
        }
        var ah = s1.finalResult();
        var ph0 = mi.info.pieceHash(pi).?;
        if (std.mem.eql(u8, &ah, &ph0)) {
            try of.writeAll(piece_buf[0..piece_length]);
        } else {
            return error.BadHash;
        }
    }
    std.log.info("fin", .{});
}

pub fn main() !void {
    do_main() catch |e| {
        try printUsage(true);
        return e;
    };
}

fn printUsage(err: bool) !void {
    var f = if (err) std.io.getStdErr() else std.io.getStdOut();
    try f.writeAll(
        \\ Usage: zbt TORRENTFILE 
        \\ Download the data specfied by TORRENTFILE to the current directory
        \\ 
    );
}

test {
    _ = bencode;
    _ = MetaInfo;
    _ = peerproto;
}