package media

import (
	"bytes"
	"errors"
	"image"
	"image/color"
	"image/gif"
	"image/jpeg"
	"image/png"
	"strings"
	"testing"
)

// pngBytes builds a w×h test image: left half red, right half a light grey, so
// a downscale that averages produces a predictable mid value and a scaler that
// silently returns the wrong region is visible.
func pngBytes(t *testing.T, w, h int) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			if x < w/2 {
				img.Set(x, y, color.RGBA{255, 0, 0, 255})
			} else {
				img.Set(x, y, color.RGBA{200, 200, 200, 255})
			}
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

func TestSniffImage(t *testing.T) {
	cases := []struct {
		name string
		data []byte
		want string
	}{
		{"png", pngBytes(t, 4, 4), "image/png"},
		{"jpeg", jpegBytes(t, 4, 4), "image/jpeg"},
		{"gif", gifBytes(t, 4, 4), "image/gif"},
	}
	for _, c := range cases {
		got, err := SniffImage(c.data)
		if err != nil {
			t.Errorf("%s: %v", c.name, err)
			continue
		}
		if got != c.want {
			t.Errorf("%s: got %q want %q", c.name, got, c.want)
		}
	}
}

// webp is common from Telegram and there is no stdlib decoder, so it must be
// refused by name rather than mis-sniffed or fed to a model as noise.
func TestSniffRefusesWebpByName(t *testing.T) {
	webp := append([]byte("RIFF\x00\x00\x00\x00WEBP"), make([]byte, 8)...)
	_, err := SniffImage(webp)
	if !errors.Is(err, ErrUnsupportedImage) {
		t.Fatalf("got %v, want ErrUnsupportedImage", err)
	}
	if !strings.Contains(err.Error(), "webp") {
		t.Errorf("error does not name the format: %v", err)
	}
}

func TestSniffRefusesGarbage(t *testing.T) {
	for _, data := range [][]byte{nil, []byte("hello"), []byte("\x00\x01\x02\x03")} {
		if _, err := SniffImage(data); !errors.Is(err, ErrUnsupportedImage) {
			t.Errorf("SniffImage(%q) = %v", data, err)
		}
	}
}

func TestPrepareImageDownscalesLongestEdge(t *testing.T) {
	im, err := PrepareImage(pngBytes(t, 2000, 1000), "web:upload", 500)
	if err != nil {
		t.Fatalf("prepare: %v", err)
	}
	if im.Width != 500 || im.Height != 250 {
		t.Errorf("got %dx%d, want 500x250", im.Width, im.Height)
	}
	if _, err := jpeg.Decode(bytes.NewReader(im.JPEG)); err != nil {
		t.Errorf("output is not decodable jpeg: %v", err)
	}
	if im.Source != "web:upload" {
		t.Errorf("source lost: %q", im.Source)
	}
}

// Tall images scale on the other axis; a scaler that only handles landscape is
// the classic version of this bug.
func TestPrepareImageHandlesPortrait(t *testing.T) {
	im, err := PrepareImage(pngBytes(t, 400, 1600), "telegram", 800)
	if err != nil {
		t.Fatalf("prepare: %v", err)
	}
	if im.Height != 800 || im.Width != 200 {
		t.Errorf("got %dx%d, want 200x800", im.Width, im.Height)
	}
}

func TestPrepareImageNeverEnlarges(t *testing.T) {
	im, err := PrepareImage(pngBytes(t, 64, 32), "telegram", 896)
	if err != nil {
		t.Fatalf("prepare: %v", err)
	}
	if im.Width != 64 || im.Height != 32 {
		t.Errorf("got %dx%d, want the original 64x32", im.Width, im.Height)
	}
}

// A degenerate strip must not scale to zero on the short axis — jpeg.Encode
// fails on a zero-height image, which would turn a weird screenshot into a
// hard error.
func TestPrepareImageClampsDegenerateAspect(t *testing.T) {
	im, err := PrepareImage(pngBytes(t, 2000, 2), "web:upload", 100)
	if err != nil {
		t.Fatalf("prepare: %v", err)
	}
	if im.Height < 1 || im.Width != 100 {
		t.Errorf("got %dx%d", im.Width, im.Height)
	}
}

// Transparent pixels composite onto white, not black: the common case is a
// screenshot or a diagram, and dark-on-black is unreadable to the model.
func TestPrepareImageFlattensAlphaOntoWhite(t *testing.T) {
	img := image.NewRGBA(image.Rect(0, 0, 8, 8)) // fully transparent
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	im, err := PrepareImage(buf.Bytes(), "web:upload", 8)
	if err != nil {
		t.Fatalf("prepare: %v", err)
	}
	decoded, err := jpeg.Decode(bytes.NewReader(im.JPEG))
	if err != nil {
		t.Fatal(err)
	}
	r, g, b, _ := decoded.At(4, 4).RGBA()
	if r>>8 < 240 || g>>8 < 240 || b>>8 < 240 {
		t.Errorf("transparent pixel became rgb(%d,%d,%d), want near-white", r>>8, g>>8, b>>8)
	}
}

func TestPrepareImageRejectsEmpty(t *testing.T) {
	if _, err := PrepareImage(nil, "x", 0); !errors.Is(err, ErrEmpty) {
		t.Errorf("got %v, want ErrEmpty", err)
	}
}

func TestDataURIIsAJPEGDataURI(t *testing.T) {
	im, err := PrepareImage(pngBytes(t, 16, 16), "x", 0)
	if err != nil {
		t.Fatal(err)
	}
	uri := im.DataURI()
	if !strings.HasPrefix(uri, "data:image/jpeg;base64,") {
		t.Fatalf("bad prefix: %.40s", uri)
	}
	if len(uri) <= len("data:image/jpeg;base64,") {
		t.Error("data uri carries no payload")
	}
}

func jpegBytes(t *testing.T, w, h int) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, nil); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

func gifBytes(t *testing.T, w, h int) []byte {
	t.Helper()
	img := image.NewPaletted(image.Rect(0, 0, w, h), []color.Color{color.Black, color.White})
	var buf bytes.Buffer
	if err := gif.Encode(&buf, img, nil); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}