package ai

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

func TestEncodeScreenshot(t *testing.T) {
	width, height := 100, 80

	// Create a test RGBA buffer (red pixels)
	rgba := make([]byte, width*height*4)
	for i := 0; i < len(rgba); i += 4 {
		rgba[i+0] = 255 // R
		rgba[i+1] = 0   // G
		rgba[i+2] = 0   // B
		rgba[i+3] = 255 // A
	}

	jpegData, err := EncodeScreenshot(rgba, width, height, 200, 200, 80)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Check JPEG magic bytes (FF D8)
	if len(jpegData) < 2 {
		t.Fatal("JPEG data too small")
	}
	if jpegData[0] != 0xFF || jpegData[1] != 0xD8 {
		t.Errorf("expected JPEG magic bytes FF D8, got %02X %02X", jpegData[0], jpegData[1])
	}

	// Decode and verify dimensions (no downscale needed in this case)
	img, err := jpeg.Decode(bytes.NewReader(jpegData))
	if err != nil {
		t.Fatalf("decode JPEG: %v", err)
	}
	bounds := img.Bounds()
	if bounds.Dx() != width || bounds.Dy() != height {
		t.Errorf("expected %dx%d, got %dx%d", width, height, bounds.Dx(), bounds.Dy())
	}
}

func TestEncodeScreenshotDownscale(t *testing.T) {
	srcWidth, srcHeight := 1920, 1080
	maxWidth, maxHeight := 1280, 720

	// Create a test RGBA buffer (gradient)
	rgba := make([]byte, srcWidth*srcHeight*4)
	for y := 0; y < srcHeight; y++ {
		for x := 0; x < srcWidth; x++ {
			idx := (y*srcWidth + x) * 4
			rgba[idx+0] = byte(x % 256)     // R
			rgba[idx+1] = byte(y % 256)     // G
			rgba[idx+2] = byte((x + y) % 256) // B
			rgba[idx+3] = 255               // A
		}
	}

	jpegData, err := EncodeScreenshot(rgba, srcWidth, srcHeight, maxWidth, maxHeight, 75)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Check JPEG magic bytes
	if jpegData[0] != 0xFF || jpegData[1] != 0xD8 {
		t.Errorf("expected JPEG magic bytes FF D8, got %02X %02X", jpegData[0], jpegData[1])
	}

	// Decode and verify dimensions are within max bounds
	img, err := jpeg.Decode(bytes.NewReader(jpegData))
	if err != nil {
		t.Fatalf("decode JPEG: %v", err)
	}
	bounds := img.Bounds()

	if bounds.Dx() > maxWidth {
		t.Errorf("output width %d exceeds max %d", bounds.Dx(), maxWidth)
	}
	if bounds.Dy() > maxHeight {
		t.Errorf("output height %d exceeds max %d", bounds.Dy(), maxHeight)
	}

	// Should maintain 16:9 ratio approximately
	expectedWidth := 1280
	expectedHeight := 720
	if bounds.Dx() != expectedWidth || bounds.Dy() != expectedHeight {
		t.Errorf("expected %dx%d, got %dx%d", expectedWidth, expectedHeight, bounds.Dx(), bounds.Dy())
	}
}

func TestEncodeScreenshotBufferTooSmall(t *testing.T) {
	_, err := EncodeScreenshot([]byte{0, 0, 0, 0}, 100, 100, 200, 200, 75)
	if err == nil {
		t.Error("expected error for buffer too small")
	}
}

func TestFitDimensions(t *testing.T) {
	tests := []struct {
		srcW, srcH, maxW, maxH int
		wantW, wantH           int
	}{
		{1920, 1080, 1280, 720, 1280, 720},    // 16:9 fits exactly
		{1920, 1080, 800, 600, 800, 450},      // width-constrained
		{1080, 1920, 800, 600, 337, 600},      // height-constrained (portrait)
		{100, 100, 200, 200, 200, 200},        // smaller than max (but func called with > check)
	}

	for _, tt := range tests {
		w, h := fitDimensions(tt.srcW, tt.srcH, tt.maxW, tt.maxH)
		if w != tt.wantW || h != tt.wantH {
			t.Errorf("fitDimensions(%d,%d,%d,%d) = %d,%d, want %d,%d",
				tt.srcW, tt.srcH, tt.maxW, tt.maxH, w, h, tt.wantW, tt.wantH)
		}
	}
}