const std = @import("std");
const builtin = @import("builtin");
const http = @import("./lib.zig");
const Stream = @import("./server.zig").Stream;

const Opcode = enum(u4) {
    continuation = 0x0,
    text = 0x1,
    binary = 0x2,
    // 3-7 are non-control frames
    close = 0x8,
    ping = 0x9,
    pong = 0xa,
    // b-f are control frames
    _,

    fn isData(self: Opcode) bool {
        return !self.isControl();
    }

    fn isControl(self: Opcode) bool {
        return @enumToInt(self) & 0x8 == 0x8;
    }
};

pub fn handshake(alloc: std.mem.Allocator, req_headers: *const http.Fields, res: *http.Response) !Socket {
    const upgrade = req_headers.get("Upgrade") orelse return error.BadHandshake;
    const connection = req_headers.get("Connection") orelse return error.BadHandshake;
    if (std.ascii.indexOfIgnoreCase(upgrade, "websocket") == null) return error.BadHandshake;
    if (std.ascii.indexOfIgnoreCase(connection, "Upgrade") == null) return error.BadHandshake;

    const key_hdr = req_headers.get("Sec-WebSocket-Key") orelse return error.BadHandshake;
    if ((try std.base64.standard.Decoder.calcSizeForSlice(key_hdr)) != 16) return error.BadHandshake;
    var key: [16]u8 = undefined;
    std.base64.standard.Decoder.decode(&key, key_hdr) catch return error.BadHandshake;

    const version = req_headers.get("Sec-WebSocket-version") orelse return error.BadHandshake;
    if (!std.mem.eql(u8, "13", version)) return error.BadHandshake;

    var headers = http.Fields.init(alloc);
    defer headers.deinit();

    try headers.put("Upgrade", "websocket");
    try headers.put("Connection", "Upgrade");
    var response_key = std.mem.zeroes([60]u8);
    std.mem.copy(u8, &response_key, key_hdr);
    std.mem.copy(u8, response_key[key_hdr.len..], "258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
    const Sha1 = std.crypto.hash.Sha1;
    var hash: [Sha1.digest_length]u8 = undefined;
    Sha1.hash(response_key[0..], &hash, .{});
    var hash_encoded: [std.base64.standard.Encoder.calcSize(Sha1.digest_length)]u8 = undefined;
    _ = std.base64.standard.Encoder.encode(&hash_encoded, &hash);
    try headers.put("Sec-WebSocket-Accept", &hash_encoded);
    const stream = try res.upgrade(.switching_protocols, &headers);

    return Socket{ .stream = stream };
}

pub const CloseReason = enum(u16) {
    ok = 1000,
    going_away = 1001,
    protocol_error = 1002,
    unsupported_frame_type = 1003,
    // reserved = 1004,

    // technically these codes can be explicitly provided but it's against spec
    no_reason_provided = 1005,
    abnormal_closure = 1006,

    invalid_frame_data = 1007, // e.g. non UTF-8 data in a text type frame
    policy_violation = 1008,
    message_too_large = 1009,
    missing_extension = 1010,
    internal_server_error = 1011,

    // unused on server side but included here for completeness
    invalid_tls_handshake = 1015,

    _,
};

const FrameInfo = struct {
    is_final: bool,
    rsv: u3,
    opcode: Opcode,
    masking_key: ?[4]u8,
    len: usize,
};

fn frameReader(reader: anytype) !FrameReader(@TypeOf(reader)) {
    return FrameReader(@TypeOf(reader)).open(reader);
}

fn FrameReader(comptime R: type) type {
    return struct {
        const Self = @This();
        info: FrameInfo,

        bytes_read: usize = 0,
        mask_idx: u2 = 0,
        underlying: R,

        const Reader = std.io.Reader(*Self, anyerror, read);

        fn open(r: R) !Self {
            var hdr_buf: [2]u8 = undefined;
            try r.readNoEof(&hdr_buf);
            const is_final = hdr_buf[0] & 0b1000_0000 != 0;
            const rsv = @intCast(u3, (hdr_buf[0] & 0b0111_0000) >> 4);
            const opcode = @intToEnum(Opcode, (hdr_buf[0] & 0b1111));
            const is_masked = hdr_buf[1] & 0b1000_0000 != 0;
            const initial_len = @intCast(u7, (hdr_buf[1] & 0b0111_1111));

            const len: usize = if (initial_len < 126)
                initial_len
            else if (initial_len == 126)
                try r.readIntBig(u16)
            else
                try r.readIntBig(usize);

            const masking_key = if (is_masked)
                try r.readBytesNoEof(4)
            else
                null;

            return Self{
                .info = .{
                    .is_final = is_final,
                    .rsv = rsv,
                    .opcode = opcode,
                    .masking_key = masking_key,
                    .len = len,
                },
                .underlying = r,
            };
        }

        fn read(self: *Self, buf: []u8) !usize {
            var r = std.io.limitedReader(self.underlying, self.info.len - self.bytes_read);
            const count = try r.read(buf);
            if (self.info.masking_key) |mask| {
                for (buf[0..count]) |*ch| {
                    ch.* ^= mask[self.mask_idx];
                    self.mask_idx +%= 1;
                }
            }
            self.bytes_read += count;
            return count;
        }

        fn reader(self: *Self) Reader {
            return .{ .context = self };
        }

        fn close(self: *Self) void {
            self.reader().skipBytes(@as(usize, 0) -% 1, .{}) catch {};
        }
    };
}

fn writeFrame(writer: anytype, header: FrameInfo, buf: []const u8) !void {
    std.debug.assert(header.len == buf.len);

    const initial_len: u7 = if (header.len < 126)
        @intCast(u7, header.len)
    else if (std.math.cast(u16, header.len)) |_|
        @as(u7, 126)
    else
        @as(u7, 127);

    var hdr_buf = [2]u8{ 0, 0 };
    hdr_buf[0] |= if (header.is_final) @as(u8, 0b1000_0000) else 0;
    hdr_buf[0] |= @as(u8, header.rsv) << 4;
    hdr_buf[0] |= @enumToInt(header.opcode);
    hdr_buf[1] |= if (header.masking_key) |_| @as(u8, 0b1000_0000) else 0;
    hdr_buf[1] |= initial_len;
    try writer.writeAll(&hdr_buf);
    if (initial_len == 126)
        try writer.writeIntBig(u16, std.math.cast(u16, buf.len).?)
    else if (initial_len == 127)
        try writer.writeIntBig(usize, buf.len);

    if (header.masking_key) |key| try writer.writeAll(&key);

    const mask = header.masking_key orelse std.mem.zeroes([4]u8);
    var mask_idx: u2 = 0;
    for (buf) |ch| {
        try writer.writeByte(mask[mask_idx] ^ ch);
    }
}

pub const MessageReader = struct {
    opcode: Opcode,
    socket: *Socket,
    frame: FrameReader(Socket.Reader),

    pub const Reader = std.io.Reader(*MessageReader, anyerror, read);

    pub fn read(self: *MessageReader, buf: []u8) !usize {
        if (self.socket.closed) return error.SocketClosed;

        var count: usize = 0;
        while (count < buf.len) {
            const c = try self.frame.read(buf);
            count += c;

            if (c == 0) {
                self.frame.close();
                if (self.frame.info.is_final) break;

                self.frame = try self.socket.waitForDataFrame();
                if (self.frame.info.opcode != .continuation) {
                    //self.socket.close();
                    return error.InvalidFrame;
                }
            }
        }

        return count;
    }

    pub fn reader(self: *MessageReader) Reader {
        return .{ .context = self };
    }

    pub fn close(self: *MessageReader) void {
        self.reader().skipBytes(@as(usize, 0) -% 1, .{}) catch {};
    }
};

pub const MessageWriter = struct {
    allocator: std.mem.Allocator,
    socket: *Socket,
    opcode: Opcode,
    buf: []u8,

    cursor: usize = 0,
    is_first: bool = true,

    pub const WriteError = Socket.WriteError;
    pub const Writer = std.io.Writer(*MessageWriter, WriteError, write);

    pub fn writer(self: *MessageWriter) Writer {
        return .{ .context = self };
    }

    pub fn write(self: *MessageWriter, buf: []const u8) !usize {
        if (self.socket.closed) return error.SocketClosed;
        var count: usize = 0;
        while (count < buf.len) {
            if (self.cursor == self.buf.len) try self.flushFrame(false);

            const max_write = std.math.min(self.buf.len - self.cursor, buf.len - count);
            std.mem.copy(u8, self.buf[self.cursor..], buf[count .. count + max_write]);

            self.cursor += max_write;
            count += max_write;
        }

        return count;
    }

    pub fn close(self: *MessageWriter) void {
        self.allocator.free(self.buf);
    }

    pub fn finish(self: *MessageWriter) !void {
        if (self.socket.closed) return error.SocketClosed;
        try self.flushFrame(true);
    }

    fn flushFrame(self: *MessageWriter, is_final: bool) !void {
        if (self.socket.closed) return error.SocketClosed;
        const opcode = if (self.is_first) self.opcode else .continuation;
        self.is_first = false;
        const header = FrameInfo{
            .is_final = is_final,
            .rsv = 0,
            .opcode = opcode,
            .masking_key = null, // TODO: use a key on client-side
            .len = self.cursor,
        };
        std.log.debug("sending frame with size {}", .{header.len});
        try writeFrame(self.socket.writer(), header, self.buf[0..self.cursor]);
        self.cursor = 0;
    }
};

pub const Socket = struct {
    stream: Stream,
    closed: bool = false,

    const ReadError = Stream.ReadError || error{SocketClosed};
    const WriteError = Stream.WriteError || error{SocketClosed};
    const Reader = std.io.Reader(*Socket, ReadError, read);
    const Writer = std.io.Writer(*Socket, WriteError, write);

    fn writer(self: *Socket) Writer {
        return .{ .context = self };
    }

    fn reader(self: *Socket) Reader {
        return .{ .context = self };
    }

    fn read(self: *Socket, buf: []u8) !usize {
        if (self.closed) return error.SocketClosed;

        return self.stream.read(buf) catch |err| {
            self.closed = true;
            return err;
        };
    }

    fn write(self: *Socket, buf: []const u8) !usize {
        if (self.closed) return error.SocketClosed;

        return self.stream.write(buf) catch |err| {
            self.closed = true;
            return err;
        };
    }

    pub fn accept(self: *Socket) !MessageReader {
        if (self.closed) return error.SocketClosed;
        const frame = try self.waitForDataFrame();
        return MessageReader{ .opcode = frame.info.opcode, .socket = self, .frame = frame };
    }

    pub const MessageOptions = struct {
        buf_size: usize = 512,
    };
    pub fn openMessage(self: *Socket, opcode: Opcode, alloc: std.mem.Allocator, options: MessageOptions) !MessageWriter {
        if (self.closed) return error.SocketClosed;
        var buf = try alloc.alloc(u8, options.buf_size);
        return MessageWriter{ .socket = self, .allocator = alloc, .buf = buf, .opcode = opcode };
    }

    fn openFrame(self: *Socket) !FrameReader(Reader) {
        std.log.debug("waiting for frame header", .{});

        const frame = try frameReader(self.reader());

        std.log.debug("Received frame of type {any}, size {}", .{ frame.info.opcode, frame.info.len });

        return frame;
    }

    fn waitForDataFrame(self: *Socket) !FrameReader(Reader) {
        while (true) {
            var frame = try frameReader(self.reader());
            if (frame.info.opcode.isData()) return frame;

            // TODO: handle control frames
            frame.close();
        }
    }
};

fn pipe(writer: anytype, reader: anytype) !usize {
    var count: usize = 0;

    var buf: [64]u8 = undefined;
    while (true) {
        const c = try reader.read(&buf);
        if (c == 0) break;

        if (try writer.write(buf[0..c]) != c) return error.EndOfStream;
        count += c;
    }
    return count;
}