/*
DESCRIPTION
  nalunit.go provides structures for a NAL unit as well as it's extensions.

AUTHORS
  Saxon Nelson-Milton <saxon@ausocean.org>, The Australian Ocean Laboratory (AusOcean)
  mrmod <mcmoranbjr@gmail.com>
*/

package h264dec

import (
	"io"

	"github.com/pkg/errors"

	"bitbucket.org/ausocean/av/codec/h264/h264dec/bits"
)

// MVCExtension describes a NAL unit header multiview video coding extension, as
// defined in section H.7.3.1.1 of the specifications.
type MVCExtension struct {
	NonIdrFlag     bool
	PriorityID     int
	ViewID         int
	TemporalID     int
	AnchorPicFlag  bool
	InterViewFlag  bool
	ReservedOneBit int
}

// NewMVCExtension parses a NAL unit header multiview video coding extension
// from br following the syntax structure specified in section H.7.3.1.1, and
// returns as a new MVCExtension.
func NewMVCExtension(br *bits.BitReader) (*MVCExtension, error) {
	e := &MVCExtension{}
	var err error

	e.NonIdrFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read NonIdrFlag")
	}

	e.PriorityID, err = br.ReadBitsInt(6)
	if err != nil {
		return nil, errors.Wrap(err, "could not read PriorityId")
	}

	e.ViewID, err = br.ReadBitsInt(10)
	if err != nil {
		return nil, errors.Wrap(err, "could not read ViewId")
	}

	e.TemporalID, err = br.ReadBitsInt(3)
	if err != nil {
		return nil, errors.Wrap(err, "could not read TemporalId")
	}

	e.AnchorPicFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read AnchorPicFlag")
	}

	e.InterViewFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read InterViewFlag")
	}

	e.ReservedOneBit, err = br.ReadBitsInt(1)
	if err != nil {
		return nil, errors.Wrap(err, "could not read ReservedOneBit")
	}

	return e, nil
}

// ThreeDAVCExtension describes a NAL unit header 3D advanced video coding
// extension, as defined in section J.7.3.1.1 of the specifications.
type ThreeDAVCExtension struct {
	ViewIdx       int
	DepthFlag     bool
	NonIdrFlag    bool
	TemporalID    int
	AnchorPicFlag bool
	InterViewFlag bool
}

// NewThreeDAVCExtension parses a NAL unit header 3D advanced video coding
// extension from br following the syntax structure specified in section
// J.7.3.1.1, and returns as a new ThreeDAVCExtension.
func NewThreeDAVCExtension(br *bits.BitReader) (*ThreeDAVCExtension, error) {
	e := &ThreeDAVCExtension{}
	var err error

	e.ViewIdx, err = br.ReadBitsInt(8)
	if err != nil {
		return nil, errors.Wrap(err, "could not read ViewIdx")
	}

	e.DepthFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read DepthFlag")
	}

	e.NonIdrFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read NonIdrFlag")
	}

	e.TemporalID, err = br.ReadBitsInt(3)
	if err != nil {
		return nil, errors.Wrap(err, "could not read TemporalId")
	}

	e.AnchorPicFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read AnchorPicFlag")
	}

	e.InterViewFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read InterViewFlag")
	}

	return e, nil
}

// SVCExtension describes a NAL unit header scalable video coding extension, as
// defined in section G.7.3.1.1 of the specifications.
type SVCExtension struct {
	IdrFlag              bool
	PriorityID           int
	NoInterLayerPredFlag bool
	DependencyID         int
	QualityID            int
	TemporalID           int
	UseRefBasePicFlag    bool
	DiscardableFlag      bool
	OutputFlag           bool
	ReservedThree2Bits   int
}

// NewSVCExtension parses a NAL unit header scalable video coding extension from
// br following the syntax structure specified in section G.7.3.1.1, and returns
// as a new SVCExtension.
func NewSVCExtension(br *bits.BitReader) (*SVCExtension, error) {
	e := &SVCExtension{}
	var err error

	e.IdrFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read IdrFlag")
	}

	e.PriorityID, err = br.ReadBitsInt(6)
	if err != nil {
		return nil, errors.Wrap(err, "could not read PriorityId")
	}

	e.NoInterLayerPredFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read NoInterLayerPredFlag")
	}

	e.DependencyID, err = br.ReadBitsInt(3)
	if err != nil {
		return nil, errors.Wrap(err, "could not read DependencyId")
	}

	e.QualityID, err = br.ReadBitsInt(4)
	if err != nil {
		return nil, errors.Wrap(err, "could not read QualityId")
	}

	e.TemporalID, err = br.ReadBitsInt(3)
	if err != nil {
		return nil, errors.Wrap(err, "could not read TemporalId")
	}

	e.UseRefBasePicFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read UseRefBasePicFlag")
	}

	e.DiscardableFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read DiscardableFlag")
	}

	e.OutputFlag, err = br.ReadBool()
	if err != nil {
		return nil, errors.Wrap(err, "could not read OutputFlag")
	}

	e.ReservedThree2Bits, err = br.ReadBitsInt(2)
	if err != nil {
		return nil, errors.Wrap(err, "could not read ReservedThree2Bits")
	}

	return e, nil
}

// NALUnit describes a network abstraction layer unit, as defined in section
// 7.3.1 of the specifications.
type NALUnit struct {
	ForbiddenZeroBit             int
	RefIdc                       int
	Type                         int
	SVCExtensionFlag             bool
	AVC3DExtensionFlag           bool
	SVCExtension                 *SVCExtension
	ThreeDAVCExtension           *ThreeDAVCExtension
	MVCExtension                 *MVCExtension
	EmulationPreventionThreeByte byte
	RBSP                         []byte
}

// NewNALUnit parses a network abstraction layer unit from br following the
// syntax structure specified in section 7.3.1, and returns as a new NALUnit.
func NewNALUnit(br *bits.BitReader) (*NALUnit, error) {
	n := &NALUnit{}

	err := readFields(br, []field{
		{&n.ForbiddenZeroBit, "ForbiddenZeroBit", 1},
		{&n.RefIdc, "NalRefIdc", 2},
		{&n.Type, "NalUnitType", 5},
	})
	if err != nil {
		return nil, err
	}

	// TODO: use consts for the NAL types here
	if n.Type == 14 || n.Type == 20 || n.Type == 21 {
		if n.Type != 21 {
			n.SVCExtensionFlag, err = br.ReadBool()
			if err != nil {
				return nil, errors.Wrap(err, "could not read SVCExtensionFlag")
			}
		} else {
			n.AVC3DExtensionFlag, err = br.ReadBool()
			if err != nil {
				return nil, errors.Wrap(err, "could not read AVC3DExtensionFlag")
			}
		}
		if n.SVCExtensionFlag {
			n.SVCExtension, err = NewSVCExtension(br)
			if err != nil {
				return nil, errors.Wrap(err, "could not parse SVCExtension")
			}
		} else if n.AVC3DExtensionFlag {
			n.ThreeDAVCExtension, err = NewThreeDAVCExtension(br)
			if err != nil {
				return nil, errors.Wrap(err, "could not parse ThreeDAVCExtension")
			}
		} else {
			n.MVCExtension, err = NewMVCExtension(br)
			if err != nil {
				return nil, errors.Wrap(err, "could not parse MVCExtension")
			}
		}
	}

	for moreRBSPData(br) {
		next3Bytes, err := br.PeekBits(24)

		// If PeekBits cannot get 3 bytes, but there still might be 2 bytes left in
		// the source, we will get an io.EOF; we wish to ignore this and continue.
		// The call to moreRBSPData will determine when we have reached the end of
		// the NAL unit.
		if err != nil && errors.Cause(err) != io.EOF {
			return nil, errors.Wrap(err, "could not Peek next 3 bytes")
		}

		if next3Bytes == 0x000003 {
			for j := 0; j < 2; j++ {
				rbspByte, err := br.ReadBits(8)
				if err != nil {
					return nil, errors.Wrap(err, "could not read rbspByte")
				}
				n.RBSP = append(n.RBSP, byte(rbspByte))
			}

			// Read Emulation prevention three byte.
			eptByte, err := br.ReadBits(8)
			if err != nil {
				return nil, errors.Wrap(err, "could not read eptByte")
			}
			n.EmulationPreventionThreeByte = byte(eptByte)
		} else {
			b, err := br.ReadBits(8)
			if err != nil {
				return nil, errors.Wrap(err, "could not read RBSP byte")
			}
			n.RBSP = append(n.RBSP, byte(b))
		}
	}

	return n, nil
}