package protocol

import (
	"bytes"
	"image"
	"image/color"
	"testing"
)

// minimalRGB builds a tiny RGBA test image with a known gradient.
func minimalRGB(w, h int) *image.RGBA {
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			img.SetRGBA(x, y, color.RGBA{R: uint8(x * 4), G: uint8(y * 4), B: 128, A: 255})
		}
	}
	return img
}

// minimalYCbCr builds a 420 subsampled image with a vertical colour split.
func minimalYCbCr(w, h int) *image.YCbCr {
	img := image.NewYCbCr(image.Rect(0, 0, w, h), image.YCbCrSubsampleRatio420)
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			r, g, b := uint8(0), uint8(0), uint8(0)
			switch {
			case x < w/3:
				r = 200
			case x < w*2/3:
				g = 200
			default:
				b = 200
			}
			img.Y[y*img.YStride+x], _, _ = color.RGBToYCbCr(r, g, b)
		}
	}
	// Subsample chroma from the known dominant colours per region.
	for y := 0; y < h/2; y++ {
		for x := 0; x < w/2; x++ {
			idx := (y*img.CStride + x)
			switch {
			case x < w/6:
				_, cb, cr := color.RGBToYCbCr(200, 0, 0)
				img.Cb[idx], img.Cr[idx] = cb, cr
			case x < w*2/6:
				_, cb, cr := color.RGBToYCbCr(0, 200, 0)
				img.Cb[idx], img.Cr[idx] = cb, cr
			default:
				_, cb, cr := color.RGBToYCbCr(0, 0, 200)
				img.Cb[idx], img.Cr[idx] = cb, cr
			}
		}
	}
	return img
}

// TestJPEGEncodeRGBA checks the RGBA path produces a decodable JPEG.
func TestJPEGEncodeRGBA(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	img := minimalRGB(64, 48)
	buf, err := enc.Encode(img, 80)
	if err != nil {
		t.Fatalf("Encode: %v", err)
	}
	if len(buf) == 0 {
		t.Fatal("empty JPEG output")
	}
	if !bytes.Equal(buf[:2], []byte{0xFF, 0xD8}) {
		t.Errorf("bad JPEG SOI marker: %x", buf[:2])
	}
	if !bytes.Equal(buf[len(buf)-2:], []byte{0xFF, 0xD9}) {
		t.Errorf("bad JPEG EOI marker")
	}
}

// TestJPEGEncodeYCbCr checks the YUV path for all three subsampling modes.
func TestJPEGEncodeYCbCr(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	img := minimalYCbCr(64, 48)
	if _, err := enc.Encode(img, 80); err != nil {
		t.Fatalf("Encode (420): %v", err)
	}

	img422 := minimalYCbCr(64, 48)
	img422.SubsampleRatio = image.YCbCrSubsampleRatio422
	if _, err := enc.Encode(img422, 80); err != nil {
		t.Fatalf("Encode (422): %v", err)
	}

	img444 := minimalYCbCr(64, 48)
	img444.SubsampleRatio = image.YCbCrSubsampleRatio444
	if _, err := enc.Encode(img444, 80); err != nil {
		t.Fatalf("Encode (444): %v", err)
	}
}

// TestJPEGEncodeQualityClamp ensures out-of-range qualities are clamped.
func TestJPEGEncodeQualityClamp(t *testing.T) {
	enc, err := NewJPEGEncoder()
	if err != nil {
		t.Fatalf("NewJPEGEncoder: %v", err)
	}
	defer enc.Close()

	img := minimalRGB(16, 16)
	for _, q := range []int{-5, 0, 101, 200} {
		if buf, err := enc.Encode(img, q); err != nil || len(buf) == 0 {
			t.Errorf("quality %d: err=%v len=%d", q, err, len(buf))
		}
	}
}