use nom::bits;
use nom::error::Error;
use nom::error::ErrorKind;
use nom::number;
use nom::sequence;
use nom::Err;
use nom::IResult;
use nom::Needed;

/******************************************************************************
 * Struct
 ******************************************************************************/

/*
 * TCP Header Format
 *
 *   0                   1                   2                   3
 *   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |          Source Port          |       Destination Port        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                        Sequence Number                        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Acknowledgment Number                      |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |  Data |           |U|A|P|R|S|F|                               |
 *  | Offset| Reserved  |R|C|S|S|Y|I|            Window             |
 *  |       |           |G|K|H|T|N|N|                               |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |           Checksum            |         Urgent Pointer        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Options                    |    Padding    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                             data                              |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

const TCP_OPTION_EOL: u8 = 0; // End of Options list
const TCP_OPTION_NOP: u8 = 1; // No operation
const TCP_OPTION_MSS: u8 = 2; // Maximum segment size
const TCP_OPTION_WSCALE: u8 = 3; // Window scale
const TCP_OPTION_SACK_PERMITTED: u8 = 4; // Selective acknowledgements permitted
const TCP_OPTION_SACK: u8 = 5; // Selective acknowledgment
const TCP_OPTION_TIMESTAMPS: u8 = 8; // Timestamps

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TcpOption {
    EOL,
    NOP,
    MSS {
        length: u8,
        mss: u16,
    },
    WSCALE {
        length: u8,
        shift_count: u8,
    },
    SACKPERMITTED,
    SACK {
        length: u8,
        left_edge: u32,
        right_edge: u32,
    },
    TIMESTAMPS {
        length: u8,
        ts_value: u32,
        ts_reply: u32,
    },
}

#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct TcpHeader {
    pub source_port: u16,
    pub dest_port: u16,

    pub seq_num: u32,
    pub ack_num: u32,

    pub data_offset: u8,
    pub reserved: u8,

    pub flag_urg: bool,
    pub flag_ack: bool,
    pub flag_psh: bool,
    pub flag_rst: bool,
    pub flag_syn: bool,
    pub flag_fin: bool,

    pub window: u16,
    pub checksum: u16,
    pub urgent_ptr: u16,
    pub options: Option<Vec<TcpOption>>,
}

/******************************************************************************
 * API
 ******************************************************************************/

fn offset_res_flags_decode(input: &[u8]) -> IResult<&[u8], (u8, u8, u8)> {
    bits::bits::<_, _, Error<_>, _, _>(sequence::tuple((
        bits::streaming::take(4u8),
        bits::streaming::take(6u8),
        bits::streaming::take(6u8),
    )))(input)
}

impl TcpOption {
    fn mss_decode(input: &[u8]) -> IResult<&[u8], TcpOption> {
        let (input, length) = number::streaming::be_u8(input)?;
        let (input, mss) = number::streaming::be_u16(input)?;
        Ok((input, TcpOption::MSS { length, mss }))
    }

    fn wscale_decode(input: &[u8]) -> IResult<&[u8], TcpOption> {
        let (input, length) = number::streaming::be_u8(input)?;
        let (input, shift_count) = number::streaming::be_u8(input)?;
        Ok((
            input,
            TcpOption::WSCALE {
                length,
                shift_count,
            },
        ))
    }

    fn sack_permitted_decode(input: &[u8]) -> IResult<&[u8], TcpOption> {
        let (input, _length) = number::streaming::be_u8(input)?;
        Ok((input, TcpOption::SACKPERMITTED))
    }

    fn sack_decode(input: &[u8]) -> IResult<&[u8], TcpOption> {
        let (input, length) = number::streaming::be_u8(input)?;
        let (input, left_edge) = number::streaming::be_u32(input)?;
        let (input, right_edge) = number::streaming::be_u32(input)?;
        Ok((
            input,
            TcpOption::SACK {
                length,
                left_edge,
                right_edge,
            },
        ))
    }

    fn timestamp_decode(input: &[u8]) -> IResult<&[u8], TcpOption> {
        let (input, length) = number::streaming::be_u8(input)?;
        let (input, ts_value) = number::streaming::be_u32(input)?;
        let (input, ts_reply) = number::streaming::be_u32(input)?;
        Ok((
            input,
            TcpOption::TIMESTAMPS {
                length,
                ts_value,
                ts_reply,
            },
        ))
    }

    fn decode(input: &[u8]) -> IResult<&[u8], TcpOption> {
        match number::streaming::be_u8(input)? {
            (input, TCP_OPTION_EOL) => Ok((input, TcpOption::EOL)),
            (input, TCP_OPTION_NOP) => Ok((input, TcpOption::NOP)),
            (input, TCP_OPTION_MSS) => TcpOption::mss_decode(input),
            (input, TCP_OPTION_WSCALE) => TcpOption::wscale_decode(input),
            (input, TCP_OPTION_SACK_PERMITTED) => TcpOption::sack_permitted_decode(input),
            (input, TCP_OPTION_SACK) => TcpOption::sack_decode(input),
            (input, TCP_OPTION_TIMESTAMPS) => TcpOption::timestamp_decode(input),
            (input, _other) => Err(Err::Failure(Error::new(input, ErrorKind::Switch))),
        }
    }
}

impl TcpHeader {
    fn fixed_header_decode(input: &[u8]) -> IResult<&[u8], TcpHeader> {
        let (input, source_port) = number::streaming::be_u16(input)?;
        let (input, dest_port) = number::streaming::be_u16(input)?;
        let (input, seq_num) = number::streaming::be_u32(input)?;
        let (input, ack_num) = number::streaming::be_u32(input)?;
        let (input, dataof_res_flags) = offset_res_flags_decode(input)?;
        let (input, window) = number::streaming::be_u16(input)?;
        let (input, checksum) = number::streaming::be_u16(input)?;
        let (input, urgent_ptr) = number::streaming::be_u16(input)?;

        Ok((
            input,
            TcpHeader {
                source_port,
                dest_port,
                seq_num,
                ack_num,
                data_offset: dataof_res_flags.0 * 4, // dataof_res_flags.0  * 32 / 8
                reserved: dataof_res_flags.1,
                flag_urg: dataof_res_flags.2 & 0b10_0000 == 0b10_0000,
                flag_ack: dataof_res_flags.2 & 0b01_0000 == 0b01_0000,
                flag_psh: dataof_res_flags.2 & 0b00_1000 == 0b00_1000,
                flag_rst: dataof_res_flags.2 & 0b00_0100 == 0b00_0100,
                flag_syn: dataof_res_flags.2 & 0b00_0010 == 0b00_0010,
                flag_fin: dataof_res_flags.2 & 0b00_0001 == 0b00_0001,
                window,
                checksum,
                urgent_ptr,
                options: None,
            },
        ))
    }

    fn options_decode(input: &[u8]) -> IResult<&[u8], Vec<TcpOption>> {
        let mut left = input;
        let mut options: Vec<TcpOption> = vec![];
        loop {
            match TcpOption::decode(left) {
                Ok((l, opt)) => {
                    left = l;
                    options.push(opt);

                    if left.len() <= 0 {
                        break;
                    }
                    if let TcpOption::EOL = opt {
                        break;
                    }
                }
                Err(e) => return Err(e),
            }
        }

        Ok((left, options))
    }

    pub fn decode(input: &[u8]) -> IResult<&[u8], TcpHeader> {
        match TcpHeader::fixed_header_decode(input) {
            Ok((left, mut header)) => {
                if header.data_offset > 20 {
                    let options_length = (header.data_offset - 20) as usize;
                    if options_length <= left.len() {
                        if let Ok((__left, options)) =
                            TcpHeader::options_decode(&left[0..options_length])
                        {
                            header.options = Some(options);
                            Ok((&left[options_length..], header))
                        } else {
                            println!("tcp header options parser error, skip tcp options");
                            Ok((&left[options_length..], header))
                        }
                    } else {
                        println!("tcp header options data not enough");
                        Err(Err::Incomplete(Needed::new(options_length - left.len())))
                    }
                } else {
                    Ok((left, header))
                }
            }
            err => err,
        }
    }
}

/******************************************************************************
 * TEST
 ******************************************************************************/

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::tcp::TcpOption::MSS;
    use crate::protocol::tcp::TcpOption::NOP;
    use crate::protocol::tcp::TcpOption::SACK;
    use crate::protocol::tcp::TcpOption::SACKPERMITTED;
    use crate::protocol::tcp::TcpOption::TIMESTAMPS;
    use crate::protocol::tcp::TcpOption::WSCALE;

    const LAST_SLICE: &'static [u8] = &[0xff];

    #[test]
    fn tcp_header_decode1() {
        /*
         * Transmission Control Protocol, Src Port: 50081, Dst Port: 443, Seq: 1, Ack: 1, Len: 1348
         *     Source Port: 50081
         *     Destination Port: 443
         *     [Stream index: 0]
         *     [Conversation completeness: Incomplete (8)]
         *     [TCP Segment Len: 1348]
         *     Sequence Number: 1    (relative sequence number)
         *     Sequence Number (raw): 1522577104
         *     [Next Sequence Number: 1349    (relative sequence number)]
         *     Acknowledgment Number: 1    (relative ack number)
         *     Acknowledgment number (raw): 3419365570
         *     1000 .... = Header Length: 32 bytes (8)
         *     Flags: 0x010 (ACK)
         *         000. .... .... = Reserved: Not set
         *         ...0 .... .... = Nonce: Not set
         *         .... 0... .... = Congestion Window Reduced (CWR): Not set
         *         .... .0.. .... = ECN-Echo: Not set
         *         .... ..0. .... = Urgent: Not set
         *         .... ...1 .... = Acknowledgment: Set
         *         .... .... 0... = Push: Not set
         *         .... .... .0.. = Reset: Not set
         *         .... .... ..0. = Syn: Not set
         *         .... .... ...0 = Fin: Not set
         *         [TCP Flags: ·······A····]
         *     Window: 2038
         *     [Calculated window size: 2038]
         *     [Window size scaling factor: -1 (unknown)]
         *     Checksum: 0xd3c2 [correct]
         *     [Checksum Status: Good]
         *     [Calculated Checksum: 0xd3c2]
         *     Urgent Pointer: 0
         *     Options: (12 bytes), No-Operation (NOP), No-Operation (NOP), Timestamps
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - Timestamps: TSval 2232684208, TSecr 3427137631
         *             Kind: Time Stamp Option (8)
         *             Length: 10
         *             Timestamp value: 2232684208
         *             Timestamp echo reply: 3427137631
         *     [Timestamps]
         *         [Time since first frame in this TCP stream: 0.000000000 seconds]
         *         [Time since previous frame in this TCP stream: 0.000000000 seconds]
         */

        let bytes = [
            0xc3, 0xa1, /* Source Port */
            0x01, 0xbb, /* Destination Port */
            0x5a, 0xc0, 0xae, 0xd0, /* Sequence Number */
            0xcb, 0xcf, 0x60, 0xc2, /* Acknowledgment Number */
            0x80, 0x10, /* DataOffset and Reserved and Flags */
            0x07, 0xf6, /* Window */
            0xd3, 0xc2, /* Checksum */
            0x00, 0x00, /* Urgent Pointer */
            0x01, /* TCP Option - No-Operation */
            0x01, /* TCP Option - No-Operation */
            0x08, 0x0a, 0x85, 0x14, 0x0e, 0xb0, 0xcc, 0x45, 0xf8,
            0x5f, /* TCP Option - Timestamps */
            0xff, /* Payload */
        ];

        let expect_options: Vec<TcpOption> = vec![
            NOP,
            NOP,
            TIMESTAMPS {
                length: 10,
                ts_value: 2232684208,
                ts_reply: 3427137631,
            },
        ];

        let expectation = TcpHeader {
            source_port: 50081,
            dest_port: 443,
            seq_num: 1522577104,
            ack_num: 3419365570,
            data_offset: 32,
            reserved: 0,
            flag_urg: false,
            flag_ack: true,
            flag_psh: false,
            flag_rst: false,
            flag_syn: false,
            flag_fin: false,
            window: 2038,
            checksum: 0xd3c2,
            urgent_ptr: 0,
            options: Some(expect_options),
        };

        assert_eq!(TcpHeader::decode(&bytes), Ok((LAST_SLICE, expectation)));
    }

    #[test]
    fn tcp_header_decode2() {
        /*
         * Transmission Control Protocol, Src Port: 58816, Dst Port: 80, Seq: 0, Len: 0
         *     Source Port: 58816
         *     Destination Port: 80
         *     [Stream index: 0]
         *     [Conversation completeness: Incomplete, DATA (15)]
         *     [TCP Segment Len: 0]
         *     Sequence Number: 0    (relative sequence number)
         *     Sequence Number (raw): 3851697578
         *     [Next Sequence Number: 1    (relative sequence number)]
         *     Acknowledgment Number: 0
         *     Acknowledgment number (raw): 0
         *     1010 .... = Header Length: 40 bytes (10)
         *     Flags: 0x002 (SYN)
         *         000. .... .... = Reserved: Not set
         *         ...0 .... .... = Nonce: Not set
         *         .... 0... .... = Congestion Window Reduced (CWR): Not set
         *         .... .0.. .... = ECN-Echo: Not set
         *         .... ..0. .... = Urgent: Not set
         *         .... ...0 .... = Acknowledgment: Not set
         *         .... .... 0... = Push: Not set
         *         .... .... .0.. = Reset: Not set
         *         .... .... ..1. = Syn: Set
         *         .... .... ...0 = Fin: Not set
         *         [TCP Flags: ··········S·]
         *     Window: 5840
         *     [Calculated window size: 5840]
         *     Checksum: 0x9de2 [correct]
         *     [Checksum Status: Good]
         *     [Calculated Checksum: 0x9de2]
         *     Urgent Pointer: 0
         *     Options: (20 bytes), Maximum segment size, SACK permitted, Timestamps, No-Operation (NOP), Window scale
         *         TCP Option - Maximum segment size: 1460 bytes
         *             Kind: Maximum Segment Size (2)
         *             Length: 4
         *             MSS Value: 1460
         *         TCP Option - SACK permitted
         *             Kind: SACK Permitted (4)
         *             Length: 2
         *         TCP Option - Timestamps: TSval 1545573, TSecr 0
         *             Kind: Time Stamp Option (8)
         *             Length: 10
         *             Timestamp value: 1545573
         *             Timestamp echo reply: 0
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - Window scale: 7 (multiply by 128)
         *             Kind: Window Scale (3)
         *             Length: 3
         *             Shift count: 7
         *             [Multiplier: 128]
         *     [Timestamps]
         *         [Time since first frame in this TCP stream: 0.000000000 seconds]
         *         [Time since previous frame in this TCP stream: 0.000000000 seconds]
         */

        let bytes = [
            0xe5, 0xc0, /* Source Port */
            0x00, 0x50, /* Destination Port */
            0xe5, 0x94, 0x3d, 0xaa, /* Sequence Number */
            0x00, 0x00, 0x00, 0x00, /* Acknowledgment Number */
            0xa0, 0x02, /* DataOffset and Reserved and Flags */
            0x16, 0xd0, /* Window */
            0x9d, 0xe2, /* Checksum */
            0x00, 0x00, /* Urgent Pointer */
            0x02, 0x04, 0x05, 0xb4, /* TCP Option - Maximum segment size */
            0x04, 0x02, /* TCP Option - SACK permitted */
            0x08, 0x0a, 0x00, 0x17, 0x95, 0x65, 0x00, 0x00, 0x00,
            0x00, /* TCP Option - Timestamps */
            0x01, /* TCP Option - No-Operation */
            0x03, 0x03, 0x07, /* TCP Option - Window scale */
            0xff, /* Payload */
        ];

        let expect_options: Vec<TcpOption> = vec![
            MSS {
                length: 4,
                mss: 1460,
            },
            SACKPERMITTED,
            TIMESTAMPS {
                length: 10,
                ts_value: 1545573,
                ts_reply: 0,
            },
            NOP,
            WSCALE {
                length: 3,
                shift_count: 7,
            },
        ];

        let expectation = TcpHeader {
            source_port: 58816,
            dest_port: 80,
            seq_num: 3851697578,
            ack_num: 0,
            data_offset: 40,
            reserved: 0,
            flag_urg: false,
            flag_ack: false,
            flag_psh: false,
            flag_rst: false,
            flag_syn: true,
            flag_fin: false,
            window: 5840,
            checksum: 0x9de2,
            urgent_ptr: 0,
            options: Some(expect_options),
        };

        assert_eq!(TcpHeader::decode(&bytes), Ok((LAST_SLICE, expectation)));
    }

    #[test]
    fn tcp_header_decode3() {
        /*
         * Transmission Control Protocol, Src Port: 58816, Dst Port: 80, Seq: 461, Ack: 17377, Len: 0
         *     Source Port: 58816
         *     Destination Port: 80
         *     [Stream index: 0]
         *     [Conversation completeness: Incomplete, DATA (15)]
         *     [TCP Segment Len: 0]
         *     Sequence Number: 461    (relative sequence number)
         *     Sequence Number (raw): 3851698039
         *     [Next Sequence Number: 461    (relative sequence number)]
         *     Acknowledgment Number: 17377    (relative ack number)
         *     Acknowledgment number (raw): 2747581568
         *     1011 .... = Header Length: 44 bytes (11)
         *     Flags: 0x010 (ACK)
         *         000. .... .... = Reserved: Not set
         *         ...0 .... .... = Nonce: Not set
         *         .... 0... .... = Congestion Window Reduced (CWR): Not set
         *         .... .0.. .... = ECN-Echo: Not set
         *         .... ..0. .... = Urgent: Not set
         *         .... ...1 .... = Acknowledgment: Set
         *         .... .... 0... = Push: Not set
         *         .... .... .0.. = Reset: Not set
         *         .... .... ..0. = Syn: Not set
         *         .... .... ...0 = Fin: Not set
         *         [TCP Flags: ·······A····]
         *     Window: 318
         *     [Calculated window size: 40704]
         *     [Window size scaling factor: 128]
         *     Checksum: 0x34b6 [correct]
         *     [Checksum Status: Good]
         *     [Calculated Checksum: 0x34b6]
         *     Urgent Pointer: 0
         *     Options: (24 bytes), No-Operation (NOP), No-Operation (NOP), Timestamps, No-Operation (NOP), No-Operation (NOP), SACK
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - Timestamps: TSval 1545583, TSecr 2375917095
         *             Kind: Time Stamp Option (8)
         *             Length: 10
         *             Timestamp value: 1545583
         *             Timestamp echo reply: 2375917095
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - No-Operation (NOP)
         *             Kind: No-Operation (1)
         *         TCP Option - SACK 18825-20273
         *             Kind: SACK (5)
         *             Length: 10
         *             left edge = 18825 (relative)
         *             right edge = 20273 (relative)
         *             [TCP SACK Count: 1]
         *     [Timestamps]
         *         [Time since first frame in this TCP stream: 0.100262000 seconds]
         *         [Time since previous frame in this TCP stream: 0.000007000 seconds]
         */

        let bytes = [
            0xe5, 0xc0, /* Source Port */
            0x00, 0x50, /* Destination Port */
            0xe5, 0x94, 0x3f, 0x77, /* Sequence Number */
            0xa3, 0xc4, 0xc4, 0x80, /* Acknowledgment Number */
            0xb0, 0x10, /* DataOffset and Reserved and Flags */
            0x01, 0x3e, /* Window */
            0x34, 0xb6, /* Checksum */
            0x00, 0x00, /* Urgent Pointer */
            0x01, /* TCP Option - No-Operation */
            0x01, /* TCP Option - No-Operation */
            0x08, 0x0a, 0x00, 0x17, 0x95, 0x6f, 0x8d, 0x9d, 0x9e,
            0x27, /* TCP Option - Timestamps */
            0x01, /* TCP Option - No-Operation */
            0x01, /* TCP Option - No-Operation */
            0x05, 0x0a, 0xa3, 0xc4, 0xca, 0x28, 0xa3, 0xc4, 0xcf,
            0xd0, /* TCP Option - SACK */
            0xff, /* Payload */
        ];

        let opts: Vec<TcpOption> = vec![
            NOP,
            NOP,
            TIMESTAMPS {
                length: 10,
                ts_value: 1545583,
                ts_reply: 2375917095,
            },
            NOP,
            NOP,
            SACK {
                length: 10,
                left_edge: 2747583016,
                right_edge: 2747584464,
            },
        ];

        let expectation = TcpHeader {
            source_port: 58816,
            dest_port: 80,
            seq_num: 3851698039,
            ack_num: 2747581568,
            data_offset: 44,
            reserved: 0,
            flag_urg: false,
            flag_ack: true,
            flag_psh: false,
            flag_rst: false,
            flag_syn: false,
            flag_fin: false,
            window: 318,
            checksum: 0x34b6,
            urgent_ptr: 0,
            options: Some(opts),
        };

        assert_eq!(TcpHeader::decode(&bytes), Ok((LAST_SLICE, expectation)));
    }
}