package codec

import "encoding/binary"

// nal_unit( NumBytesInNALunit ) {
//     forbidden_zero_bit All         f(1)
//     nal_ref_idc All             u(2)
//     nal_unit_type                u(5)
// }

type H264NaluHdr struct {
	Forbidden_zero_bit uint8
	Nal_ref_idc        uint8
	Nal_unit_type      uint8
}

func (hdr *H264NaluHdr) Decode(bs *BitStream) {
	hdr.Forbidden_zero_bit = bs.GetBit()
	hdr.Nal_ref_idc = bs.Uint8(2)
	hdr.Nal_unit_type = bs.Uint8(5)
}

type SliceHeader struct {
	First_mb_in_slice    uint64
	Slice_type           uint64
	Pic_parameter_set_id uint64
	Frame_num            uint64
}

//调用方根据sps中的log2_max_frame_num_minus4的值来解析Frame_num
func (sh *SliceHeader) Decode(bs *BitStream) {
	sh.First_mb_in_slice = bs.ReadUE()
	sh.Slice_type = bs.ReadUE()
	sh.Pic_parameter_set_id = bs.ReadUE()
}

type SPS struct {
	Profile_idc                          uint8
	Constraint_set0_flag                 uint8
	Constraint_set1_flag                 uint8
	Constraint_set2_flag                 uint8
	Constraint_set3_flag                 uint8
	Constraint_set4_flag                 uint8
	Constraint_set5_flag                 uint8
	Reserved_zero_2bits                  uint8
	Level_idc                            uint8
	Seq_parameter_set_id                 uint64
	Chroma_format_idc                    uint64
	Separate_colour_plane_flag           uint8
	Bit_depth_luma_minus8                uint64
	Bit_depth_chroma_minus8              uint64
	Log2_max_frame_num_minus4            uint64
	Pic_order_cnt_type                   uint64
	Max_num_ref_frames                   uint64
	Gaps_in_frame_num_value_allowed_flag uint8
	Pic_width_in_mbs_minus1              uint64
	Pic_height_in_map_units_minus1       uint64
	Frame_mbs_only_flag                  uint8
	Direct_8x8_inference_flag            uint8
	Frame_cropping_flag                  uint8
	Frame_crop_left_offset               uint64
	Frame_crop_right_offset              uint64
	Frame_crop_top_offset                uint64
	Frame_crop_bottom_offset             uint64
	Vui_parameters_present_flag          uint8
}

func (sps *SPS) Decode(bs *BitStream) {
	sps.Profile_idc = bs.Uint8(8)
	sps.Constraint_set0_flag = bs.GetBit()
	sps.Constraint_set1_flag = bs.GetBit()
	sps.Constraint_set2_flag = bs.GetBit()
	sps.Constraint_set3_flag = bs.GetBit()
	sps.Constraint_set4_flag = bs.GetBit()
	sps.Constraint_set5_flag = bs.GetBit()
	sps.Reserved_zero_2bits = bs.Uint8(2)
	sps.Level_idc = bs.Uint8(8)
	sps.Seq_parameter_set_id = bs.ReadUE()
	if sps.Profile_idc == 100 || sps.Profile_idc == 110 ||
		sps.Profile_idc == 122 || sps.Profile_idc == 244 ||
		sps.Profile_idc == 44 || sps.Profile_idc == 83 ||
		sps.Profile_idc == 86 || sps.Profile_idc == 118 || sps.Profile_idc == 128 {
		sps.Chroma_format_idc = bs.ReadUE()
		if sps.Chroma_format_idc == 3 {
			sps.Separate_colour_plane_flag = bs.Uint8(1) //separate_colour_plane_flag
		}
		sps.Bit_depth_luma_minus8 = bs.ReadUE()   //bit_depth_luma_minus8
		sps.Bit_depth_chroma_minus8 = bs.ReadUE() //bit_depth_chroma_minus8
		bs.SkipBits(1)                            //qpprime_y_zero_transform_bypass_flag
		seq_scaling_matrix_present_flag := bs.GetBit()
		if seq_scaling_matrix_present_flag == 1 {
			//seq_scaling_list_present_flag[i]
			if sps.Chroma_format_idc == 3 {
				bs.SkipBits(12)
			} else {
				bs.SkipBits(8)
			}
		}
	}
	sps.Log2_max_frame_num_minus4 = bs.ReadUE()
	sps.Pic_order_cnt_type = bs.ReadUE()
	if sps.Pic_order_cnt_type == 0 {
		bs.ReadUE() // log2_max_pic_order_cnt_lsb_minus4
	} else if sps.Pic_order_cnt_type == 1 {
		bs.SkipBits(1) //delta_pic_order_always_zero_flag
		bs.ReadSE()    //offset_for_non_ref_pic
		bs.ReadSE()    //offset_for_top_to_bottom_field
		num_ref_frames_in_pic_order_cnt_cycle := bs.ReadUE()
		for i := 0; i < int(num_ref_frames_in_pic_order_cnt_cycle); i++ {
			bs.ReadSE() //offset_for_ref_frame
		}
	}
	sps.Max_num_ref_frames = bs.ReadUE()
	sps.Gaps_in_frame_num_value_allowed_flag = bs.GetBit()
	sps.Pic_width_in_mbs_minus1 = bs.ReadUE()
	sps.Pic_height_in_map_units_minus1 = bs.ReadUE()
	sps.Frame_mbs_only_flag = bs.GetBit()
	if sps.Frame_mbs_only_flag == 0 {
		bs.SkipBits(1) // mb_adaptive_frame_field_flag
	}
	sps.Direct_8x8_inference_flag = bs.GetBit()
	sps.Frame_cropping_flag = bs.GetBit()
	if sps.Frame_cropping_flag == 1 {
		sps.Frame_crop_left_offset = bs.ReadUE()   //frame_crop_left_offset
		sps.Frame_crop_right_offset = bs.ReadUE()  //frame_crop_right_offset
		sps.Frame_crop_top_offset = bs.ReadUE()    //frame_crop_top_offset
		sps.Frame_crop_bottom_offset = bs.ReadUE() //frame_crop_bottom_offset
	}
	sps.Vui_parameters_present_flag = bs.GetBit()
}

type PPS struct {
	Pic_parameter_set_id                         uint64
	Seq_parameter_set_id                         uint64
	Entropy_coding_mode_flag                     uint8
	Bottom_field_pic_order_in_frame_present_flag uint8
	Num_slice_groups_minus1                      uint64
}

func (pps *PPS) Decode(bs *BitStream) {
	pps.Pic_parameter_set_id = bs.ReadUE()
	pps.Seq_parameter_set_id = bs.ReadUE()
	pps.Entropy_coding_mode_flag = bs.GetBit()
	pps.Bottom_field_pic_order_in_frame_present_flag = bs.GetBit()
	pps.Num_slice_groups_minus1 = bs.ReadUE()
}

type SEIReaderWriter interface {
	Read(size uint16, bs *BitStream)
	Write(bsw *BitStreamWriter)
}

type UserDataUnregistered struct {
	UUID     []byte
	UserData []byte
}

func (udu *UserDataUnregistered) Read(size uint16, bs *BitStream) {
	udu.UUID = bs.GetBytes(16)
	udu.UserData = bs.GetBytes(int(size - 16))
}

func (udu *UserDataUnregistered) Write(bsw *BitStreamWriter) {
	bsw.PutBytes(udu.UUID)
	bsw.PutBytes(udu.UserData)
}

type SEI struct {
	PayloadType uint16
	PayloadSize uint16
	Sei_payload SEIReaderWriter
}

func (sei *SEI) Decode(bs *BitStream) {
	for bs.NextBits(8) == 0xFF {
		sei.PayloadType += 255
	}
	sei.PayloadType += uint16(bs.Uint8(8))
	for bs.NextBits(8) == 0xFF {
		sei.PayloadSize += 255
	}
	sei.PayloadSize += uint16(bs.Uint8(8))
	if sei.PayloadType == 5 {
		sei.Sei_payload = new(UserDataUnregistered)
		sei.Sei_payload.Read(sei.PayloadSize, bs)
	}
}

func (sei *SEI) Encode(bsw *BitStreamWriter) []byte {
	payloadType := sei.PayloadType
	payloadSize := sei.PayloadSize
	for payloadType >= 0xFF {
		bsw.PutByte(0xFF)
		payloadType -= 255
	}
	bsw.PutByte(uint8(payloadType))
	for payloadSize >= 0xFF {
		bsw.PutByte(0xFF)
		payloadSize -= 255
	}
	bsw.PutByte(uint8(payloadSize))
	sei.Sei_payload.Write(bsw)
	return bsw.Bits()
}

func GetSPSIdWithStartCode(sps []byte) uint64 {
	start, sc := FindStartCode(sps, 0)
	return GetSPSId(sps[start+int(sc):])
}

func GetSPSId(sps []byte) uint64 {
	sps = sps[1:]
	bs := NewBitStream(sps)
	bs.SkipBits(24)
	return bs.ReadUE()
}

func GetPPSIdWithStartCode(pps []byte) uint64 {
	start, sc := FindStartCode(pps, 0)
	return GetPPSId(pps[start+int(sc):])
}

func GetPPSId(pps []byte) uint64 {
	pps = pps[1:]
	bs := NewBitStream(pps)
	return bs.ReadUE()
}

//https://stackoverflow.com/questions/12018535/get-the-width-height-of-the-video-from-h-264-nalu
//int Width = ((pic_width_in_mbs_minus1 +1)*16) - frame_crop_right_offset *2 - frame_crop_left_offset *2;
//int Height = ((2 - frame_mbs_only_flag)* (pic_height_in_map_units_minus1 +1) * 16) - (frame_crop_bottom_offset* 2) - (frame_crop_top_offset* 2);
func GetH264Resolution(sps []byte) (width uint32, height uint32) {
	start, sc := FindStartCode(sps, 0)
	bs := NewBitStream(sps[start+int(sc)+1:])
	var s SPS
	s.Decode(bs)

	widthInSample := (uint32(s.Pic_width_in_mbs_minus1) + 1) * 16
	widthCrop := uint32(s.Frame_crop_left_offset)*2 - uint32(s.Frame_crop_right_offset)*2
	width = widthInSample - widthCrop

	heightInSample := ((2 - uint32(s.Frame_mbs_only_flag)) * (uint32(s.Pic_height_in_map_units_minus1) + 1) * 16)
	heightCrop := uint32(s.Frame_crop_bottom_offset)*2 - uint32(s.Frame_crop_top_offset)*2
	height = heightInSample - heightCrop

	return
}

// aligned(8) class AVCDecoderConfigurationRecord {
// 	unsigned int(8) configurationVersion = 1;
// 	unsigned int(8) AVCProfileIndication;
// 	unsigned int(8) profile_compatibility;
// 	unsigned int(8) AVCLevelIndication;
// 	bit(6) reserved = ‘111111’b;
// 	unsigned int(2) lengthSizeMinusOne;
// 	bit(3) reserved = ‘111’b;
// 	unsigned int(5) numOfSequenceParameterSets;
// 	for (i=0; i< numOfSequenceParameterSets;  i++) {
// 		unsigned int(16) sequenceParameterSetLength ;
// 		bit(8*sequenceParameterSetLength) sequenceParameterSetNALUnit;
// 	}
// 	unsigned int(8) numOfPictureParameterSets;
// 	for (i=0; i< numOfPictureParameterSets;  i++) {
// 		unsigned int(16) pictureParameterSetLength;
// 		bit(8*pictureParameterSetLength) pictureParameterSetNALUnit;
// 	}
// if( profile_idc  ==  100  ||  profile_idc  ==  110  ||
//     profile_idc  ==  122  ||  profile_idc  ==  144 )
// {
//     bit(6) reserved = ‘111111’b;
//     unsigned int(2) chroma_format;
//     bit(5) reserved = ‘11111’b;
//     unsigned int(3) bit_depth_luma_minus8;
//     bit(5) reserved = ‘11111’b;
//     unsigned int(3) bit_depth_chroma_minus8;
//     unsigned int(8) numOfSequenceParameterSetExt;
//     for (i=0; i< numOfSequenceParameterSetExt; i++) {
//      unsigned int(16) sequenceParameterSetExtLength;
//      bit(8*sequenceParameterSetExtLength) sequenceParameterSetExtNALUnit;
//     }
//  }

//  }
// }
// bits
// 8   version ( always 0x01 )
// 8   avc profile ( sps[0][1] )
// 8   avc compatibility ( sps[0][2] )
// 8   avc level ( sps[0][3] )
// 6   reserved ( all bits on )
// 2   NALULengthSizeMinusOne
// 3   reserved ( all bits on )
// 5   number of SPS NALUs (usually 1)

// repeated once per SPS:
// 16         SPS size
// variable   SPS NALU data

// 8   number of PPS NALUs (usually 1)
// repeated once per PPS:
//   16       PPS size
//   variable PPS NALU data

func CreateH264AVCCExtradata(spss [][]byte, ppss [][]byte) []byte {
	extradata := make([]byte, 6, 256)
	for i, sps := range spss {
		start, sc := FindStartCode(sps, 0)
		spss[i] = sps[start+int(sc):]
	}

	for i, pps := range ppss {
		start, sc := FindStartCode(pps, 0)
		ppss[i] = pps[start+int(sc):]
	}

	extradata[0] = 0x01
	extradata[1] = spss[0][1]
	extradata[2] = spss[0][2]
	extradata[3] = spss[0][3]
	extradata[4] = 0xFF
	extradata[5] = 0xE0 | uint8(len(spss))
	for _, sps := range spss {
		spssize := make([]byte, 2)
		binary.BigEndian.PutUint16(spssize, uint16(len(sps)))
		extradata = append(extradata, spssize...)
		extradata = append(extradata, sps...)
	}
	extradata = append(extradata, uint8(len(ppss)))
	for _, pps := range ppss {
		ppssize := make([]byte, 2)
		binary.BigEndian.PutUint16(ppssize, uint16(len(pps)))
		extradata = append(extradata, ppssize...)
		extradata = append(extradata, pps...)
	}
	var h264sps SPS
	h264sps.Decode(NewBitStream(spss[0][1:]))
	if h264sps.Profile_idc == 100 ||
		h264sps.Profile_idc == 110 ||
		h264sps.Profile_idc == 122 ||
		h264sps.Profile_idc == 144 {
		tmp := make([]byte, 4)
		tmp[0] = 0xFC | uint8(h264sps.Chroma_format_idc&0x03)
		tmp[1] = 0xF8 | uint8(h264sps.Bit_depth_luma_minus8&0x07)
		tmp[2] = 0xF8 | uint8(h264sps.Bit_depth_chroma_minus8&0x07)
		tmp[3] = 0
		extradata = append(extradata, tmp...)
	}

	return extradata
}

func CovertExtradata(extraData []byte) ([][]byte, [][]byte) {
	spsnum := extraData[5] & 0x1F
	spss := make([][]byte, spsnum)
	offset := 6
	for i := 0; i < int(spsnum); i++ {
		spssize := binary.BigEndian.Uint16(extraData[offset:])
		sps := make([]byte, spssize+4)
		copy(sps, []byte{0x00, 0x00, 0x00, 0x01})
		copy(sps[4:], extraData[offset+2:offset+2+int(spssize)])
		offset += 2 + int(spssize)
		spss[i] = sps
	}
	ppsnum := extraData[offset]
	ppss := make([][]byte, ppsnum)
	offset++
	for i := 0; i < int(ppsnum); i++ {
		ppssize := binary.BigEndian.Uint16(extraData[offset:])
		pps := make([]byte, ppssize+4)
		copy(pps, []byte{0x00, 0x00, 0x00, 0x01})
		copy(pps[4:], extraData[offset+2:offset+2+int(ppssize)])
		offset += 2 + int(ppssize)
		ppss[i] = pps
	}
	return spss, ppss
}

func ConvertAnnexBToAVCC(annexb []byte) []byte {
	start, sc := FindStartCode(annexb, 0)
	if sc == START_CODE_4 {
		binary.BigEndian.PutUint32(annexb[start:], uint32(len(annexb)-4))
		return annexb
	} else {
		avcc := make([]byte, 1+len(annexb))
		binary.BigEndian.PutUint32(avcc, uint32(len(annexb)-3))
		copy(avcc[4:], annexb[start+3:])
		return avcc
	}
}

func CovertAVCCToAnnexB(avcc []byte) {
	avcc[0] = 0x00
	avcc[1] = 0x00
	avcc[2] = 0x00
	avcc[3] = 0x01
}