package field

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"image"
	"image/png"
	"io"
	"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)
}

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
}

// 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 {
	small := h.Resample(size, size)
	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*size)
	for y := 0; y < size; 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, size, px, png.BestSpeed)
}

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