package frame

import (
	"errors"
	"runtime"
	"unsafe"
)

type FrameUint16 struct {
	properties Properties
	Pts        int64
	Y          []uint16
	Cb         []uint16
	Cr         []uint16
}

func NewUint16FrameFromBytes(properties Properties, pts int64, data []byte) (*FrameUint16, error) {
	if frameLength, _ := properties.ColorSpace.FrameSize(properties.Width, properties.Height); frameLength != len(data) {
		return nil, errors.New("wrong length of data")
	}

	if properties.ColorSpace.BitDepth >= 16 {
		return nil, errors.New("wrong bit depth")
	}

	buf := make([]uint16, len(data)/2)
	copy(buf, unsafe.Slice((*uint16)(unsafe.Pointer(&data[0])), len(data)/2))
	runtime.KeepAlive(data)

	iY := properties.ColorSpace.ChromaSampling.PlaneLumaSamples(properties.Width, properties.Height)
	iCb := properties.ColorSpace.ChromaSampling.PlaneCbSamples(properties.Width, properties.Height)
	iCr := properties.ColorSpace.ChromaSampling.PlaneCrSamples(properties.Width, properties.Height)

	return &FrameUint16{
		properties: properties,
		Y:          buf[:iY],
		Cb:         buf[iY : iY+iCb],
		Cr:         buf[iY+iCb : iY+iCb+iCr],
		Pts:        pts,
	}, nil
}

func (i *FrameUint16) Get16(x, y int) (Y uint16, Cb uint16, Cr uint16) {
	cy, cb, cr := i.GetNative(x, y)

	return cy << (16 - i.properties.ColorSpace.BitDepth), cb << (16 - i.properties.ColorSpace.BitDepth), cr << (16 - i.properties.ColorSpace.BitDepth)
}

func (i *FrameUint16) Get8(x, y int) (Y uint8, Cb uint8, Cr uint8) {
	cy, cb, cr := i.GetNative(x, y)

	return uint8(cy >> (i.properties.ColorSpace.BitDepth - 8)), uint8(cb >> (i.properties.ColorSpace.BitDepth - 8)), uint8(cr >> (i.properties.ColorSpace.BitDepth - 8))
}

func (i *FrameUint16) Properties() Properties {
	return i.properties
}

func (i *FrameUint16) PTS() int64 {
	return i.Pts
}

func (i *FrameUint16) GetNative(x, y int) (Y uint16, Cb uint16, Cr uint16) {
	Yindex := x + y*i.properties.Width

	Cwidth := (i.properties.Width * int(i.properties.ColorSpace.ChromaSampling.A)) / int(i.properties.ColorSpace.ChromaSampling.J)
	if i.properties.ColorSpace.ChromaSampling.B == 0 {
		y /= 2
	}
	Cindex := (x*int(i.properties.ColorSpace.ChromaSampling.A))/int(i.properties.ColorSpace.ChromaSampling.J) + y*Cwidth
	Y = i.Y[Yindex]
	Cb = i.Cb[Cindex]
	Cr = i.Cr[Cindex]
	return
}

func (i *FrameUint16) GetNativeLuma() []uint16 {
	return i.Y
}

func (i *FrameUint16) GetNativeCb() []uint16 {
	return i.Cb
}

func (i *FrameUint16) GetNativeCr() []uint16 {
	return i.Cr
}