package ai

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

// EncodeScreenshot converts raw RGBA pixel data to a JPEG image.
// If the source dimensions exceed maxWidth x maxHeight, the image is
// downscaled using nearest-neighbor sampling (fast, no external deps).
// Returns the JPEG bytes.
func EncodeScreenshot(rgba []byte, srcWidth, srcHeight, maxWidth, maxHeight, quality int) ([]byte, error) {
	expectedLen := srcWidth * srcHeight * 4
	if len(rgba) < expectedLen {
		return nil, fmt.Errorf("RGBA buffer too small: got %d bytes, expected %d for %dx%d", len(rgba), expectedLen, srcWidth, srcHeight)
	}

	if quality <= 0 || quality > 100 {
		quality = 75
	}

	// Create source image from RGBA buffer
	src := image.NewRGBA(image.Rect(0, 0, srcWidth, srcHeight))
	copy(src.Pix, rgba[:expectedLen])

	// Determine output dimensions
	dstWidth, dstHeight := srcWidth, srcHeight
	if srcWidth > maxWidth || srcHeight > maxHeight {
		dstWidth, dstHeight = fitDimensions(srcWidth, srcHeight, maxWidth, maxHeight)
	}

	var img image.Image = src

	// Downscale if needed using nearest-neighbor sampling
	if dstWidth != srcWidth || dstHeight != srcHeight {
		dst := image.NewRGBA(image.Rect(0, 0, dstWidth, dstHeight))
		for y := 0; y < dstHeight; y++ {
			srcY := y * srcHeight / dstHeight
			for x := 0; x < dstWidth; x++ {
				srcX := x * srcWidth / dstWidth
				srcIdx := (srcY*srcWidth + srcX) * 4
				dstIdx := (y*dstWidth + x) * 4
				dst.Pix[dstIdx+0] = src.Pix[srcIdx+0] // R
				dst.Pix[dstIdx+1] = src.Pix[srcIdx+1] // G
				dst.Pix[dstIdx+2] = src.Pix[srcIdx+2] // B
				dst.Pix[dstIdx+3] = src.Pix[srcIdx+3] // A
			}
		}
		img = dst
	}

	// Encode to JPEG
	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: quality}); err != nil {
		return nil, fmt.Errorf("encode JPEG: %w", err)
	}

	return buf.Bytes(), nil
}

// fitDimensions calculates the largest dimensions that fit within max bounds
// while preserving aspect ratio.
func fitDimensions(srcW, srcH, maxW, maxH int) (int, int) {
	ratio := float64(srcW) / float64(srcH)
	w, h := maxW, int(float64(maxW)/ratio)
	if h > maxH {
		h = maxH
		w = int(float64(maxH) * ratio)
	}
	if w <= 0 {
		w = 1
	}
	if h <= 0 {
		h = 1
	}
	return w, h
}