package media

import (
	"bytes"
	"encoding/binary"
	"errors"
	"hash/crc32"
	"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()
}

// A decode bomb is a small file. Nothing bounded pixels before decoding, so a
// 20000x20000 PNG of flat colour — a few hundred kilobytes on the wire, well
// under the byte cap — decoded to 1.6 GB and then allocated another 1.6 GB to
// flatten, in the process that owns the database and the socket.
func TestPrepareImageRefusesADecodeBomb(t *testing.T) {
	// The header is what is checked, so the test writes a real header and
	// truncated pixel data: reaching the decode at all is the failure.
	var buf bytes.Buffer
	if err := png.Encode(&buf, image.NewGray(image.Rect(0, 0, 1, 1))); err != nil {
		t.Fatal(err)
	}
	bomb := forgePNGSize(t, buf.Bytes(), 20000, 20000)
	_, err := PrepareImage(bomb, "telegram", 0)
	if !errors.Is(err, ErrTooManyPixels) {
		t.Fatalf("err = %v, want ErrTooManyPixels", err)
	}
	// A phone photo is not a bomb.
	if _, err := PrepareImage(pngBytes(t, 64, 48), "telegram", 0); err != nil {
		t.Fatalf("an ordinary image was refused: %v", err)
	}
}

// forgePNGSize rewrites the IHDR width and height (and its CRC) of a valid PNG,
// which is how a header claiming 400 megapixels is produced without writing
// 400 megapixels.
func forgePNGSize(t *testing.T, src []byte, w, h uint32) []byte {
	t.Helper()
	out := append([]byte(nil), src...)
	// 8 byte signature, 4 byte length, 4 byte "IHDR", then width and height.
	const ihdr = 8 + 4 + 4
	binary.BigEndian.PutUint32(out[ihdr:], w)
	binary.BigEndian.PutUint32(out[ihdr+4:], h)
	crc := crc32.ChecksumIEEE(out[8+4 : ihdr+13])
	binary.BigEndian.PutUint32(out[ihdr+13:], crc)
	return out
}

// Transparency still composites onto white, which is what makes a screenshot
// readable. The old code did that with a full-size intermediate; the scaler
// walks the source instead and must give the same answer.
func TestPrepareImageFlattensOntoWhite(t *testing.T) {
	img := image.NewRGBA(image.Rect(0, 0, 8, 8))
	// Fully transparent everywhere: over white, that is white.
	data := encodePNG(t, img)
	out, err := PrepareImage(data, "test", 4)
	if err != nil {
		t.Fatal(err)
	}
	dec, err := jpeg.Decode(bytes.NewReader(out.JPEG))
	if err != nil {
		t.Fatal(err)
	}
	r, g, b, _ := dec.At(2, 2).RGBA()
	if r>>8 < 240 || g>>8 < 240 || b>>8 < 240 {
		t.Fatalf("transparent pixel came out %d,%d,%d, want white", r>>8, g>>8, b>>8)
	}
}

func encodePNG(t *testing.T, img image.Image) []byte {
	t.Helper()
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}