package oklab

import (
	"encoding/binary"
	"image"
	"image/color"
	"image/draw"
	"math"
)

// Not many images can say they have 256-bit color;
// I've definitely gone overboard with this.
const bytesPerPixel = 8 * 4

// Image represents an image that encodes its pixels using OKLab color.
//
// Each pixel is 32 bytes in size, encoding the Lightness, ChromaA, ChromaB, and Alpha components
// of each pixel in that order as big-endian float64 values.
type Image struct {
	Pix    []uint8
	Stride int
	Rect   image.Rectangle
}

// ColorModel implements the [image.Image] interface.
func (p *Image) ColorModel() color.Model {
	return Model
}

// Bounds implements the [image.Image] interface.
func (p *Image) Bounds() image.Rectangle {
	return p.Rect
}

// At implements the [image.Image] interface.
func (p *Image) At(x, y int) color.Color {
	return p.OkLabAt(x, y)
}

// RGBA64At implements the [image.RGBA64Image] interface.
func (p *Image) RGBA64At(x, y int) color.RGBA64 {
	r, g, b, a := p.OkLabAt(x, y).RGBA()
	return color.RGBA64{uint16(r), uint16(g), uint16(b), uint16(a)}
}

// OkLabAt decodes the pixel at the given coordinates into an OkLab color.
//
// The zero value is returned for coordinates outside of the image.
func (p *Image) OkLabAt(x, y int) Color {
	if !image.Pt(x, y).In(p.Rect) {
		return Color{}
	}

	i := p.PixOffset(x, y)
	in := p.Pix[i : i+bytesPerPixel]

	return Color{
		Lightness: math.Float64frombits(binary.BigEndian.Uint64(in[0:8])),
		ChromaA:   math.Float64frombits(binary.BigEndian.Uint64(in[8:16])),
		ChromaB:   math.Float64frombits(binary.BigEndian.Uint64(in[16:24])),
		A:         math.Float64frombits(binary.BigEndian.Uint64(in[24:32])),
	}
}

// PixOffset returns the offset into the [Pix] slice a pixel begins.
//
// The given coordinates must be inside the image.
func (p *Image) PixOffset(x, y int) int {
	return (y-p.Rect.Min.Y)*p.Stride + (x-p.Rect.Min.X)*bytesPerPixel
}

// Set implements the [draw.Image] interface.
func (p *Image) Set(x, y int, c color.Color) {
	if !image.Pt(x, y).In(p.Rect) {
		return
	}

	p.setOkLab(x, y, FromColor(c))
}

// Set implements the [draw.RGBA64Image] interface.
func (p *Image) SetRGBA64(x, y int, c color.RGBA64) {
	if !image.Pt(x, y).In(p.Rect) {
		return
	}

	p.setOkLab(x, y, FromRGBA(uint32(c.R), uint32(c.G), uint32(c.B), uint32(c.A)))
}

func (p *Image) setOkLab(x, y int, c Color) {
	i := p.PixOffset(x, y)
	out := p.Pix[i : i+bytesPerPixel]

	binary.BigEndian.PutUint64(out[0:8], math.Float64bits(c.Lightness))
	binary.BigEndian.PutUint64(out[8:16], math.Float64bits(c.ChromaA))
	binary.BigEndian.PutUint64(out[16:24], math.Float64bits(c.ChromaB))
	binary.BigEndian.PutUint64(out[24:32], math.Float64bits(c.A))
}

// SetOkLab encodes the given color into the pixel at the given coordinates.
//
// Does nothing if the coordinates are outside the image.
func (p *Image) SetOkLab(x, y int, c Color) {
	if !image.Pt(x, y).In(p.Rect) {
		return
	}

	p.setOkLab(x, y, c)
}

// SubImage returns a new [Image] that shares the same underlying pixel data.
func (p *Image) SubImage(r image.Rectangle) *Image {
	r = r.Intersect(p.Rect)

	if r.Empty() {
		return &Image{}
	}

	start := p.PixOffset(r.Min.X, r.Min.Y)
	end := p.PixOffset(r.Max.X-1, r.Max.Y-1) + bytesPerPixel

	return &Image{
		Pix:    p.Pix[start:end:end],
		Stride: p.Stride,
		Rect:   r,
	}
}

// Opaque returns true if every pixel in the image would have an alpha value of 0xffff when
// converted to [color.RGBA64].
func (p *Image) Opaque() bool {
	w, h := p.Rect.Dx(), p.Rect.Dy()
	for y := 0; y < h; y++ {
		row := p.Pix[y*p.Stride : y*p.Stride+w*bytesPerPixel]
		for i, stop := 0, w*bytesPerPixel; i < stop; i += bytesPerPixel {
			// we consider a pixel opaque if the NRGBA/RGBA methods
			// would have returned 0xffff.
			if uint32(math.Float64frombits(binary.BigEndian.Uint64(row[i+24:i+32]))*0xffff+0.5) != 0xffff {
				return false
			}
		}
	}

	return true
}

// NewImage returns a new [Image] of the specified size.
func NewImage(r image.Rectangle) *Image {
	return &Image{
		Pix:    make([]uint8, r.Dx()*r.Dy()*bytesPerPixel),
		Stride: r.Dx() * bytesPerPixel,
		Rect:   r,
	}
}

var (
	// type assertions
	_ interface {
		image.Image
		draw.Image
		image.RGBA64Image
		draw.RGBA64Image
	} = (*Image)(nil)
)