package generator

import (
	"encoding/binary"
	"fmt"
)

type YCbCr10 struct {
	Y  uint32
	Cb uint32
	Cr uint32
}

type PixelSampler func(x, y, width, height, frameIndex int) YCbCr10

type CPUGenerator struct {
	width   int
	height  int
	sampler PixelSampler
}

func NewCPUGenerator(
	width, height uint,
	sampler PixelSampler,
) (*CPUGenerator, error) {
	if width == 0 || height == 0 {
		return nil, fmt.Errorf(
			"width and height must be greater than 0, got: %dx%d",
			width,
			height,
		)
	}
	if width%6 != 0 {
		return nil, fmt.Errorf(
			"width must be divisible by 6, got: %d",
			width,
		)
	}
	if sampler == nil {
		return nil, fmt.Errorf("pixel sampler is nil")
	}

	return &CPUGenerator{
		width:   int(width),
		height:  int(height),
		sampler: sampler,
	}, nil
}

func (g *CPUGenerator) GenerateFrame(dest []byte, frameIndex int) error {
	frameSize := (g.width / 6) * g.height * 16
	if len(dest) < frameSize {
		return fmt.Errorf(
			"cpu: destination is too small: got %d bytes, need %d",
			len(dest),
			frameSize,
		)
	}

	const componentMask uint32 = 0x3ff
	offset := 0
	for y := 0; y < g.height; y++ {
		for x := 0; x < g.width; x += 6 {
			var pixels [6]YCbCr10
			for i := range pixels {
				pixels[i] = g.sampler(x+i, y, g.width, g.height, frameIndex)
			}

			word0 := pixels[0].Cb&componentMask |
				(pixels[0].Y&componentMask)<<10 |
				(pixels[0].Cr&componentMask)<<20
			word1 := pixels[1].Y&componentMask |
				(pixels[2].Cb&componentMask)<<10 |
				(pixels[2].Y&componentMask)<<20
			word2 := pixels[2].Cr&componentMask |
				(pixels[3].Y&componentMask)<<10 |
				(pixels[4].Cb&componentMask)<<20
			word3 := pixels[4].Y&componentMask |
				(pixels[4].Cr&componentMask)<<10 |
				(pixels[5].Y&componentMask)<<20

			binary.LittleEndian.PutUint32(dest[offset:], word0)
			binary.LittleEndian.PutUint32(dest[offset+4:], word1)
			binary.LittleEndian.PutUint32(dest[offset+8:], word2)
			binary.LittleEndian.PutUint32(dest[offset+12:], word3)
			offset += 16
		}
	}

	return nil
}

func (g *CPUGenerator) Close() error { return nil }