package h264dec

import (
	"bytes"
	"fmt"

	"bitbucket.org/ausocean/av/codec/h264/h264dec/bits"
	"github.com/pkg/errors"
)

var (
	DefaultScalingMatrix4x4 = [][]int{
		{6, 13, 20, 28, 13, 20, 28, 32, 20, 28, 32, 37, 28, 32, 37, 42},
		{10, 14, 20, 24, 14, 20, 24, 27, 20, 24, 27, 30, 24, 27, 30, 34},
	}

	DefaultScalingMatrix8x8 = [][]int{
		{6, 10, 13, 16, 18, 23, 25, 27,
			10, 11, 16, 18, 23, 25, 27, 29,
			13, 16, 18, 23, 25, 27, 29, 31,
			16, 18, 23, 25, 27, 29, 31, 33,
			18, 23, 25, 27, 29, 31, 33, 36,
			23, 25, 27, 29, 31, 33, 36, 38,
			25, 27, 29, 31, 33, 36, 38, 40,
			27, 29, 31, 33, 36, 38, 40, 42},
		{9, 13, 15, 17, 19, 21, 22, 24,
			13, 13, 17, 19, 21, 22, 24, 25,
			15, 17, 19, 21, 22, 24, 25, 27,
			17, 19, 21, 22, 24, 25, 27, 28,
			19, 21, 22, 24, 25, 27, 28, 30,
			21, 22, 24, 25, 27, 28, 30, 32,
			22, 24, 25, 27, 28, 30, 32, 33,
			24, 25, 27, 28, 30, 32, 33, 35},
	}
	Default4x4IntraList = []int{6, 13, 13, 20, 20, 20, 38, 38, 38, 38, 32, 32, 32, 37, 37, 42}
	Default4x4InterList = []int{10, 14, 14, 20, 20, 20, 24, 24, 24, 24, 27, 27, 27, 30, 30, 34}
	Default8x8IntraList = []int{
		6, 10, 10, 13, 11, 13, 16, 16, 16, 16, 18, 18, 18, 18, 18, 23,
		23, 23, 23, 23, 23, 25, 25, 25, 25, 25, 25, 25, 27, 27, 27, 27,
		27, 27, 27, 27, 29, 29, 29, 29, 29, 29, 29, 31, 31, 31, 31, 31,
		31, 33, 33, 33, 33, 33, 36, 36, 36, 36, 38, 38, 38, 40, 40, 42}
	Default8x8InterList = []int{
		9, 13, 13, 15, 13, 15, 17, 17, 17, 17, 19, 19, 19, 19, 19, 21,
		21, 21, 21, 21, 21, 22, 22, 22, 22, 22, 22, 22, 24, 24, 24, 24,
		24, 24, 24, 24, 25, 25, 25, 25, 25, 25, 25, 27, 27, 27, 27, 27,
		27, 28, 28, 28, 28, 28, 30, 30, 30, 30, 32, 32, 32, 33, 33, 35}
	ScalingList4x4 = map[int][]int{
		0:  Default4x4IntraList,
		1:  Default4x4IntraList,
		2:  Default4x4IntraList,
		3:  Default4x4InterList,
		4:  Default4x4InterList,
		5:  Default4x4InterList,
		6:  Default8x8IntraList,
		7:  Default8x8InterList,
		8:  Default8x8IntraList,
		9:  Default8x8InterList,
		10: Default8x8IntraList,
		11: Default8x8InterList,
	}
	ScalingList8x8 = ScalingList4x4
)

// SPS describes a sequence parameter set as defined by section 7.3.2.1.1 in
// the Specifications.
type SPS struct {
	Profile                         uint8
	Constraint0                     bool
	Constraint1                     bool
	Constraint2                     bool
	Constraint3                     bool
	Constraint4                     bool
	Constraint5                     bool
	LevelIDC                        uint8
	SPSID                           uint64
	ChromaFormatIDC                 uint64
	SeparateColorPlaneFlag          bool
	BitDepthLumaMinus8              uint64
	BitDepthChromaMinus8            uint64
	QPPrimeYZeroTransformBypassFlag bool
	SeqScalingMatrixPresentFlag     bool
	SeqScalingListPresentFlag       []bool
	ScalingList4x4                  [][]uint64
	UseDefaultScalingMatrix4x4Flag  []bool
	ScalingList8x8                  [][]uint64
	UseDefaultScalingMatrix8x8Flag  []bool
	Log2MaxFrameNumMinus4           uint64
	PicOrderCountType               uint64
	Log2MaxPicOrderCntLSBMin4       uint64
	DeltaPicOrderAlwaysZeroFlag     bool
	OffsetForNonRefPic              int64
	OffsetForTopToBottomField       int64
	NumRefFramesInPicOrderCntCycle  uint64
	OffsetForRefFrameList           []int
	MaxNumRefFrames                 uint64
	GapsInFrameNumValueAllowed      bool
	PicWidthInMBSMinus1             uint64
	PicHeightInMapUnitsMinus1       uint64
	FrameMBSOnlyFlag                bool
	MBAdaptiveFrameFieldFlag        bool
	Direct8x8InferenceFlag          bool
	FrameCroppingFlag               bool
	FrameCropLeftOffset             uint64
	FrameCropRightOffset            uint64
	FrameCropTopOffset              uint64
	FrameCropBottomOffset           uint64
	VUIParametersPresentFlag        bool
	VUIParameters                   *VUIParameters
}

// NewSPS parses a sequence parameter set raw byte sequence from br following
// the syntax structure specified in section 7.3.2.1.1, and returns as a new
// SPS.
func NewSPS(rbsp []byte, showPacket bool) (*SPS, error) {
	logger.Printf("debug: SPS RBSP %d bytes %d bits\n", len(rbsp), len(rbsp)*8)
	logger.Printf("debug: \t%#v\n", rbsp[0:8])
	sps := SPS{}
	br := bits.NewBitReader(bytes.NewReader(rbsp))
	r := newFieldReader(br)

	sps.Profile = uint8(r.readBits(8))
	sps.Constraint0 = r.readBits(1) == 1
	sps.Constraint1 = r.readBits(1) == 1
	sps.Constraint2 = r.readBits(1) == 1
	sps.Constraint3 = r.readBits(1) == 1
	sps.Constraint4 = r.readBits(1) == 1
	sps.Constraint5 = r.readBits(1) == 1
	r.readBits(2) // 2 reserved bits.
	sps.LevelIDC = uint8(r.readBits(8))
	sps.SPSID = r.readUe()
	sps.ChromaFormatIDC = r.readUe()

	// This should be done only for certain ProfileIDC:
	isProfileIDC := []int{100, 110, 122, 244, 44, 83, 86, 118, 128, 138, 139, 134, 135}
	// SpecialProfileCase1
	if isInList(isProfileIDC, int(sps.Profile)) {
		if sps.ChromaFormatIDC == chroma444 {
			// TODO: should probably deal with error here.
			sps.SeparateColorPlaneFlag = r.readBits(1) == 1
		}

		sps.BitDepthLumaMinus8 = r.readUe()
		sps.BitDepthChromaMinus8 = r.readUe()
		sps.QPPrimeYZeroTransformBypassFlag = r.readBits(1) == 1
		sps.SeqScalingMatrixPresentFlag = r.readBits(1) == 1

		if sps.SeqScalingMatrixPresentFlag {
			max := 12
			if sps.ChromaFormatIDC != chroma444 {
				max = 8
			}
			logger.Printf("debug: \tbuilding Scaling matrix for %d elements\n", max)
			for i := 0; i < max; i++ {
				sps.SeqScalingListPresentFlag = append(sps.SeqScalingListPresentFlag, r.readBits(1) == 1)

				if sps.SeqScalingListPresentFlag[i] {
					if i < 6 {
						scalingList(
							br,
							ScalingList4x4[i],
							16,
							DefaultScalingMatrix4x4[i])
						// 4x4: Page 75 bottom
					} else {
						// 8x8 Page 76 top
						scalingList(
							br,
							ScalingList8x8[i],
							64,
							DefaultScalingMatrix8x8[i-6])
					}
				}
			}
		}
	} // End SpecialProfileCase1

	// showSPS()
	// return sps
	// Possibly wrong due to no scaling list being built
	sps.Log2MaxFrameNumMinus4 = r.readUe()
	sps.PicOrderCountType = r.readUe()

	if sps.PicOrderCountType == 0 {
		sps.Log2MaxPicOrderCntLSBMin4 = r.readUe()
	} else if sps.PicOrderCountType == 1 {
		sps.DeltaPicOrderAlwaysZeroFlag = r.readBits(1) == 1
		sps.OffsetForNonRefPic = int64(r.readSe())
		sps.OffsetForTopToBottomField = int64(r.readSe())
		sps.NumRefFramesInPicOrderCntCycle = r.readUe()

		for i := 0; i < int(sps.NumRefFramesInPicOrderCntCycle); i++ {
			sps.OffsetForRefFrameList = append(sps.OffsetForRefFrameList, r.readSe())
		}

	}

	sps.MaxNumRefFrames = r.readUe()
	sps.GapsInFrameNumValueAllowed = r.readBits(1) == 1
	sps.PicWidthInMBSMinus1 = r.readUe()
	sps.PicHeightInMapUnitsMinus1 = r.readUe()
	sps.FrameMBSOnlyFlag = r.readBits(1) == 1

	if !sps.FrameMBSOnlyFlag {
		sps.MBAdaptiveFrameFieldFlag = r.readBits(1) == 1
	}

	sps.Direct8x8InferenceFlag = r.readBits(1) == 1
	sps.FrameCroppingFlag = r.readBits(1) == 1

	if sps.FrameCroppingFlag {
		sps.FrameCropLeftOffset = r.readUe()
		sps.FrameCropRightOffset = r.readUe()
		sps.FrameCropTopOffset = r.readUe()
		sps.FrameCropBottomOffset = r.readUe()
	}

	sps.VUIParametersPresentFlag = r.readBits(1) == 1

	if sps.VUIParametersPresentFlag {

	} // End VuiParameters Annex E.1.1

	return &sps, nil
}

// SPS describes a sequence parameter set as defined by section E.1.1 in the
// Specifications.
type VUIParameters struct {
	AspectRatioInfoPresentFlag         bool
	AspectRatioIDC                     uint8
	SARWidth                           uint32
	SARHeight                          uint32
	OverscanInfoPresentFlag            bool
	OverscanAppropriateFlag            bool
	VideoSignalTypePresentFlag         bool
	VideoFormat                        uint8
	VideoFullRangeFlag                 bool
	ColorDescriptionPresentFlag        bool
	ColorPrimaries                     uint8
	TransferCharacteristics            uint8
	MatrixCoefficients                 uint8
	ChromaLocInfoPresentFlag           bool
	ChromaSampleLocTypeTopField        uint64
	ChromaSampleLocTypeBottomField     uint64
	TimingInfoPresentFlag              bool
	NumUnitsInTick                     uint32
	TimeScale                          uint32
	FixedFrameRateFlag                 bool
	NALHRDParametersPresentFlag        bool
	NALHRDParameters                   *HRDParameters
	VCLHRDParametersPresentFlag        bool
	VCLHRDParameters                   *HRDParameters
	LowDelayHRDFlag                    bool
	PicStructPresentFlag               bool
	BitstreamRestrictionFlag           bool
	MotionVectorsOverPicBoundariesFlag bool
	MaxBytesPerPicDenom                uint64
	MaxBitsPerMBDenom                  uint64
	Log2MaxMVLengthHorizontal          uint64
	Log2MaxMVLengthVertical            uint64
	MaxNumReorderFrames                uint64
	MaxDecFrameBuffering               uint64
}

// NewVUIParameters parses video usability information parameters from br
// following the syntax structure specified in section E.1.1, and returns as a
// new VUIParameters.
func NewVUIParameters(br *bits.BitReader) (*VUIParameters, error) {
	p := &VUIParameters{}
	r := newFieldReader(br)

	p.AspectRatioInfoPresentFlag = r.readBits(1) == 1

	if p.AspectRatioInfoPresentFlag {
		p.AspectRatioIDC = uint8(r.readBits(8))

		EXTENDED_SAR := 999
		if int(p.AspectRatioIDC) == EXTENDED_SAR {
			p.SARWidth = uint32(r.readBits(16))
			p.SARHeight = uint32(r.readBits(16))
		}
	}

	p.OverscanInfoPresentFlag = r.readBits(1) == 1

	if p.OverscanInfoPresentFlag {
		p.OverscanAppropriateFlag = r.readBits(1) == 1
	}

	p.VideoSignalTypePresentFlag = r.readBits(1) == 1

	if p.VideoSignalTypePresentFlag {
		p.VideoFormat = uint8(r.readBits(3))
	}

	if p.VideoSignalTypePresentFlag {
		p.VideoFullRangeFlag = r.readBits(1) == 1
		p.ColorDescriptionPresentFlag = r.readBits(1) == 1

		if p.ColorDescriptionPresentFlag {
			p.ColorPrimaries = uint8(r.readBits(8))
			p.TransferCharacteristics = uint8(r.readBits(8))
			p.MatrixCoefficients = uint8(r.readBits(8))
		}
	}
	p.ChromaLocInfoPresentFlag = r.readBits(1) == 1

	if p.ChromaLocInfoPresentFlag {
		p.ChromaSampleLocTypeTopField = uint64(r.readUe())
		p.ChromaSampleLocTypeBottomField = uint64(r.readUe())
	}

	p.TimingInfoPresentFlag = r.readBits(1) == 1

	if p.TimingInfoPresentFlag {
		p.NumUnitsInTick = uint32(r.readBits(32))
		p.TimeScale = uint32(r.readBits(32))
		p.FixedFrameRateFlag = r.readBits(1) == 1
	}

	p.NALHRDParametersPresentFlag = r.readBits(1) == 1

	var err error
	if p.NALHRDParametersPresentFlag {
		p.NALHRDParameters, err = NewHRDParameters(br)
		if err != nil {
			return nil, errors.Wrap(err, "could not get hrdParameters")
		}
	}

	p.VCLHRDParametersPresentFlag = r.readBits(1) == 1

	if p.VCLHRDParametersPresentFlag {
		p.VCLHRDParameters, err = NewHRDParameters(br)
		if err != nil {
			return nil, errors.Wrap(err, "could not get hrdParameters")
		}
	}
	if p.NALHRDParametersPresentFlag || p.VCLHRDParametersPresentFlag {
		p.LowDelayHRDFlag = r.readBits(1) == 1
	}

	p.PicStructPresentFlag = r.readBits(1) == 1
	p.BitstreamRestrictionFlag = r.readBits(1) == 1

	if p.BitstreamRestrictionFlag {
		p.MotionVectorsOverPicBoundariesFlag = r.readBits(1) == 1
		p.MaxBytesPerPicDenom = r.readUe()
		p.MaxBitsPerMBDenom = r.readUe()
		p.Log2MaxMVLengthHorizontal = r.readUe()
		p.Log2MaxMVLengthVertical = r.readUe()
		p.MaxNumReorderFrames = r.readUe()
		p.MaxDecFrameBuffering = r.readUe()
	}
	return p, nil
}

// HRDParameters describes hypothetical reference decoder parameters as defined
// by section E.1.2 in the specifications.
type HRDParameters struct {
	CPBCntMinus1                    uint64
	BitRateScale                    uint8
	CPBSizeScale                    uint8
	BitRateValueMinus1              []uint64
	CPBSizeValueMinus1              []uint64
	CBRFlag                         []bool
	InitialCPBRemovalDelayLenMinus1 uint8
	CPBRemovalDelayLenMinus1        uint8
	DPBOutputDelayLenMinus1         uint8
	TimeOffsetLen                   uint8
}

// NewHRDParameters parses hypothetical reference decoder parameter from br
// following the syntax structure specified in section E.1.2, and returns as a
// new HRDParameters.
func NewHRDParameters(br *bits.BitReader) (*HRDParameters, error) {
	h := &HRDParameters{}
	r := newFieldReader(br)

	h.CPBCntMinus1 = r.readUe()
	h.BitRateScale = uint8(r.readBits(4))
	h.CPBSizeScale = uint8(r.readBits(4))

	// SchedSelIdx E1.2
	for sseli := 0; sseli <= int(h.CPBCntMinus1); sseli++ {
		h.BitRateValueMinus1 = append(h.BitRateValueMinus1, r.readUe())
		h.CPBSizeValueMinus1 = append(h.CPBSizeValueMinus1, r.readUe())

		if v, _ := br.ReadBits(1); v == 1 {
			h.CBRFlag = append(h.CBRFlag, true)
		} else {
			h.CBRFlag = append(h.CBRFlag, false)
		}

		h.InitialCPBRemovalDelayLenMinus1 = uint8(r.readBits(5))
		h.CPBRemovalDelayLenMinus1 = uint8(r.readBits(5))
		h.DPBOutputDelayLenMinus1 = uint8(r.readBits(5))
		h.TimeOffsetLen = uint8(r.readBits(5))
	}

	if r.err() != nil {
		return nil, fmt.Errorf("error from fieldReader: %v", r.err())
	}
	return h, nil
}

func isInList(l []int, term int) bool {
	for _, m := range l {
		if m == term {
			return true
		}
	}
	return false
}

func scalingList(br *bits.BitReader, scalingList []int, sizeOfScalingList int, defaultScalingMatrix []int) error {
	lastScale := 8
	nextScale := 8
	for i := 0; i < sizeOfScalingList; i++ {
		if nextScale != 0 {
			deltaScale, err := readSe(br)
			if err != nil {
				return errors.Wrap(err, "could not parse deltaScale")
			}
			nextScale = (lastScale + deltaScale + 256) % 256
			if i == 0 && nextScale == 0 {
				// Scaling list should use the default list for this point in the matrix
				_ = defaultScalingMatrix
			}
		}
		if nextScale == 0 {
			scalingList[i] = lastScale
		} else {
			scalingList[i] = nextScale
		}
		lastScale = scalingList[i]
	}
	return nil
}