// xgif is a helper package to create GIF animations based on PNG images
// The resulting animations have the following properties:
// 1. All frames have the same size (max(pngs.width), max(pngs.height))
// 2. All PNGs are centered in the given frame
// 3. The frame background is plain white
// Note that to convert from a PNG to a GIF compatible image (Bitmap), the PNG image must be quantized (colors are aggregated in median buckets)
// so that it has at most 255 colors.
// This is required because GIFs support only 256 colors and we must keep 1 slot for the white background.
package xgif

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"image/gif"
	"image/png"
	"math"

	"golang.org/x/sync/errgroup"

	"github.com/d2lang/d2/internal/testutil"
	d2png "github.com/d2lang/d2/lib/png"

	"github.com/ericpauley/go-quantize/quantize"
	"github.com/mxschmitt/playwright-go"

	"github.com/d2lang/util-go/go2"
)

const INFINITE_LOOP = 0

var BG_COLOR = color.White

const fps = 30
const workers = 16

func ConvertAnimatedSVGToPNGs(browser playwright.Browser, svg []byte, durationMs int) ([][]byte, error) {
	totalFrames := (durationMs / 1000) * fps
	out := make([][]byte, totalFrames)

	batchSize := int(math.Ceil(float64(totalFrames) / float64(workers)))

	g := new(errgroup.Group)
	g.SetLimit(workers)

	for w := 0; w < workers; w++ {
		w := w
		g.Go(func() error {
			context, err := browser.NewContext(playwright.BrowserNewContextOptions{
				DeviceScaleFactor: playwright.Float(1.0),
			})
			if err != nil {
				return err
			}
			defer context.Close()

			page, err := context.NewPage()
			if err != nil {
				return err
			}
			defer page.Close()

			if err := d2png.MountSVG(page, string(svg)); err != nil {
				return err
			}

			for frameIndex := w * batchSize; frameIndex < min((w+1)*batchSize, totalFrames); frameIndex++ {
				t := float64(frameIndex) / float64(fps)
				if err := d2png.SetAnimationTime(page, t); err != nil {
					return err
				}
				_, err = page.Evaluate(`() => new Promise(r => requestAnimationFrame(() => r()))`)
				if err != nil {
					return err
				}

				pngBuf, err := d2png.ScreenshotSVG(page)
				if err != nil {
					return err
				}

				out[frameIndex] = pngBuf
			}
			return nil
		})
	}

	if err := g.Wait(); err != nil {
		return nil, err
	}

	return out, nil
}

func AnimatePNGs(pngs [][]byte, animIntervalMs int) ([]byte, error) {
	var width, height int
	pngImgs := make([]image.Image, len(pngs))
	for i, pngBytes := range pngs {
		img, err := png.Decode(bytes.NewBuffer(pngBytes))
		if err != nil {
			return nil, err
		}
		pngImgs[i] = img
		bounds := img.Bounds()
		width = go2.Max(width, bounds.Dx())
		height = go2.Max(height, bounds.Dy())
	}

	interval := int(math.Round(100.0 / float64(fps)))
	anim := &gif.GIF{
		LoopCount: INFINITE_LOOP,
		Config: image.Config{
			Width:  width,
			Height: height,
		},
	}

	anim.Image = make([]*image.Paletted, len(pngImgs))
	anim.Delay = make([]int, len(pngImgs))

	g := new(errgroup.Group)
	g.SetLimit(workers)

	for i := 0; i < len(pngImgs); i++ {
		i := i
		g.Go(func() error {
			pngImage := pngImgs[i]
			// 1. convert the PNG into a GIF compatible image (Bitmap) by quantizing it to 255 colors
			buf := bytes.NewBuffer(nil)
			err := gif.Encode(buf, pngImage, &gif.Options{
				NumColors: 256, // GIFs can have up to 256 colors
				Quantizer: quantize.MedianCutQuantizer{},
			})
			if err != nil {
				return err
			}
			gifImg, err := gif.Decode(buf)
			if err != nil {
				return err
			}
			palettedImg, ok := gifImg.(*image.Paletted)
			if !ok {
				return fmt.Errorf("decoded gif image could not be cast as *image.Paletted")
			}

			// 2. make GIF frames of the same size, keeping images centered and with a white background
			bounds := pngImage.Bounds()
			top := (height - bounds.Dy()) / 2
			bottom := top + bounds.Dy()
			left := (width - bounds.Dx()) / 2
			right := left + bounds.Dx()

			var bgIndex int
			if len(palettedImg.Palette) == 256 {
				bgIndex = findWhiteIndex(palettedImg.Palette)
				palettedImg.Palette[bgIndex] = BG_COLOR
			} else {
				bgIndex = len(palettedImg.Palette)
				palettedImg.Palette = append(palettedImg.Palette, BG_COLOR)
			}
			frame := image.NewPaletted(image.Rect(0, 0, width, height), palettedImg.Palette)
			for x := 0; x < width; x++ {
				for y := 0; y < height; y++ {
					if x <= left || y <= top || x >= right || y >= bottom {
						frame.SetColorIndex(x, y, uint8(bgIndex))
					} else {
						frame.SetColorIndex(x, y, palettedImg.ColorIndexAt(x-left, y-top))
					}
				}
			}

			anim.Image[i] = frame
			anim.Delay[i] = interval
			return nil
		})
	}
	if err := g.Wait(); err != nil {
		return nil, err
	}

	buf := bytes.NewBuffer(nil)
	err := gif.EncodeAll(buf, anim)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func findWhiteIndex(palette color.Palette) int {
	nearestIndex := 0
	nearestScore := 0.
	for i, c := range palette {
		r, g, b, _ := c.RGBA()
		if r == 255 && g == 255 && b == 255 {
			return i
		}

		avg := float64(r+g+b) / 255.
		if avg > nearestScore {
			nearestScore = avg
			nearestIndex = i
		}
	}
	return nearestIndex
}

// Validate checks the frame count, loop behavior, dimensions, and interval of
// gifBytes.
//
// Deprecated: Validate is a test-only helper specific to D2's output and will
// be removed after one compatibility release. Downstream tests should decode
// the GIF and validate the properties they rely on directly.
func Validate(gifBytes []byte, nFrames int, intervalMS int) error {
	return testutil.ValidateGIF(gifBytes, nFrames, intervalMS)
}
