// This is Free Software under GNU Affero General Public License v >= 3.0
// without warranty, see README.md and license for details.
//
// SPDX-License-Identifier: AGPL-3.0-or-later
// License-Filename: LICENSES/AGPL-3.0.txt
//
// Copyright (C) 2024 by via donau
//   – Österreichische Wasserstraßen-Gesellschaft mbH
// Software engineering by Intevation GmbH
//
// Author(s):
//  * Sascha L. Teichmann <sascha.teichmann@intevation.de>

package mesh

import (
	"bufio"
	"cmp"
	"encoding/binary"
	"errors"
	"io"
	"math"
	"slices"
	"sort"

	"gemma.intevation.de/gemma/pkg/common"
)

const quantScale = 1_000

func quant(x float64) int64 { return int64(math.Round(x * quantScale)) }

func (s *STRTree) serializeV2(w io.Writer) error {
	// Write header
	if _, err := w.Write([]byte(magicHeader)); err != nil {
		return err
	}
	// Write version
	if err := binary.Write(w, binary.LittleEndian, uint32(2)); err != nil {
		return err
	}
	// Write the mesh itself
	if err := s.tin.serializeV2(w); err != nil {
		return err
	}
	if err := binary.Write(w, binary.LittleEndian, uint8(s.Entries)); err != nil {
		return err
	}
	if err := s.serializeIndexV2(w); err != nil {
		return err
	}
	return s.serializeBBoxesV2(w)
}

func (t *Tin) serializeV2(w io.Writer) error {
	if err := binary.Write(w, binary.LittleEndian, t.EPSG); err != nil {
		return err
	}
	if err := t.Min.Write(w); err != nil {
		return err
	}
	if err := t.Max.Write(w); err != nil {
		return err
	}
	if err := t.serializeVerticesV2(w); err != nil {
		return err
	}
	return t.serializeTrianglesV2(w)
}

func zigZag(buf []byte, w io.Writer, x int64) error {
	for b := buf[:binary.PutVarint(buf, x)]; len(b) > 0; {
		m, err := w.Write(buf)
		if err != nil {
			return err
		}
		b = b[m:]
	}
	return nil
}

func (t *Tin) serializeVerticesV2(w io.Writer) error {
	vertices := t.Vertices
	if err := binary.Write(w, binary.LittleEndian, uint32(len(vertices))); err != nil {
		return err
	}
	m := t.Min
	buf := make([]byte, binary.MaxVarintLen64)
	for dim := range 3 {
		delta := common.Delta()
		for i := range vertices {
			var v float64
			switch dim {
			case 0:
				v = vertices[i].X - m.X
			case 1:
				v = vertices[i].Y - m.Y
			case 2:
				v = vertices[i].Z - m.Z
			}
			q := quant(v)
			d := delta(q)
			if err := zigZag(buf, w, d); err != nil {
				return err
			}
		}
	}
	return nil
}

func (t *Tin) serializeTrianglesV2(w io.Writer) error {
	triangles := t.Triangles
	if err := binary.Write(w, binary.LittleEndian, uint32(len(triangles))); err != nil {
		return err
	}
	buf := make([]byte, binary.MaxVarintLen64)
	delta := common.Delta()
	for _, tri := range triangles {
		for _, idx := range tri {
			d := delta(int64(idx))
			if err := zigZag(buf, w, d); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *STRTree) serializeIndexV2(w io.Writer) error {
	index := s.index
	if err := binary.Write(w, binary.LittleEndian, uint32(len(index))); err != nil {
		return err
	}
	buf := make([]byte, binary.MaxVarintLen64)
	delta := common.Delta()
	for _, idx := range index {
		d := delta(int64(idx))
		if err := zigZag(buf, w, d); err != nil {
			return err
		}
	}
	return nil
}

func (s *STRTree) serializeBBoxesV2(w io.Writer) error {
	boxes := s.bboxes
	if err := binary.Write(w, binary.LittleEndian, uint32(len(boxes))); err != nil {
		return err
	}
	var (
		m      = s.Min()
		x1     = func(b Box2D) int64 { return quant(b.X1 - m.X) }
		y1     = func(b Box2D) int64 { return quant(b.Y1 - m.Y) }
		width  = func(b Box2D) int64 { return quant(b.X2 - b.X1) }
		height = func(b Box2D) int64 { return quant(b.Y2 - b.Y1) }
		buf    = make([]byte, binary.MaxVarintLen64)
	)
	for _, proj := range []func(Box2D) int64{x1, y1, width, height} {
		delta := common.Delta()
		for _, b := range boxes {
			v := delta(proj(b))
			if err := zigZag(buf, w, v); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *STRTree) deserializeV2(r *bufio.Reader) error {
	// TODO: implement me!
	return errors.New("not implemented, yet")
}

// optimizeForSerializationV2 prepares the mesh to be stored in V2.
func (s *STRTree) optimizeForSerializationV2() {
	s.removeUnusedTriangles()
	s.sortVertices()
	s.sortIndices()
}

// removeUnusedTriangles removes all triangles from the
// TIN that are not registered in the spatial index.
func (s *STRTree) removeUnusedTriangles() {
	used := make([]bool, len(s.tin.Triangles))
	unused := len(used)
	s.allTriangleIndices(func(indices []int32) {
		for _, idx := range indices {
			used[idx] = true
			unused-- // They only occur once in the spatial index.
		}
	})
	if unused <= 0 {
		return
	}
	newTris := make([][]int32, 0, len(s.tin.Triangles)-unused)
	remap := map[int32]int32{}
	for idx, tri := range s.tin.Triangles {
		if used[idx] {
			remap[int32(idx)] = int32(len(newTris))
			newTris = append(newTris, tri)
		}
	}
	s.tin.Triangles = newTris
	// Update the spatial index as the gaps are now closed.
	s.allTriangleIndices(func(indices []int32) {
		for i, idx := range indices {
			indices[i] = remap[idx]
		}
	})
}

// sortVertices so that the deltas between neighbors are minimized.
func (s *STRTree) sortVertices() {
	vertices := s.tin.Vertices
	indices := make([]int32, len(vertices))
	for i := range len(indices) {
		indices[i] = int32(i)
	}

	// Take aspect ratio of bbox into account.
	width := s.tin.Max.X - s.tin.Min.X
	height := s.tin.Max.Y - s.tin.Min.Y
	if width == 0 || height == 0 {
		return
	}

	// v = len(vertices)
	// r = max(w, h)/min(w, h)
	// n * t = v <=> n = v/t
	// n / t = r <=> n = r*t
	// v/t = r*t
	// v = r*t^2
	// v/r = t^2
	// t = sqrt(v/r)
	r := float64(max(width, height)) / float64(min(width, height))
	t := int(math.Ceil(math.Sqrt(float64(len(vertices)) / r)))

	var d1, d2 func(int32, int32) int

	if width > height {
		// Sort in X first and slices alternating up and down
		slices.SortFunc(indices, func(a, b int32) int {
			return cmp.Compare(vertices[a].X, vertices[b].X)
		})
		d1 = func(a, b int32) int { return cmp.Compare(vertices[a].Y, vertices[b].Y) }
		d2 = func(a, b int32) int { return cmp.Compare(vertices[b].Y, vertices[a].Y) }
	} else {
		// Sort in Y first and slices alternating left and right
		slices.SortFunc(indices, func(a, b int32) int {
			return cmp.Compare(vertices[a].Y, vertices[b].Y)
		})
		d1 = func(a, b int32) int { return cmp.Compare(vertices[a].X, vertices[b].X) }
		d2 = func(a, b int32) int { return cmp.Compare(vertices[b].X, vertices[a].X) }
	}

	for p := indices; len(p) > 0; d1, d2 = d2, d1 {
		l := min(len(p), t)
		slices.SortStableFunc(p[:l], d2)
		p = p[l:]
	}

	// We really need to to sort the vertices actually, too.
	sort.Slice(vertices, func(i, j int) bool { return indices[i] < indices[j] })

	reindices := make([]int32, len(indices))
	for i, idx := range indices {
		reindices[idx] = int32(i)
	}

	// Update the indices in the triangles.
	for _, tri := range s.tin.Triangles {
		for i, idx := range tri {
			tri[i] = reindices[idx]
		}
	}
}