package frame

import (
	"errors"
	"sync"
)

type Pool struct {
	p          sync.Pool
	properties Properties
	frameSize  int
}

func NewPool(properties Properties) (*Pool, error) {
	p := &Pool{
		properties: properties,
	}

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

	frameSize, err := properties.ColorSpace.FrameSize(properties.Width, properties.Height)
	if err != nil {
		return nil, err
	}

	p.frameSize = frameSize

	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)

	// 16-bit frame
	if properties.ColorSpace.BitDepth > 8 {
		p.p.New = func() any {
			buf := make([]uint16, p.frameSize/2)
			return &fUint16{
				ret:        p.Put,
				properties: properties,
				Y:          buf[:iY],
				Cb:         buf[iY : iY+iCb],
				Cr:         buf[iY+iCb : iY+iCb+iCr],
			}
		}
	} else {
		p.p.New = func() any {
			buf := make([]uint8, p.frameSize)
			return &fUint8{
				ret:        p.Put,
				properties: properties,
				Y:          buf[:iY],
				Cb:         buf[iY : iY+iCb],
				Cr:         buf[iY+iCb : iY+iCb+iCr],
			}
		}
	}

	return p, nil
}

func (p *Pool) Properties() Properties {
	return p.properties
}

func (p *Pool) Get(pts, nextPts int64) Frame {
	switch tf := p.p.Get().(type) {
	case *fUint16:
		tf.Pts = pts
		tf.NextPts = nextPts
		return tf
	case *fUint8:
		tf.Pts = pts
		tf.NextPts = nextPts
		return tf
	default:
		panic("unsupported type")
	}
}

func (p *Pool) Put(f Frame) {
	switch tf := f.(type) {
	case *fUint16:
		if tf.properties != p.properties {
			panic("unsupported properties")
		}
		if (len(tf.Y)+len(tf.Cb)+len(tf.Cr))*2 != p.frameSize {
			panic("unsupported data size")
		}
		p.p.Put(tf)
	case *fUint8:
		if tf.properties != p.properties {
			panic("unsupported properties")
		}
		if (len(tf.Y) + len(tf.Cb) + len(tf.Cr)) != p.frameSize {
			panic("unsupported data size")
		}
		p.p.Put(tf)
	default:
		panic("unsupported type")
	}
}