package srt

import (
	"errors"
	"encoding/binary"
)

// arbitrary indexing
const (
	DATA uint8 = iota
	HANDSHAKE
	ACK
	NAK
	ACKACK
	SHUTDOWN
	DROP
)

// see SRT protocol RFC for information
type ControlHeader struct {
	ctrl_type uint16
	ctrl_subtype uint16
	tsi uint32
}

type DataHeader struct {
	seq_num uint32
	msg_flags uint8
	msg_num uint32
}

type HandshakeCIF struct {
	version uint32
	enc_field uint16
	ext_field uint16
	init_seq_num uint32
	mtu uint32
	max_flow uint32
	hs_type uint32
	sock_id uint32
	syn_cookie uint32
	peer_ip [16]byte
	hs_extensions []*HandshakeExtension
}

type HandshakeExtension struct {
	ext_type uint16
	ext_len uint32
	ext_contents interface{}
}

type HSEMSG struct {
	version uint32
	flags uint32
	recv_delay uint16
	send_delay uint16
}

type KMMSG struct {
	key_type uint8
	key_len uint8
	salt [16]byte
	wrapped_key []byte
}

type ACKCIF struct {
	last_acked uint32
	rtt uint32
	var_rtt uint32
	buff_size uint32
	pkt_recv_rate uint32
	bw uint32
	rcv_rate uint32
}

type NACKCIF struct {
	lost_pkts []*pckts_range
}

type pckts_range struct {
	start uint32
	end uint32
}

// should be safe to ignore
type DROPCIF struct {
	to_drop pckts_range
}

// header and cif are interfaces to allow easier typing, fitting above structs
type Packet struct {
	packet_type uint8
	timestamp uint32
	dest_sock uint32
	header_info interface{} 
	cif interface{}
}

// completely pointless since only implementing receiver, here anyway
func marshall_data_packet(packet *Packet, header []byte) ([]byte, error) {
	info, ok_head := packet.header_info.(*DataHeader)
	data, ok_data := packet.cif.([]byte)
	if !ok_head || !ok_data {
		return header, errors.New("data packet does not have data header or data")
		}
	binary.BigEndian.PutUint32(header[:4], info.seq_num)
	head2 := (uint32(info.msg_flags) << 26) + info.msg_num
	binary.BigEndian.PutUint32(header[4:8], head2)
	return append(header, data...), nil
}

func marshall_ctrl_packet(packet *Packet, header []byte) ([]byte, error) {
	ctrl_header, ok_head := packet.header_info.(*ControlHeader)
	if !ok_head {
		return header, errors.New("control packet does not have ctrl header")
	}
	binary.BigEndian.PutUint16(header[:2], ctrl_header.ctrl_type)
	binary.BigEndian.PutUint16(header[2:4], ctrl_header.ctrl_subtype)
	binary.BigEndian.PutUint32(header[4:8], ctrl_header.tsi)
	header[0] |= 0x80
	switch packet.packet_type {
	case HANDSHAKE:
		data, ok_data := packet.cif.(*HandshakeCIF)
		if !ok_data {
			return header, errors.New("Handshake has no data")
		}
		cif := marshall_hs_cif(data)
		return append(header, cif...), nil
	case SHUTDOWN:
		cif := make([]byte, 4)
		return append(header, cif...), nil
	case ACK:
		data, ok_data := packet.cif.(*ACKCIF)
		if !ok_data {
			return header, errors.New("ACK has no data")
		}
		cif := marshall_ack_cif(data)
		return append(header, cif...), nil
	case NAK:
		data, ok_data := packet.cif.(*NACKCIF)
		if !ok_data {
			return header, errors.New("NAK has no data")
		}
		cif := marshall_nack_cif(data)
		return append(header, cif...), nil
	}
	return header, errors.New("Control packet type not recognized")
}

func marshall_nack_cif(data *NACKCIF) ([]byte) {
	var loss_list_bytes []byte
	for _, pkts := range data.lost_pkts {
		if pkts.start == pkts.end {
			curr_bytes := make([]byte, 4)
			binary.BigEndian.PutUint32(curr_bytes, pkts.start)
			loss_list_bytes = append(loss_list_bytes, curr_bytes...)
		} else {
			curr_bytes := make([]byte, 8)
			binary.BigEndian.PutUint32(curr_bytes[:4], pkts.start)
			binary.BigEndian.PutUint32(curr_bytes[4:8], pkts.end)
			if (pkts.end - pkts.start) % uint32(1 << 31) != 1 {
				curr_bytes[0] |= 0x80
			}
			loss_list_bytes = append(loss_list_bytes, curr_bytes...)
		}
	}
	return loss_list_bytes
}

// locations and length are determined by protocol,
// no real point abstracting that
// could be cleaner with reflects, but added work for very little real gain
func marshall_ack_cif(data *ACKCIF) ([]byte) {
	cif := make([]byte, 28)
	binary.BigEndian.PutUint32(cif[:4], data.last_acked)
	binary.BigEndian.PutUint32(cif[4:8], data.rtt)
	binary.BigEndian.PutUint32(cif[8:12], data.var_rtt)
	binary.BigEndian.PutUint32(cif[12:16], data.buff_size)
	binary.BigEndian.PutUint32(cif[16:20], data.pkt_recv_rate)
	binary.BigEndian.PutUint32(cif[20:24], data.bw)
	binary.BigEndian.PutUint32(cif[24:28], data.rcv_rate)

	return cif
}

// same as above
func marshall_hs_cif(data *HandshakeCIF) ([]byte) {
	cif := make([]byte, 48)
	binary.BigEndian.PutUint32(cif[:4], data.version)
	binary.BigEndian.PutUint16(cif[4:6], data.enc_field)
	binary.BigEndian.PutUint16(cif[6:8], data.ext_field)
	binary.BigEndian.PutUint32(cif[8:12], data.init_seq_num)
	binary.BigEndian.PutUint32(cif[12:16], data.mtu)
	binary.BigEndian.PutUint32(cif[16:20], data.max_flow)
	binary.BigEndian.PutUint32(cif[20:24], data.hs_type)
	binary.BigEndian.PutUint32(cif[24:28], data.sock_id)
	binary.BigEndian.PutUint32(cif[28:32], data.syn_cookie)
	for i := 0; i < 16; i++ {
		cif[32 + i] = data.peer_ip[i]
	}
	for _, extension := range data.hs_extensions {
		ext_buff := make([]byte, int(extension.ext_len) + 4)
		binary.BigEndian.PutUint16(ext_buff[:2], extension.ext_type)
		binary.BigEndian.PutUint16(ext_buff[2:4], uint16(extension.ext_len / 4))
		switch extension.ext_type {
		case 2:
			contents := extension.ext_contents.(*HSEMSG)
			binary.BigEndian.PutUint32(ext_buff[4:8], contents.version)
			binary.BigEndian.PutUint32(ext_buff[8:12], contents.flags)
			binary.BigEndian.PutUint16(ext_buff[12:14], contents.recv_delay)
			binary.BigEndian.PutUint16(ext_buff[14:16], contents.send_delay)
		case 4:
			contents, ok := extension.ext_contents.(*KMMSG)
			if !ok { // handle km_state error
				copy(ext_buff[4:8], extension.ext_contents.([]byte))
			} else {
				binary.BigEndian.PutUint32(ext_buff[4:8], uint32(0x12202900) | uint32(contents.key_type))
				binary.BigEndian.PutUint32(ext_buff[12:16], uint32(0x02000200))
				binary.BigEndian.PutUint32(ext_buff[16:20], uint32(0x0400) | uint32(contents.key_len / 4))
				for i := 0; i < 16; i++ {
					ext_buff[20 + i] = contents.salt[i]
				}
				copy(ext_buff[36:], contents.wrapped_key)
			}
		default:
			copy(ext_buff[4:], extension.ext_contents.([]byte))
		}
		cif = append(cif, ext_buff...)
	}
	return cif
}

func MarshallPacket(packet *Packet, agent *SRTManager) ([]byte, error) {
	header := make([]byte, 16)
	binary.BigEndian.PutUint32(header[8:12], packet.timestamp)
	binary.BigEndian.PutUint32(header[12:16], packet.dest_sock)
	if packet.packet_type == DATA {
		return marshall_data_packet(packet, header)
	} else {
		return marshall_ctrl_packet(packet, header)
	}
}

func parse_data_packet(pkt *Packet, buffer []byte) (error) {
	info := new(DataHeader)
	info.seq_num = binary.BigEndian.Uint32(buffer[:4])
	info.msg_flags = buffer[4] >> 2 // unsused since live streaming makes it irrelevant, kk for encrypt eventually
	info.msg_num = binary.BigEndian.Uint32(buffer[4:8]) & 0x03ffffff

	pkt.header_info = info
	if len(buffer) > 16 {
		data := make([]byte, len(buffer) - 16)
		copy(data, buffer[16:])
		pkt.cif = data
		return nil
	}
	return errors.New("Data Packet has no data")
}

func parse_ctrl_packet(pkt *Packet, buffer []byte) (error) {
	info := new(ControlHeader)
	info.ctrl_type = binary.BigEndian.Uint16(buffer[:2]) & 0x7fff
	info.ctrl_subtype = binary.BigEndian.Uint16(buffer[2:4])
	info.tsi = binary.BigEndian.Uint32(buffer[4:8])

	pkt.header_info = info
	
	switch info.ctrl_type {
	case 0:
		pkt.packet_type = HANDSHAKE
		if len(buffer) >= 64 {
			cif := new(HandshakeCIF)
			pkt.cif = cif
			return parse_hs_cif(cif, buffer[16:])
		}
		return errors.New("HS not long enough")
	case 6:
		pkt.packet_type = ACKACK
		return nil
	case 5:
		pkt.packet_type = SHUTDOWN
		return nil
	default:
		return errors.New("Unexpected control type")
	}
}

func parse_hs_cif(cif *HandshakeCIF, buffer []byte) (error) {
	cif.version = binary.BigEndian.Uint32(buffer[:4])
	cif.enc_field = binary.BigEndian.Uint16(buffer[4:6])
	cif.ext_field = binary.BigEndian.Uint16(buffer[6:8])
	cif.init_seq_num = binary.BigEndian.Uint32(buffer[8:12])
	cif.mtu = binary.BigEndian.Uint32(buffer[12:16])
	cif.max_flow = binary.BigEndian.Uint32(buffer[16:20])
	cif.hs_type = binary.BigEndian.Uint32(buffer[20:24])
	cif.sock_id = binary.BigEndian.Uint32(buffer[24:28])
	cif.syn_cookie = binary.BigEndian.Uint32(buffer[28:32])
	for i := 0; i < 16; i++ {
		cif.peer_ip[i] = buffer[32 + i]
	}
	extensions := buffer[48:]
	for len(extensions) != 0 {
		if len(extensions) <= 4 {
			return errors.New("Extension present, shorter than header")
		}
		ext := new(HandshakeExtension)
		ext.ext_type = binary.BigEndian.Uint16(extensions[:2])
		ext.ext_len = uint32(binary.BigEndian.Uint16(extensions[2:4])) * 4
		if len(extensions) < 4 + int(ext.ext_len) {
			return errors.New("Extension shorter than advertised")
		}
		switch ext.ext_type {
		case 1:
			content := new(HSEMSG)
			content.version = binary.BigEndian.Uint32(extensions[4:8])
			content.flags = binary.BigEndian.Uint32(extensions[8:12])
			content.recv_delay = binary.BigEndian.Uint16(extensions[12:14])
			content.send_delay = binary.BigEndian.Uint16(extensions[14:16])
			ext.ext_contents = content
		case 3:
			content := new(KMMSG)
			content.key_type = extensions[7] & 0x3
			content.key_len = extensions[19] * 4
			for i := 0; i < 16; i++ {
				content.salt[i] = extensions[20 + i]
			}
			// -36 from actual content len, extensions includes headers as well
			wrap_key_len := 4 + ext.ext_len - 36
			content.wrapped_key = make([]byte, wrap_key_len)
			copy(content.wrapped_key, extensions[36:36 + wrap_key_len])
			ext.ext_contents = content
		default:
			content := make([]byte, ext.ext_len)
			copy(content, extensions[4:4 + ext.ext_len])
			ext.ext_contents = content
		}
		cif.hs_extensions = append(cif.hs_extensions, ext)
		extensions = extensions[4 + ext.ext_len:]
	}
	return nil
}

func ParsePacket(buffer []byte) (*Packet, error) {
	if len(buffer) < 16 {
		return nil, errors.New("packet too short")
	}
	pkt := new(Packet)
	pkt.timestamp = binary.BigEndian.Uint32(buffer[8:12])
	pkt.dest_sock = binary.BigEndian.Uint32(buffer[12:16])
	
	if buffer[0] >> 7 == 0 {
		err := parse_data_packet(pkt, buffer)
		if err != nil {
			return pkt, err
		}
	} else {
		err := parse_ctrl_packet(pkt, buffer)
		if err != nil {
			return pkt, err
		}
	}
	return pkt, nil
}