package generator

import (
	"encoding/binary"
	"path/filepath"
	"testing"
)

func TestWGPUGenerator(t *testing.T) {
	const width, height = 1920, 1080
	g, err := NewWGPUGenerator(width, height, filepath.Join("..", "..", "kernels", "v210_bars.wgsl"))
	if err != nil {
		t.Fatalf("init: %v", err)
	}
	defer g.Close()

	buf := make([]byte, width*height*8/3)
	if err := g.GenerateFrame(buf, 0); err != nil {
		t.Fatalf("GenerateFrame: %v", err)
	}
	if err := g.GenerateFrame(buf, 1); err != nil {
		t.Fatalf("GenerateFrame 2: %v", err)
	}

	want := [7][3]uint32{
		{721, 512, 512}, {674, 176, 543}, {581, 589, 176},
		{534, 253, 207}, {251, 771, 817}, {204, 435, 848}, {111, 848, 481},
	}
	barOf := func(x int) int {
		if b := x * 7 / width; b < 7 {
			return b
		}
		return 6
	}
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			p := y*width + x
			off := (p / 6) * 16
			w0 := binary.LittleEndian.Uint32(buf[off:])
			w1 := binary.LittleEndian.Uint32(buf[off+4:])
			w2 := binary.LittleEndian.Uint32(buf[off+8:])
			w3 := binary.LittleEndian.Uint32(buf[off+12:])
			var yv, cb, cr uint32
			switch p % 6 {
			case 0:
				yv, cb, cr = (w0>>10)&0x3FF, w0&0x3FF, (w0>>20)&0x3FF
			case 1:
				yv = w1 & 0x3FF
			case 2:
				yv, cb, cr = (w1>>20)&0x3FF, (w1>>10)&0x3FF, w2&0x3FF
			case 3:
				yv = (w2 >> 10) & 0x3FF
			case 4:
				yv, cb, cr = w3&0x3FF, (w2>>20)&0x3FF, (w3>>10)&0x3FF
			case 5:
				yv = (w3 >> 20) & 0x3FF
			}
			b := want[barOf(x)]
			if yv != b[0] {
				t.Fatalf("pixel (%d,%d): Y=%d want %d", x, y, yv, b[0])
			}
			if x%2 == 0 && (cb != b[1] || cr != b[2]) {
				t.Fatalf("pixel (%d,%d): Cb=%d Cr=%d want %d/%d", x, y, cb, cr, b[1], b[2])
			}
		}
	}
}