package h264

import (
	"io"

	"github.com/pkg/errors"

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

type MVCExtension struct {
	NonIdrFlag     bool
	PriorityID     int
	ViewID         int
	TemporalID     int
	AnchorPicFlag  bool
	InterViewFlag  bool
	ReservedOneBit int
}

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
}

type ThreeDAVCExtension struct {
	ViewIdx       int
	DepthFlag     bool
	NonIdrFlag    bool
	TemporalID    int
	AnchorPicFlag bool
	InterViewFlag bool
}

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
}

type SVCExtension struct {
	IdrFlag              bool
	PriorityId           int
	NoInterLayerPredFlag bool
	DependencyId         int
	QualityId            int
	TemporalId           int
	UseRefBasePicFlag    bool
	DiscardableFlag      bool
	OutputFlag           bool
	ReservedThree2Bits   int
}

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
}

type NALUnit struct {
	ForbiddenZeroBit             int
	RefIdc                       int
	Type                         int
	SVCExtensionFlag             int
	AVC3dExtensionFlag           int
	SVCExtension                 *SVCExtension
	ThreeDAVCExtension           *ThreeDAVCExtension
	MVCExtension                 *MVCExtension
	EmulationPreventionThreeByte byte
	RBSP                         []byte
}

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
	}

	var headBytes, rbspBytes = 1, 0

	// 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")
			}
			n.HeaderBytes += 3
		} else if n.AVC3DExtensionFlag {
			n.ThreeDAVCExtension, err = NewThreeDAVCExtension(br)
			if err != nil {
				return nil, errors.Wrap(err, "could not parse ThreeDAVCExtension")
			}
			n.HeaderBytes += 2
		} else {
			n.MVCExtension, err = NewMVCExtension(br)
			if err != nil {
				return nil, errors.Wrap(err, "could not parse MVCExtension")
			}
			n.HeaderBytes += 3
		}
	}

	for moreRBSPData(br) {
		next3Bytes, err := br.PeekBits(24)
		if err != nil && errors.Cause(err) != io.EOF {
			return nil, errors.Wrap("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))
			}
			i += 2

			// 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
}