package backend

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"hash/crc32"
	"image"
	"image/color"
	"image/gif"
	"image/jpeg"
	"image/png"
	"sync"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/datastore"
)

type capturingImageExtractor struct {
	mu   sync.Mutex
	req  ImageExtractRequest
	text string
}

func (e *capturingImageExtractor) ExtractImageText(_ context.Context, req ImageExtractRequest) (string, ImageExtractUsage, error) {
	e.mu.Lock()
	e.req = ImageExtractRequest{
		ContentType: req.ContentType,
		Data:        append([]byte(nil), req.Data...),
	}
	e.mu.Unlock()
	return e.text, ImageExtractUsage{}, nil
}

func (e *capturingImageExtractor) request() ImageExtractRequest {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.req
}

func TestNormalizeImageForExtractOversizedPNG(t *testing.T) {
	source := noisyPNG(t, 512, 512)
	if len(source) <= 16<<10 {
		t.Fatalf("source PNG size=%d, want > %d", len(source), 16<<10)
	}

	normalized, contentType, ok := normalizeImageForExtract(context.Background(), source, 16<<10)
	if !ok {
		t.Fatal("expected oversized PNG to normalize")
	}
	if contentType != "image/jpeg" {
		t.Fatalf("content type=%q, want image/jpeg", contentType)
	}
	if len(normalized) > 16<<10 {
		t.Fatalf("normalized size=%d, want <= %d", len(normalized), 16<<10)
	}
	if _, format, err := image.Decode(bytes.NewReader(normalized)); err != nil || format != "jpeg" {
		t.Fatalf("normalized image decode format=%q err=%v", format, err)
	}
}

func TestNormalizeImageForExtractOversizedJPEGAndGIF(t *testing.T) {
	sourceImage := noisyImage(512, 512)
	cases := []struct {
		name   string
		format string
		encode func(*bytes.Buffer) error
	}{
		{
			name:   "jpeg",
			format: "jpeg",
			encode: func(buf *bytes.Buffer) error {
				return jpeg.Encode(buf, sourceImage, &jpeg.Options{Quality: 100})
			},
		},
		{
			name:   "gif",
			format: "gif",
			encode: func(buf *bytes.Buffer) error {
				return gif.Encode(buf, sourceImage, nil)
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var source bytes.Buffer
			if err := tc.encode(&source); err != nil {
				t.Fatal(err)
			}
			if source.Len() <= 8<<10 {
				t.Fatalf("source %s size=%d, want > %d", tc.format, source.Len(), 8<<10)
			}

			normalized, contentType, ok := normalizeImageForExtract(context.Background(), source.Bytes(), 8<<10)
			if !ok {
				t.Fatalf("expected oversized %s to normalize", tc.format)
			}
			if contentType != "image/jpeg" {
				t.Fatalf("content type=%q, want image/jpeg", contentType)
			}
			if len(normalized) > 8<<10 {
				t.Fatalf("normalized size=%d, want <= %d", len(normalized), 8<<10)
			}
			if _, format, err := image.Decode(bytes.NewReader(normalized)); err != nil || format != "jpeg" {
				t.Fatalf("normalized image decode format=%q err=%v", format, err)
			}
		})
	}
}

func TestNormalizeImageForExtractLeavesSmallPayloadUntouched(t *testing.T) {
	source := []byte("small-image-payload")
	normalized, contentType, ok := normalizeImageForExtract(context.Background(), source, 1024)
	if !ok {
		t.Fatal("small payload should not require normalization")
	}
	if contentType != "" {
		t.Fatalf("content type=%q, want empty", contentType)
	}
	if !bytes.Equal(normalized, source) {
		t.Fatal("small payload was changed")
	}
}

func TestNormalizeImageForExtractHonorsCanceledContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, _, ok := normalizeImageForExtract(ctx, noisyPNG(t, 512, 512), 16<<10); ok {
		t.Fatal("expected canceled normalization to stop")
	}
}

func TestImageExtractNormalizeSlotHasSingleCapacity(t *testing.T) {
	select {
	case imageExtractNormalizeSlot <- struct{}{}:
	case <-time.After(time.Second):
		t.Fatal("normalization slot was already held")
	}
	defer func() { <-imageExtractNormalizeSlot }()

	select {
	case imageExtractNormalizeSlot <- struct{}{}:
		t.Fatal("normalization slot accepted a second concurrent holder")
	default:
	}
}

func TestSupportedImageNormalizationContentTypeStripsParameters(t *testing.T) {
	for _, contentType := range []string{
		"image/jpeg; charset=binary",
		"IMAGE/PNG; name=upload",
		" image/gif ; foo=bar",
	} {
		if !supportedImageNormalizationContentType(contentType) {
			t.Fatalf("content type %q should be normalizable", contentType)
		}
	}
	if supportedImageNormalizationContentType("image/webp; charset=binary") {
		t.Fatal("image/webp should remain outside the normalization allowlist")
	}
}

func TestOpaqueImageCompositesTransparencyOntoWhite(t *testing.T) {
	src := image.NewNRGBA(image.Rect(0, 0, 1, 1))
	src.SetNRGBA(0, 0, color.NRGBA{R: 0, G: 0, B: 0, A: 0})
	got := color.NRGBAModel.Convert((opaqueImage{Image: src}).At(0, 0)).(color.NRGBA)
	if got != (color.NRGBA{R: 255, G: 255, B: 255, A: 255}) {
		t.Fatalf("opaque pixel=%+v, want white opaque pixel", got)
	}
}

func TestNormalizeImageForExtractReturnsFalseWhenStillTooLarge(t *testing.T) {
	source := noisyPNG(t, 512, 512)
	if _, _, ok := normalizeImageForExtract(context.Background(), source, 64); ok {
		t.Fatal("expected normalization to fail when provider limit is too small")
	}
}

func TestNormalizeImageForExtractRejectsUndecodableInput(t *testing.T) {
	if _, _, ok := normalizeImageForExtract(context.Background(), bytes.Repeat([]byte("not-an-image"), 128), 32); ok {
		t.Fatal("expected undecodable input to be rejected")
	}
}

func TestNormalizeImageForExtractRejectsPixelBombBeforeDecode(t *testing.T) {
	source := oversizedPixelBombPNG(t)
	if _, _, ok := normalizeImageForExtract(context.Background(), source, 16<<10); ok {
		t.Fatal("expected image above decoded pixel cap to be rejected")
	}
}

func TestProcessImageExtractTaskRejectsPixelBombWithoutProviderCall(t *testing.T) {
	extractor := &capturingImageExtractor{text: "must not run"}
	b := newTestBackendWithOptions(t, Options{
		AsyncImageExtract: AsyncImageExtractOptions{
			Enabled:       true,
			MaxImageBytes: 16 << 10,
			Extractor:     extractor,
		},
	})
	source := oversizedPixelBombPNG(t)
	fileID := insertImageFileForExtractTest(t, b, "/img/pixel-bomb.png", "image/png", source)
	result, err := b.ProcessImageExtractTask(context.Background(), ImageExtractTaskSpec{
		FileID: fileID, Path: "/img/pixel-bomb.png", ContentType: "image/png", Revision: 1,
	})
	if err != nil {
		t.Fatal(err)
	}
	if result != ImageExtractResultTooLarge {
		t.Fatalf("result=%q, want %q", result, ImageExtractResultTooLarge)
	}
	if req := extractor.request(); len(req.Data) != 0 {
		t.Fatalf("provider received %d bytes, want no call", len(req.Data))
	}
}

func TestNormalizeImageForExtractRetriesWideJPEGDimensionAtSmallerScale(t *testing.T) {
	source := oversizedDecodablePNG(t, 65536, 1, 20<<10)
	config, format, err := image.DecodeConfig(bytes.NewReader(source))
	if err != nil || format != "png" {
		t.Fatalf("wide fixture DecodeConfig format=%q err=%v", format, err)
	}
	if config.Width < 65536 || int64(config.Width)*int64(config.Height) > maxImageExtractDecodedPixels {
		t.Fatalf("wide fixture dimensions=%dx%d do not match test boundary", config.Width, config.Height)
	}

	normalized, contentType, ok := normalizeImageForExtract(context.Background(), source, 16<<10)
	if !ok {
		t.Fatal("expected wide image to normalize after reducing JPEG dimensions")
	}
	if contentType != "image/jpeg" || len(normalized) > 16<<10 {
		t.Fatalf("normalized content type/size=%q/%d, want image/jpeg and <= %d", contentType, len(normalized), 16<<10)
	}
}

func TestLoadImageBytesForExtractUsesRuntimeSourceLimit(t *testing.T) {
	b := &Dat9Backend{}
	want := []byte("source bytes are read from disk-backed storage")
	data, err := b.loadImageBytesForExtract(context.Background(), &datastore.File{
		StorageType: datastore.StorageDB9,
		SizeBytes:   64<<20 + 1,
		ContentBlob: want,
	}, 16<<10, 1<<30, true)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(data, want) {
		t.Fatalf("data=%q, want %q", data, want)
	}
}

func TestProcessImageExtractTaskNormalizesPayloadWithoutChangingFile(t *testing.T) {
	extractor := &capturingImageExtractor{text: "normalized image"}
	b := newTestBackendWithOptions(t, Options{
		AsyncImageExtract: AsyncImageExtractOptions{
			Enabled:       true,
			MaxImageBytes: 16 << 10,
			Extractor:     extractor,
		},
	})

	source := noisyPNG(t, 512, 512)
	fileID := insertImageFileForExtractTest(t, b, "/img/oversized.png", "image/png", source)
	before, err := b.Store().GetFile(context.Background(), fileID)
	if err != nil {
		t.Fatal(err)
	}
	if before.ChecksumSHA256 == "" {
		t.Fatal("test source must have a checksum")
	}
	result, err := b.ProcessImageExtractTask(context.Background(), ImageExtractTaskSpec{
		FileID:      fileID,
		Path:        "/img/oversized.png",
		ContentType: "image/png",
		Revision:    1,
	})
	if err != nil {
		t.Fatal(err)
	}
	if result != ImageExtractResultWritten {
		t.Fatalf("result=%q, want %q", result, ImageExtractResultWritten)
	}

	req := extractor.request()
	if req.ContentType != "image/jpeg" {
		t.Fatalf("provider content type=%q, want image/jpeg", req.ContentType)
	}
	if len(req.Data) == 0 || len(req.Data) > 16<<10 {
		t.Fatalf("provider payload size=%d, want 1..%d", len(req.Data), 16<<10)
	}

	stored, err := b.Store().GetFile(context.Background(), fileID)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(stored.ContentBlob, source) {
		t.Fatal("image extraction changed the stored source bytes")
	}
	if stored.SizeBytes != before.SizeBytes || stored.ChecksumSHA256 != before.ChecksumSHA256 {
		t.Fatalf("image extraction changed source metadata: before size/checksum=%d/%s after=%d/%s", before.SizeBytes, before.ChecksumSHA256, stored.SizeBytes, stored.ChecksumSHA256)
	}
}

func TestProcessImageExtractTaskWithoutWorkspaceDoesNotLoadOversizedSource(t *testing.T) {
	extractor := &capturingImageExtractor{text: "must not run"}
	b := newTestBackendWithS3AndOptions(t, Options{DatabaseAutoEmbedding: true})
	source := noisyPNG(t, 512, 512)
	const providerLimit = 16 << 10
	if len(source) <= providerLimit {
		t.Fatalf("fixture bytes=%d, want > %d", len(source), providerLimit)
	}
	fileID := insertS3ImageFileForExtractTest(t, b, "/img/no-workspace.png", "image/png", source)
	result, err := b.ProcessImageExtractTaskWithRuntime(context.Background(), ImageExtractTaskSpec{
		FileID: fileID, Path: "/img/no-workspace.png", ContentType: "image/png", Revision: 1,
	}, ImageExtractRuntime{
		Extractor: extractor, MaxSourceBytes: 1 << 30, MaxImageBytes: providerLimit,
		MaxExtractTextBytes: DefaultImageExtractMaxTextBytes, RequireRasterImageBytes: true,
	})
	if result != ImageExtractResultLoadError || !errors.Is(err, ErrMediaExtractWorkspaceRequired) || len(extractor.request().Data) != 0 {
		t.Fatalf("result=%q err=%v provider bytes=%d, want retryable workspace-required without provider call", result, err, len(extractor.request().Data))
	}
}

func TestProcessImageExtractTaskWithoutWorkspaceKeepsUnsupportedFormatTerminal(t *testing.T) {
	b := newTestBackendWithS3AndOptions(t, Options{DatabaseAutoEmbedding: true})
	source := bytes.Repeat([]byte("tiff"), 5<<10)
	fileID := insertS3ImageFileForExtractTest(t, b, "/img/no-workspace.tiff", "image/tiff", source)
	result, err := b.ProcessImageExtractTaskWithRuntime(context.Background(), ImageExtractTaskSpec{
		FileID: fileID, Path: "/img/no-workspace.tiff", ContentType: "image/tiff", Revision: 1,
	}, ImageExtractRuntime{
		Extractor: &capturingImageExtractor{text: "must not run"}, MaxSourceBytes: 1 << 30, MaxImageBytes: 16 << 10,
		MaxExtractTextBytes: DefaultImageExtractMaxTextBytes, RequireRasterImageBytes: true,
	})
	if err != nil || result != ImageExtractResultTooLarge {
		t.Fatalf("result=%q err=%v, want terminal too_large for unsupported normalization format", result, err)
	}
}

func TestProcessImageExtractTaskSmallS3ImageBypassesSaturatedWorkspace(t *testing.T) {
	b := newTestBackendWithS3(t)
	workspace, err := NewMediaExtractWorkspace(MediaExtractWorkspaceOptions{
		RootDir:            t.TempDir(),
		MaxConcurrentTasks: 1,
	})
	if err != nil {
		t.Fatal(err)
	}
	held, err := workspace.NewTask(context.Background(), "held")
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = held.Close() }()

	fileID := insertS3ImageFileForExtractTest(t, b, "/img/small-s3.png", "image/png", fakePNG)
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()
	result, err := b.ProcessImageExtractTaskWithRuntime(ctx, ImageExtractTaskSpec{
		FileID: fileID, Path: "/img/small-s3.png", ContentType: "image/png", Revision: 1,
	}, ImageExtractRuntime{
		Extractor:               &staticImageExtractor{text: "small S3 image"},
		Timeout:                 time.Second,
		MaxImageBytes:           1 << 20,
		MaxExtractTextBytes:     1 << 10,
		RequireRasterImageBytes: true,
		Workspace:               workspace,
	})
	if err != nil {
		t.Fatal(err)
	}
	if result != ImageExtractResultWritten {
		t.Fatalf("result=%q, want %q", result, ImageExtractResultWritten)
	}
}

func TestProcessImageExtractTaskReleasesS3WorkspaceBeforeProviderCall(t *testing.T) {
	b := newTestBackendWithS3(t)
	workspace, err := NewMediaExtractWorkspace(MediaExtractWorkspaceOptions{
		RootDir:            t.TempDir(),
		MaxConcurrentTasks: 1,
	})
	if err != nil {
		t.Fatal(err)
	}
	source := noisyPNG(t, 512, 512)
	const maxImageBytes = 16 << 10
	if len(source) <= maxImageBytes {
		t.Fatalf("source size=%d, want > %d", len(source), maxImageBytes)
	}
	fileID := insertS3ImageFileForExtractTest(t, b, "/img/large-s3.png", "image/png", source)
	extractor := &gatedImageExtractor{
		started: make(chan struct{}, 1),
		release: make(chan struct{}),
	}
	var releaseProvider sync.Once
	defer releaseProvider.Do(func() { close(extractor.release) })

	resultCh := make(chan ImageExtractResult, 1)
	errCh := make(chan error, 1)
	go func() {
		result, extractErr := b.ProcessImageExtractTaskWithRuntime(context.Background(), ImageExtractTaskSpec{
			FileID: fileID, Path: "/img/large-s3.png", ContentType: "image/png", Revision: 1,
		}, ImageExtractRuntime{
			Extractor:               extractor,
			Timeout:                 5 * time.Second,
			MaxImageBytes:           maxImageBytes,
			MaxExtractTextBytes:     1 << 10,
			RequireRasterImageBytes: true,
			Workspace:               workspace,
		})
		resultCh <- result
		errCh <- extractErr
	}()

	select {
	case <-extractor.started:
	case <-time.After(5 * time.Second):
		t.Fatal("provider call did not start")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	next, err := workspace.NewTask(ctx, "next")
	cancel()
	if err != nil {
		t.Fatalf("workspace remained held during provider call: %v", err)
	}
	if err := next.Close(); err != nil {
		t.Fatal(err)
	}
	releaseProvider.Do(func() { close(extractor.release) })
	if err := <-errCh; err != nil {
		t.Fatal(err)
	}
	if result := <-resultCh; result != ImageExtractResultWritten {
		t.Fatalf("result=%q, want %q", result, ImageExtractResultWritten)
	}
}

func noisyPNG(t *testing.T, width, height int) []byte {
	t.Helper()
	img := noisyImage(width, height)
	var encoded bytes.Buffer
	if err := png.Encode(&encoded, img); err != nil {
		t.Fatal(err)
	}
	return encoded.Bytes()
}

func noisyImage(width, height int) *image.RGBA {
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			img.SetRGBA(x, y, color.RGBA{
				R: uint8((x*37 + y*17) % 256),
				G: uint8((x*13 + y*43) % 256),
				B: uint8((x*71 + y*29) % 256),
				A: 255,
			})
		}
	}
	return img
}

func oversizedDecodablePNG(t *testing.T, width, height, paddingBytes int) []byte {
	t.Helper()
	img := image.NewGray(image.Rect(0, 0, width, height))
	var encoded bytes.Buffer
	if err := png.Encode(&encoded, img); err != nil {
		t.Fatal(err)
	}

	// Keep the fixture valid while making its encoded source larger than the
	// provider limit, so removing the pixel guard changes the test outcome.
	data := encoded.Bytes()
	i := bytes.LastIndex(data, []byte{0, 0, 0, 0, 'I', 'E', 'N', 'D'})
	if i < 0 {
		t.Fatal("encoded PNG has no IEND chunk")
	}
	text := make([]byte, paddingBytes)
	copy(text, []byte("padding\x00"))
	chunk := make([]byte, 4+4+len(text)+4)
	binary.BigEndian.PutUint32(chunk[0:4], uint32(len(text)))
	copy(chunk[4:8], []byte("tEXt"))
	copy(chunk[8:8+len(text)], text)
	binary.BigEndian.PutUint32(chunk[8+len(text):], crc32.ChecksumIEEE(chunk[4:8+len(text)]))

	result := make([]byte, 0, len(data)+len(chunk))
	result = append(result, data[:i]...)
	result = append(result, chunk...)
	result = append(result, data[i:]...)
	if _, format, err := image.Decode(bytes.NewReader(result)); err != nil || format != "png" {
		t.Fatalf("oversized PNG fixture decode format=%q err=%v", format, err)
	}
	return result
}

func oversizedPixelBombPNG(t *testing.T) []byte {
	t.Helper()
	source := oversizedDecodablePNG(t, 4100, 4100, 20<<10)
	config, format, err := image.DecodeConfig(bytes.NewReader(source))
	if err != nil || format != "png" {
		t.Fatalf("pixel-bomb fixture DecodeConfig format=%q err=%v", format, err)
	}
	if int64(config.Width)*int64(config.Height) <= maxImageExtractDecodedPixels {
		t.Fatalf("pixel-bomb fixture dimensions=%dx%d do not exceed pixel cap", config.Width, config.Height)
	}
	if len(source) <= 16<<10 {
		t.Fatalf("pixel-bomb fixture size=%d is not oversized", len(source))
	}
	return source
}