package field

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"image"
	"image/png"
	"io"
	"math"
	"os"
	"path/filepath"
)

// Greyscale PNG in and out, plus the raw 16-bit little-endian .r16 that World Machine, Gaea and the engine's
// own exporter write. The numpy pipeline hand-rolled all of this because the engine's Python has no PIL;
// image/png covers it, so the only thing worth carrying over is the behaviour, not the code.

// WriteGray16 writes values already in 0..65535. Compression matters at this size: one 7141 x 7141 map is
// 102 MB of samples, so the level is a parameter and the caller pays for what it needs. The height map is
// imported by the editor and worth compressing; the derivative maps are rebuilt from a seed and are not.
func WriteGray16(path string, w, h int, values []uint16, level png.CompressionLevel) error {
	if len(values) != w*h {
		return fmt.Errorf("%s: %d values for a %dx%d image", path, len(values), w, h)
	}
	img := image.NewGray16(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		row := img.Pix[y*img.Stride : y*img.Stride+w*2]
		src := values[y*w : y*w+w]
		for x, v := range src {
			binary.BigEndian.PutUint16(row[x*2:], v) // PNG is big-endian; image.Gray16 stores it that way too
		}
	}
	return encode(path, img, level)
}

// WriteGray8 writes values already in 0..255.
func WriteGray8(path string, w, h int, values []uint8, level png.CompressionLevel) error {
	if len(values) != w*h {
		return fmt.Errorf("%s: %d values for a %dx%d image", path, len(values), w, h)
	}
	img := image.NewGray(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		copy(img.Pix[y*img.Stride:y*img.Stride+w], values[y*w:y*w+w])
	}
	return encode(path, img, level)
}

// WriteRGB writes tightly packed 8-bit RGB, three bytes a pixel. It is what a categorical map wants: a class
// raster has no ramp to run through a palette, only a colour per class.
// WriteRGBA writes colour with a separate alpha plane, which is what an overlay sheet is: strokes on a
// transparent background, where blank is decided by alpha rather than by a reserved colour.
//
// Non-premultiplied (NRGBA), deliberately. A mark's colour has to come back out of the file exactly as it
// went in, because the classifier reads exact colours against a tolerance; premultiplying would scale every
// channel by the alpha and round on the way, and a mark would classify as something else or as nothing.
func WriteRGBA(path string, w, h int, px []uint8, alpha []uint8, level png.CompressionLevel) error {
	if len(px) != w*h*3 {
		return fmt.Errorf("%s: %d bytes for a %dx%d RGB image", path, len(px), w, h)
	}
	if len(alpha) != w*h {
		return fmt.Errorf("%s: %d alpha bytes for a %dx%d image", path, len(alpha), w, h)
	}
	img := image.NewNRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		row := img.Pix[y*img.Stride:]
		src := px[y*w*3:]
		a := alpha[y*w:]
		for x := 0; x < w; x++ {
			row[x*4], row[x*4+1], row[x*4+2], row[x*4+3] = src[x*3], src[x*3+1], src[x*3+2], a[x]
		}
	}
	return encode(path, img, level)
}

func WriteRGB(path string, w, h int, px []uint8, level png.CompressionLevel) error {
	if len(px) != w*h*3 {
		return fmt.Errorf("%s: %d bytes for a %dx%d RGB image", path, len(px), w, h)
	}
	img := image.NewNRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		row := img.Pix[y*img.Stride:]
		src := px[y*w*3:]
		for x := 0; x < w; x++ {
			row[x*4], row[x*4+1], row[x*4+2], row[x*4+3] = src[x*3], src[x*3+1], src[x*3+2], 255
		}
	}
	return encode(path, img, level)
}

func encode(path string, img image.Image, level png.CompressionLevel) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	// Written beside the target and renamed, so a killed run never leaves create_world.py a half-written PNG
	// to import. The numpy pipeline learned this the hard way with a killed commandlet.
	tmp := path + ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return err
	}
	bw := bufio.NewWriterSize(f, 1<<20)
	enc := png.Encoder{CompressionLevel: level}
	if err := enc.Encode(bw, img); err != nil {
		f.Close()
		os.Remove(tmp)
		return err
	}
	if err := bw.Flush(); err != nil {
		f.Close()
		os.Remove(tmp)
		return err
	}
	if err := f.Close(); err != nil {
		os.Remove(tmp)
		return err
	}
	return os.Rename(tmp, path)
}

// ReadHeightmap reads a 16-bit greyscale PNG, an 8-bit greyscale PNG (widened, as the numpy pipeline widened
// it) or a raw 16-bit little-endian .r16/.raw, and returns values in 0..65535 with the image's dimensions.
// A raw file needs its width when it is not square.
func ReadHeightmap(path string, width int) (values []uint16, w, h int, err error) {
	switch ext := filepath.Ext(path); ext {
	case ".r16", ".raw":
		return readRaw(path, width)
	default:
		return readPNG(path)
	}
}

func readPNG(path string) ([]uint16, int, int, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, 0, 0, err
	}
	defer f.Close()
	img, err := png.Decode(bufio.NewReaderSize(f, 1<<20))
	if err != nil {
		return nil, 0, 0, fmt.Errorf("%s: %w", path, err)
	}
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	out := make([]uint16, w*h)
	switch src := img.(type) {
	case *image.Gray16:
		for y := 0; y < h; y++ {
			row := src.Pix[y*src.Stride:]
			for x := 0; x < w; x++ {
				out[y*w+x] = binary.BigEndian.Uint16(row[x*2:])
			}
		}
	case *image.Gray:
		// Widened the way the numpy reader widened it: 8-bit 255 must become 65535, not 65280, or a DEM
		// comes in a whisker short of its own ceiling.
		for y := 0; y < h; y++ {
			row := src.Pix[y*src.Stride:]
			for x := 0; x < w; x++ {
				out[y*w+x] = uint16(row[x]) * 257
			}
		}
	default:
		// Anything else (RGB, paletted) goes through the generic accessor, which already returns 16-bit.
		for y := 0; y < h; y++ {
			for x := 0; x < w; x++ {
				r, g, bl, _ := img.At(b.Min.X+x, b.Min.Y+y).RGBA()
				out[y*w+x] = uint16((r*299 + g*587 + bl*114) / 1000)
			}
		}
	}
	return out, w, h, nil
}

func readRaw(path string, width int) ([]uint16, int, int, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, 0, 0, err
	}
	if len(raw)%2 != 0 {
		return nil, 0, 0, fmt.Errorf("%s: %d bytes is not a whole number of 16-bit samples", path, len(raw))
	}
	n := len(raw) / 2
	w := width
	if w <= 0 {
		w = isqrt(n)
		if w*w != n {
			return nil, 0, 0, fmt.Errorf("%s: %d samples is not square; give the width", path, n)
		}
	}
	if n%w != 0 {
		return nil, 0, 0, fmt.Errorf("%s: %d samples do not divide by width %d", path, n, w)
	}
	out := make([]uint16, n)
	for i := 0; i < n; i++ {
		out[i] = binary.LittleEndian.Uint16(raw[i*2:])
	}
	return out, w, n / w, nil
}

func isqrt(n int) int {
	r := 0
	for (r+1)*(r+1) <= n {
		r++
	}
	return r
}

// WriteHillshade writes an 8-bit relief shade of a height field, at any scale.
//
// It is not WriteThumbnail with a bigger number, and the difference is the whole reason it exists.
// WriteThumbnail's shading term is a raw gradient over the cell size, which is fine on a 512-pixel picture of
// a whole map where the gradients are small, and saturates to pure black and white the moment it is used on
// real ground at two metres a cell. The result reads as flat-topped terraces with hard edges - a mountainside
// rendered as a staircase - and it is convincing enough to be mistaken for a defect in the terrain. It was.
//
// This is the standard DEM hillshade instead: the surface normal against a light from the north-west at 45
// degrees, which is bounded by construction and says the same thing at any cell size.
func WriteHillshade(path string, h *Field, size int, exaggeration float64) error {
	sizeH := aspectH(h, size)
	small := h
	if size != h.W || sizeH != h.H {
		small = h.Resample(size, sizeH)
	}
	exag := exaggeration
	if exag <= 0 {
		exag = 1
	}
	px := make([]uint8, size*sizeH)
	Rows(sizeH, func(y0, y1 int) {
		for y := y0; y < y1; y++ {
			for x := 0; x < size; x++ {
				gx := float64(small.AtClamped(x+1, y)-small.AtClamped(x-1, y)) * exag
				gy := float64(small.AtClamped(x, y+1)-small.AtClamped(x, y-1)) * exag
				slope := math.Atan(math.Hypot(gx, gy) / (2 * small.CellM))
				aspect := math.Atan2(gy, -gx)
				lum := math.Cos(slope)*math.Cos(math.Pi/4) +
					math.Sin(slope)*math.Sin(math.Pi/4)*math.Cos(3*math.Pi/4-aspect)
				v := 0.25 + 0.75*math.Max(0, lum)
				if v > 1 {
					v = 1
				}
				px[y*size+x] = uint8(v * 255)
			}
		}
	})
	return WriteGray8(path, size, sizeH, px, png.BestSpeed)
}

// WriteThumbnail writes a small 8-bit preview of a height field, hillshaded so the drainage is actually
// visible: a flat grey ramp hides exactly the thing this generator exists to produce.
func WriteThumbnail(path string, h *Field, size int) error {
	sizeH := aspectH(h, size)
	small := h.Resample(size, sizeH)
	lo, hi := small.MinMax()
	span := float64(hi - lo)
	if span < 1e-6 {
		span = 1
	}
	// Light from the north-west at 45 degrees, the convention every DEM hillshade uses.
	px := make([]uint8, size*sizeH)
	for y := 0; y < sizeH; y++ {
		for x := 0; x < size; x++ {
			gx := float64(small.AtClamped(x+1, y) - small.AtClamped(x-1, y))
			gy := float64(small.AtClamped(x, y+1) - small.AtClamped(x, y-1))
			shade := (0.5*gx + 0.5*gy) / (2 * small.CellM)
			lum := 0.35 + 0.65*(float64(small.At(x, y))-float64(lo))/span
			lum += shade * 0.35
			if lum < 0 {
				lum = 0
			} else if lum > 1 {
				lum = 1
			}
			px[y*size+x] = uint8(lum * 255)
		}
	}
	return WriteGray8(path, size, sizeH, px, png.BestSpeed)
}

var _ io.Writer = (*bufio.Writer)(nil)