package generator

import (
	"encoding/binary"
	"fmt"
)

type YCbCr10 struct {
	Y  uint32
	Cb uint32
	Cr uint32
}

// FrameRenderer writes a complete frame or patches part of an existing frame.
type FrameRenderer func(dst []byte, width, height, frameIndex int) error

type CPUGenerator struct {
	width  int
	height int
	base   []byte
	patch  FrameRenderer
}

func NewCPUGenerator(
	width, height uint,
	baseRenderer FrameRenderer,
	patch FrameRenderer,
) (*CPUGenerator, error) {
	if width == 0 || height == 0 {
		return nil, fmt.Errorf("cpu: width and height must be greater than zero, got %dx%d", width, height)
	}
	if width%6 != 0 {
		return nil, fmt.Errorf("cpu: width must be divisible by 6, got %d", width)
	}
	if baseRenderer == nil {
		return nil, fmt.Errorf("cpu: base renderer is nil")
	}

	g := &CPUGenerator{
		width:  int(width),
		height: int(height),
		base:   make([]byte, int(width/6*height)*16),
		patch:  patch,
	}
	if err := baseRenderer(g.base, g.width, g.height, 0); err != nil {
		return nil, fmt.Errorf("cpu: render base frame: %w", err)
	}
	return g, nil
}

func (g *CPUGenerator) GenerateFrame(dst []byte, frameIndex int) error {
	if len(dst) < len(g.base) {
		return fmt.Errorf(
			"cpu: destination is too small: got %d bytes, need %d",
			len(dst),
			len(g.base),
		)
	}

	copy(dst, g.base)

	if g.patch != nil {
		if err := g.patch(dst[:len(g.base)], g.width, g.height, frameIndex); err != nil {
			return fmt.Errorf("cpu: patch frame %d: %w", frameIndex, err)
		}
	}

	return nil
}

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

func packV210Block(dst []byte, pixels [6]YCbCr10) {
	const mask uint32 = 0x3ff
	word0 := pixels[0].Cb&mask | (pixels[0].Y&mask)<<10 | (pixels[0].Cr&mask)<<20
	word1 := pixels[1].Y&mask | (pixels[2].Cb&mask)<<10 | (pixels[2].Y&mask)<<20
	word2 := pixels[2].Cr&mask | (pixels[3].Y&mask)<<10 | (pixels[4].Cb&mask)<<20
	word3 := pixels[4].Y&mask | (pixels[4].Cr&mask)<<10 | (pixels[5].Y&mask)<<20

	binary.LittleEndian.PutUint32(dst, word0)
	binary.LittleEndian.PutUint32(dst[4:], word1)
	binary.LittleEndian.PutUint32(dst[8:], word2)
	binary.LittleEndian.PutUint32(dst[12:], word3)
}