extern crate rand;
extern crate sha1;

pub mod base64;

pub use base64::B64Encode;
pub use base64::B64;
use rand::Rng;

macro_rules! get_fin {
    ($a: expr) => {
        ($a >> 7) & 0b1
    };
}

macro_rules! get_fin_enum {
    ($a: expr) => {
        match $a {
            0 => FIN::NotFinal,
            1 => FIN::Final,
            _ => panic!(),
        }
    };
}

macro_rules! get_fin_enum_value {
    ($a: expr) => {
        match $a {
            FIN::NotFinal => 0,
            FIN::Final => 1,
            _ => panic!(),
        }
    };
}

macro_rules! get_opcode {
    ($a: expr) => {
        $a & 0b1111
    };
}

macro_rules! get_opcode_enum {
    ($a: expr) => {
        match $a {
            0x0 => OPCODE::Fragment,
            0x1 => OPCODE::Txt,
            0x2 => OPCODE::Bin,
            0x3 => OPCODE::Rsv10,
            0x4 => OPCODE::Rsv11,
            0x5 => OPCODE::Rsv12,
            0x6 => OPCODE::Rsv13,
            0x7 => OPCODE::Rsv14,
            0x8 => OPCODE::Close,
            0x9 => OPCODE::Ping,
            0xA => OPCODE::Pong,
            0xB => OPCODE::Rsv20,
            0xC => OPCODE::Rsv21,
            0xD => OPCODE::Rsv22,
            0xE => OPCODE::Rsv23,
            0xF => OPCODE::Rsv24,
            _ => panic!(),
        }
    };
}

macro_rules! get_opcode_enum_value {
    ($a: expr) => {
        match $a {
            OPCODE::Fragment => 0x0,
            OPCODE::Txt => 0x1,
            OPCODE::Bin => 0x2,
            OPCODE::Rsv10 => 0x3,
            OPCODE::Rsv11 => 0x4,
            OPCODE::Rsv12 => 0x5,
            OPCODE::Rsv13 => 0x6,
            OPCODE::Rsv14 => 0x7,
            OPCODE::Close => 0x8,
            OPCODE::Ping => 0x9,
            OPCODE::Pong => 0xA,
            OPCODE::Rsv20 => 0xB,
            OPCODE::Rsv21 => 0xC,
            OPCODE::Rsv22 => 0xD,
            OPCODE::Rsv23 => 0xE,
            OPCODE::Rsv24 => 0xF,
            _ => panic!(),
        }
    };
}

macro_rules! get_payload_length {
    ($a: expr) => {
        ($a & 0b_0111_1111) as usize
    };
}

macro_rules! decode_payload_length_16 {
    ($a: expr, $b: expr) => {
        (($a as usize) << 8 & ($b as usize))
    };
}

macro_rules! decode_payload_length_64 {
    ($a: expr, $b: expr, $c: expr, $d: expr, $e: expr, $f: expr, $g: expr, $h: expr) => { (
        ($a as usize) << 56 &
        ($b as usize) << 48 &
        ($c as usize) << 40 &
        ($d as usize) << 32 &
        ($e as usize) << 24 &
        ($f as usize) << 16 &
        ($g as usize) << 8 &
        ($h as usize)
    ) }
}

macro_rules! is_masked {
    ($a: expr) => {
        $a >> 7 == 1
    };
}

#[derive(Debug, PartialEq)]
pub enum OPCODE {
    Fragment,
    Txt,
    Bin,
    Rsv10,
    Rsv11,
    Rsv12,
    Rsv13,
    Rsv14,
    Close,
    Ping,
    Pong,
    Rsv20,
    Rsv21,
    Rsv22,
    Rsv23,
    Rsv24,
}

#[derive(Debug, PartialEq)]
pub enum FIN {
    NotFinal,
    Final,
}

pub fn gen_mask() -> [u8; 4] {
    let mut rng = rand::thread_rng();
    let mut vec = [0u8; 4];
    for x in vec.iter_mut() {
        *x = rng.gen();
    }
    vec
}

pub fn gen_sec_websocket_key() -> String {
    let mut rng = rand::thread_rng();
    let mut vec = [0u8; 16];
    for x in vec.iter_mut() {
        *x = rng.gen();
    }
    B64::<String>::encode(&vec)
}

pub fn get_sec_websocket_accept(sec_web_socket_key: &str) -> String {
    let mut hash = sha1::Sha1::new();
    hash.update(
        (sec_web_socket_key.to_owned() + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").as_bytes(),
    );
    B64::<String>::encode(&hash.digest().bytes())
}

#[derive(Debug)]
pub struct Frame {
    fin_rsv_opcode: u8,
    mask_payload_length: u8,
    payload_length: Vec<u8>,
    mask: Option<[u8; 4]>,
    payload: Vec<u8>,
}

impl Frame {
    pub fn new(fin: FIN, opcode: OPCODE, data: &[u8]) -> Frame {
        let fin_rsv_opcode: u8 = (get_fin_enum_value!(fin) << 7) | (get_opcode_enum_value!(opcode));
        let payload_length: Vec<u8>;

        let mask_payload_length: u8 = match data.len() {
            len if len > 125 && len <= (u16::MAX as usize) => {
                payload_length = vec![(len >> 8 & 0b_1111_1111) as u8, (len & 0b_1111_1111) as u8];
                126u8
            }
            len if len > (u16::MAX as usize) => {
                payload_length = vec![
                    (len >> 56 & 0b_1111_1111) as u8,
                    (len >> 48 & 0b_1111_1111) as u8,
                    (len >> 40 & 0b_1111_1111) as u8,
                    (len >> 32 & 0b_1111_1111) as u8,
                    (len >> 24 & 0b_1111_1111) as u8,
                    (len >> 16 & 0b_1111_1111) as u8,
                    (len >> 8 & 0b_1111_1111) as u8,
                    (len & 0b_1111_1111) as u8,
                ];
                127u8
            }
            len => {
                payload_length = vec![];
                len as u8
            }
        };
        let payload = data.to_vec();

        Frame {
            fin_rsv_opcode,
            mask_payload_length,
            payload_length,
            mask: None,
            payload,
        }
    }
    pub fn fin(&self) -> FIN {
        get_fin_enum!(get_fin!(self.fin_rsv_opcode))
    }
    pub fn opcode(&self) -> OPCODE {
        get_opcode_enum!(get_opcode!(self.fin_rsv_opcode))
    }
    pub fn masking(&mut self, mask: &[u8; 4]) {
        if let Some(mask_) = self.mask {
            if mask_ != *mask {
                self.unmasking();
                self.masking(mask);
            }
        } else {
            for i in 0..self.payload.len() {
                self.payload[i] ^= mask[i % 4];
            }
            self.mask = Some(*mask);
            self.mask_payload_length ^= 0b_1000_0000;
        }
    }
    pub fn unmasking(&mut self) {
        if let Some(mask) = self.mask {
            for i in 0..self.payload.len() {
                self.payload[i] ^= mask[i % 4];
            }
            self.mask = None;
            self.mask_payload_length &= 0b_0111_1111;
        }
    }
    pub fn to_vec(&self) -> Vec<u8> {
        let mut vec = Vec::<u8>::with_capacity(self.payload.len() + 14);

        vec.push(self.fin_rsv_opcode);
        vec.push(self.mask_payload_length);

        vec.extend_from_slice(self.payload_length.as_slice());

        if let Some(mask) = self.mask {
            vec.extend_from_slice(&mask);
        }

        vec.extend_from_slice(self.payload.as_slice());

        vec
    }
}

pub fn decode_frame(v: &[u8], mut index: usize) -> Frame {
    let fin_rsv_opcode = v[index];
    let mask_payload_length = v[index + 1];

    index += 2;

    let mut temp_payload_length = get_payload_length!(mask_payload_length);
    let payload_length = match temp_payload_length {
        126 => {
            index += 2;
            temp_payload_length = decode_payload_length_16!(v[2], v[3]);
            vec![v[2], v[3]]
        }
        127 => {
            index += 8;
            temp_payload_length =
                decode_payload_length_64!(v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9]);
            vec![v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9]]
        }
        _ => {
            vec![]
        }
    };

    let mask: Option<[u8; 4]> = if is_masked!(mask_payload_length) {
        index += 4;
        Some([v[index - 4], v[index - 3], v[index - 2], v[index - 1]])
    } else {
        None
    };

    let payload = v[index..(index + temp_payload_length)].to_vec();

    index += temp_payload_length;

    Frame {
        fin_rsv_opcode,
        mask_payload_length,
        payload_length,
        mask,
        payload,
    }
}