package mpeg

import (
	"github.com/wader/fq/format"
	"github.com/wader/fq/format/registry"
	"github.com/wader/fq/pkg/decode"
)

func init() {
	registry.MustRegister(decode.Format{
		Name:        format.AVC_SPS,
		Description: "H.264/AVC Sequence Parameter Set",
		DecodeFn:    avcSPSDecode,
	})
}

func avcVuiParameters(d *decode.D) {
	aspectRatioInfoPresentFlag := d.FieldBool("aspect_ratio_info_present_flag")
	if aspectRatioInfoPresentFlag {
		aspectRatioIdc := d.FieldU8("aspect_ratio_idc")
		const extendedSAR = 255
		if aspectRatioIdc == extendedSAR {
			d.FieldU16("width")
			d.FieldU16("height")
		}
	}
	overscanInfoPresentFlag := d.FieldBool("overscan_info_present_flag")
	if overscanInfoPresentFlag {
		d.FieldBool("overscan_appropriate_flag")
	}
	videoSignalTypePresentFlag := d.FieldBool("video_signal_type_present_flag")
	if videoSignalTypePresentFlag {
		d.FieldU3("video_format")
		d.FieldBool("video_full_range_flag")
		colourDescriptionPresentFlag := d.FieldBool("colour_description_present_flag")
		if colourDescriptionPresentFlag {
			d.FieldU8("colour_primaries")
			d.FieldU8("transfer_characteristics")
			d.FieldU8("matrix_coefficients")
		}
	}
	chromaLocInfoPresentFlag := d.FieldBool("chroma_loc_info_present_flag")
	if chromaLocInfoPresentFlag {
		d.FieldUFn("chroma_sample_loc_type_top_field", uEV)
		d.FieldUFn("chroma_sample_loc_type_bottom_field", uEV)
	}

	timingInfoPresentFlag := d.FieldBool("timing_info_present_flag")

	if timingInfoPresentFlag {
		d.FieldU32("num_units_in_tick")
		d.FieldU32("time_scale")
		d.FieldBool("fixed_frame_rate_flag")
	}
	nalHrdParametersPresentFlag := d.FieldBool("nal_hrd_parameters_present_flag")
	if nalHrdParametersPresentFlag {
		d.FieldStruct("nal_hrd_parameters", avcHdrParameters)
	}
	vclHrdParametersPresentFlag := d.FieldBool("vcl_hrd_parameters_present_flag")
	if vclHrdParametersPresentFlag {
		d.FieldStruct("vcl_hrd_parameters", avcHdrParameters)
	}
	if nalHrdParametersPresentFlag || vclHrdParametersPresentFlag {
		d.FieldBool("low_delay_hrd_flag")
	}
	d.FieldBool("pic_struct_present_flag")
	bitstreamRestrictionFlag := d.FieldBool("bitstream_restriction_flag")
	if bitstreamRestrictionFlag {
		d.FieldBool("motion_vectors_over_pic_boundaries_flag")
		d.FieldUFn("max_bytes_per_pic_denom", uEV)
		d.FieldUFn("max_bits_per_mb_denom", uEV)
		d.FieldUFn("log2_max_mv_length_horizontal", uEV)
		d.FieldUFn("log2_max_mv_length_vertical", uEV)
		d.FieldUFn("max_num_reorder_frames", uEV)
		d.FieldUFn("max_dec_frame_buffering", uEV)
	}
}

func avcHdrParameters(d *decode.D) {
	cpbCntMinus1 := d.FieldUFn("cpb_cnt_minus1", uEV)
	_ = cpbCntMinus1
	d.FieldU4("bit_rate_scale")
	d.FieldU4("cpb_size_scale")
	d.FieldArray("sched_sels", func(d *decode.D) {
		for i := uint64(0); i <= cpbCntMinus1; i++ {
			d.FieldStruct("sched_sel", func(d *decode.D) {
				d.FieldUFn("bit_rate_value_minus1", uEV)
				d.FieldUFn("cpb_size_value_minus1", uEV)
				d.FieldBool("cbr_flag")
			})
		}
	})
	d.FieldU5("initial_cpb_removal_delay_length_minus1")
	d.FieldU5("cpb_removal_delay_length_minus1")
	d.FieldU5("dpb_output_delay_length_minus1")
	d.FieldU5("time_offset_length")
}

func avcSPSDecode(d *decode.D, in interface{}) interface{} {
	profileIdc := d.FieldU8("profile_idc", d.MapUToStrSym(avcProfileNames))
	d.FieldBool("constraint_set0_flag")
	d.FieldBool("constraint_set1_flag")
	d.FieldBool("constraint_set2_flag")
	d.FieldBool("constraint_set3_flag")
	d.FieldBool("constraint_set4_flag")
	d.FieldBool("constraint_set5_flag")
	d.FieldU2("reserved_zero_2bits")
	d.FieldU8("level_idc", d.MapUToStrSym(avcLevelNames))
	d.FieldUFn("seq_parameter_set_id", uEV)

	switch profileIdc {
	// TODO: ffmpeg has some more (legacy values?)
	case 100, 110, 122, 244, 44, 83, 86, 118, 128, 138, 139, 134, 135:
		chromaFormatIdc := d.FieldUFn("chroma_format_idc", uEV)
		if chromaFormatIdc == 3 {
			d.FieldBool("separate_colour_plane_flag")
		}

		d.FieldUFn("bit_depth_luma_minus8", uEV)
		d.FieldUFn("bit_depth_chroma_minus8", uEV)
		d.FieldBool("qpprime_y_zero_transform_bypass_flag")
		seqScalingMatrixPresentFlag := d.FieldBool("seq_scaling_matrix_present_flag")
		// TODO:
		_ = seqScalingMatrixPresentFlag
	}

	d.FieldUFn("log2_max_frame_num_minus4", uEV)

	picOrderCntType := d.FieldUFn("pic_order_cnt_type", uEV)
	if picOrderCntType == 0 {
		d.FieldUFn("log2_max_pic_order_cnt_lsb_minus4", uEV)
	} else if picOrderCntType == 1 {
		d.FieldBool("delta_pic_order_always_zero_flag")
		d.FieldSFn("offset_for_non_ref_pic", sEV)
		d.FieldSFn("offset_for_top_to_bottom_field", sEV)
		numRefFramesInPicOrderCntCycle := d.FieldUFn("num_ref_frames_in_pic_order_cnt_cycle", uEV)
		d.FieldArray("offset_for_ref_frames", func(d *decode.D) {
			for i := uint64(0); i < numRefFramesInPicOrderCntCycle; i++ {
				sEV(d)
			}
		})
	}

	d.FieldUFn("max_num_ref_frames", uEV)
	d.FieldBool("gaps_in_frame_num_value_allowed_flag")
	d.FieldUFn("pic_width_in_mbs_minus1", uEV)
	d.FieldUFn("pic_height_in_map_units_minus1", uEV)
	frameMbsOnlyFlag := d.FieldBool("frame_mbs_only_flag")
	if !frameMbsOnlyFlag {
		d.FieldBool("mb_adaptive_frame_field_flag")
	}
	d.FieldBool("direct_8x8_inference_flag")
	frameCroppingFlag := d.FieldBool("frame_cropping_flag")
	if frameCroppingFlag {
		d.FieldUFn("frame_crop_left_offset", uEV)
		d.FieldUFn("frame_crop_right_offset", uEV)
		d.FieldUFn("frame_crop_top_offset", uEV)
		d.FieldUFn("frame_crop_bottom_offset", uEV)
	}
	vuiParametersPresentFlag := d.FieldBool("vui_parameters_present_flag")
	if vuiParametersPresentFlag {
		d.FieldStruct("vui_parameters", avcVuiParameters)
	}

	d.FieldRawLen("rbsp_trailing_bits", d.BitsLeft())

	return nil
}