package utils

import (
	"bytes"
	"encoding/base64"
	"image"
	"image/color"
	"image/png"
	"io"

	"github.com/charmbracelet/x/ansi"
	"github.com/charmbracelet/x/ansi/iterm2"
	"github.com/charmbracelet/x/ansi/kitty"
	"github.com/charmbracelet/x/ansi/sixel"
)

// ImageProtocol selects the terminal graphics protocol.
type ImageProtocol int

const (
	ProtocolKitty ImageProtocol = iota
	ProtocolSixel
	ProtocolITerm
)

// SpectrogramDisplaySize is the default pixel dimension for spectrogram images.
// 448px suits Retina/HiDPI screens (224 logical pixels at 2x).
const SpectrogramDisplaySize = 448

// ClampImageSize clamps a dimension to [224, 448].
func ClampImageSize(size int) int {
	return max(224, min(896, size))
}

// WriteImage writes an image using the specified terminal graphics protocol.
func WriteImage(img image.Image, w io.Writer, protocol ImageProtocol) error {
	switch protocol {
	case ProtocolSixel:
		return WriteSixelImage(img, w)
	case ProtocolITerm:
		return WriteITermImage(img, w)
	default:
		return WriteKittyImage(img, w)
	}
}

// ClearImages clears previously displayed images.
// For kitty, deletes all image placements. For sixel/iTerm2, no-op (inline text).
func ClearImages(w io.Writer, protocol ImageProtocol) error {
	switch protocol {
	case ProtocolKitty:
		return ClearKittyImages(w)
	default:
		return nil
	}
}

// ClearKittyImages clears all previously displayed Kitty images
func ClearKittyImages(w io.Writer) error {
	_, err := io.WriteString(w, ansi.KittyGraphics(nil, "a=d"))
	return err
}

// WriteKittyImage writes an image to the writer using the Kitty graphics protocol.
// The image is encoded as PNG, base64'd, and sent via chunked Kitty escape sequences.
func WriteKittyImage(img image.Image, w io.Writer) error {
	return kitty.EncodeGraphics(w, img, &kitty.Options{
		Format:       kitty.PNG,
		Action:       kitty.TransmitAndPut,
		Transmission: kitty.Direct,
		Chunk:        true,
	})
}

// WriteSixelImage writes an image using the Sixel graphics protocol.
func WriteSixelImage(img image.Image, w io.Writer) error {
	var buf bytes.Buffer
	enc := &sixel.Encoder{}
	if err := enc.Encode(&buf, img); err != nil {
		return err
	}
	_, err := io.WriteString(w, ansi.SixelGraphics(0, 1, 0, buf.Bytes()))
	return err
}

// WriteITermImage writes an image using the iTerm2 Inline Image Protocol.
func WriteITermImage(img image.Image, w io.Writer) error {
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return err
	}
	b64 := base64.StdEncoding.EncodeToString(buf.Bytes())
	_, err := io.WriteString(w, ansi.ITerm2(iterm2.File{
		Inline:  true,
		Content: []byte(b64),
	}))
	return err
}

// CreateGrayscaleImage creates an image.Image from a 2D uint8 array.
// The array is organized as [rows][cols] where rows = frequency bins.
func CreateGrayscaleImage(data [][]uint8) image.Image {
	if len(data) == 0 || len(data[0]) == 0 {
		return nil
	}

	height := len(data)
	width := len(data[0])

	img := image.NewGray(image.Rect(0, 0, width, height))

	for y := range height {
		off := y * img.Stride
		row := data[y]
		copy(img.Pix[off:off+width], row)
	}

	return img
}

// CreateRGBImage creates an image.Image from a 2D RGBPixel array.
// The array is organized as [rows][cols] where rows = frequency bins.
func CreateRGBImage(data [][]RGBPixel) image.Image {
	if len(data) == 0 || len(data[0]) == 0 {
		return nil
	}

	height := len(data)
	width := len(data[0])

	img := image.NewRGBA(image.Rect(0, 0, width, height))

	for y := range height {
		off := y * img.Stride
		row := data[y]
		for x := range width {
			i := off + x*4
			img.Pix[i] = row[x].R
			img.Pix[i+1] = row[x].G
			img.Pix[i+2] = row[x].B
			img.Pix[i+3] = 255
		}
	}

	return img
}

// resizeScale holds precomputed scale factors for nearest-neighbor resizing.
type resizeScale struct {
	srcWidth, srcHeight int
	scaleX, scaleY      float64
}

func newResizeScale(img image.Image, newWidth, newHeight int) resizeScale {
	bounds := img.Bounds()
	return resizeScale{
		srcWidth:  bounds.Dx(),
		srcHeight: bounds.Dy(),
		scaleX:    float64(bounds.Dx()) / float64(newWidth),
		scaleY:    float64(bounds.Dy()) / float64(newHeight),
	}
}

// srcCoord maps a destination pixel coordinate to source coordinate, clamped to bounds.
func (s resizeScale) srcCoord(dstX, dstY int) (srcX, srcY int) {
	srcX = int(float64(dstX) * s.scaleX)
	srcY = int(float64(dstY) * s.scaleY)
	if srcX >= s.srcWidth {
		srcX = s.srcWidth - 1
	}
	if srcY >= s.srcHeight {
		srcY = s.srcHeight - 1
	}
	return
}

// resizeGray resizes a Gray image using nearest-neighbor interpolation.
func resizeGray(src *image.Gray, s resizeScale, newWidth, newHeight int) *image.Gray {
	result := image.NewGray(image.Rect(0, 0, newWidth, newHeight))
	for y := range newHeight {
		dstOff := y * result.Stride
		_, srcY := s.srcCoord(0, y)
		srcRowOff := srcY * src.Stride
		for x := range newWidth {
			srcX, _ := s.srcCoord(x, 0)
			result.Pix[dstOff+x] = src.Pix[srcRowOff+srcX]
		}
	}
	return result
}

// resizeRGBA resizes an RGBA image using nearest-neighbor interpolation.
func resizeRGBA(src *image.RGBA, s resizeScale, newWidth, newHeight int) *image.RGBA {
	result := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
	for y := range newHeight {
		dstOff := y * result.Stride
		_, srcY := s.srcCoord(0, y)
		srcRowOff := srcY * src.Stride
		for x := range newWidth {
			srcX, _ := s.srcCoord(x, 0)
			si := srcRowOff + srcX*4
			di := dstOff + x*4
			result.Pix[di] = src.Pix[si]
			result.Pix[di+1] = src.Pix[si+1]
			result.Pix[di+2] = src.Pix[si+2]
			result.Pix[di+3] = src.Pix[si+3]
		}
	}
	return result
}

// resizeGeneric resizes any image using nearest-neighbor interpolation (slow fallback).
func resizeGeneric(img image.Image, s resizeScale, newWidth, newHeight int) *image.RGBA {
	bounds := img.Bounds()
	result := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
	for y := range newHeight {
		for x := range newWidth {
			srcX, srcY := s.srcCoord(x, y)
			c := img.At(srcX+bounds.Min.X, srcY+bounds.Min.Y)
			r, g, b, _ := c.RGBA()
			result.SetRGBA(x, y, color.RGBA{
				R: uint8(r >> 8),
				G: uint8(g >> 8),
				B: uint8(b >> 8),
				A: 255,
			})
		}
	}
	return result
}

// ResizeImage resizes an image using nearest-neighbor interpolation.
// For higher quality, use golang.org/x/image/draw, but this keeps dependencies minimal.
func ResizeImage(img image.Image, newWidth, newHeight int) image.Image {
	s := newResizeScale(img, newWidth, newHeight)

	if srcGray, ok := img.(*image.Gray); ok {
		return resizeGray(srcGray, s, newWidth, newHeight)
	}
	if srcRGBA, ok := img.(*image.RGBA); ok {
		return resizeRGBA(srcRGBA, s, newWidth, newHeight)
	}
	return resizeGeneric(img, s, newWidth, newHeight)
}

// WritePNG writes an image to a writer in PNG format using fast compression.
func WritePNG(img image.Image, w io.Writer) error {
	enc := &png.Encoder{CompressionLevel: png.BestSpeed}
	return enc.Encode(w, img)
}