// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package io

import (
	"math"
	"os"
	"path/filepath"
	"slices"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Coverage-guided fuzzing over the four binary readers. Every input is
// written to one shared temp file per worker process and handed to the
// reader; a panic, a hang or an allocation blow-up fails the input, and
// so does a returned result whose shape disagrees with its payload,
// because a silent mismatch is the worst outcome a parser can produce.
// Seeds are the real fixtures plus files built by the own writers, so
// the mutator starts from genuinely valid bytes rather than from magic
// constants. Run one target at a time, for example:
//
//	go test -run '^$' -fuzz FuzzLoadHDF5 -fuzztime 2m ./io/
//
// Inputs that expose a defect land in testdata/fuzz/<FuzzName>/ and
// stay there as regression seeds.

// fuzzWrite seeds the corpus with one file built by a writer.
func fuzzWrite(f *testing.F, name string, build func(path string) error) {
	path := filepath.Join(f.TempDir(), name)
	if err := build(path); err != nil {
		f.Fatalf("build seed %s: %v", name, err)
	}
	data, err := os.ReadFile(path)
	if err != nil {
		f.Fatalf("read seed %s: %v", name, err)
	}
	f.Add(data)
}

// fuzzFloats builds a float64 vector of n recognisably distinct values.
func fuzzFloats(f *testing.F, n int) *core.Array {
	f.Helper()
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = float64(i+1) + 0.5
	}
	a, err := core.FromFloats(vals, n)
	if err != nil {
		f.Fatalf("FromFloats(%d): %v", n, err)
	}
	return a
}

// fuzzInts builds an int64 vector of n recognisably distinct values.
func fuzzInts(f *testing.F, n int) *core.Array {
	f.Helper()
	vals := make([]int64, n)
	for i := range vals {
		vals[i] = int64(i) - 2
	}
	a, err := core.FromInts(vals, n)
	if err != nil {
		f.Fatalf("FromInts(%d): %v", n, err)
	}
	return a
}

// fuzzShapeProduct multiplies the shape with the overflow guard the
// parsers themselves are held to: an attacker-controlled shape must
// never wrap around into a plausible small length.
func fuzzShapeProduct(t *testing.T, shape []int) int {
	n := 1
	for _, d := range shape {
		if d < 0 || d > (1<<31-1)/n {
			t.Fatalf("hostile shape %v accepted", shape)
		}
		n *= d
	}
	return n
}

func FuzzLoadHDF5(f *testing.F) {
	fuzzWrite(f, "fixture.h5", func(p string) error {
		data, err := os.ReadFile("testdata/h5/fixture.h5")
		if err != nil {
			return err
		}
		return os.WriteFile(p, data, 0o644)
	})
	fuzzWrite(f, "fletcher.h5", func(p string) error {
		data, err := os.ReadFile("testdata/h5/fletcher.h5")
		if err != nil {
			return err
		}
		return os.WriteFile(p, data, 0o644)
	})
	fuzzWrite(f, "latest.h5", func(p string) error {
		data, err := os.ReadFile("testdata/h5/latest.h5")
		if err != nil {
			return err
		}
		return os.WriteFile(p, data, 0o644)
	})
	path := filepath.Join(f.TempDir(), "fuzz.h5")
	f.Fuzz(func(t *testing.T, data []byte) {
		if err := os.WriteFile(path, data, 0o644); err != nil {
			t.Fatal(err)
		}
		sets, err := LoadHDF5(path)
		if err != nil {
			return
		}
		for _, s := range sets {
			if n := fuzzShapeProduct(t, s.Shape); n != s.Values.Len() {
				t.Fatalf("dataset %s: shape product %d != payload %d", s.Path, n, s.Values.Len())
			}
		}
	})
}

func FuzzLoadNetCDF(f *testing.F) {
	fuzzWrite(f, "classic.nc", func(p string) error {
		dims := []NetCDFDim{{Name: "lat", Length: 3}, {Name: "lon", Length: 4}}
		vars := []NetCDFVar{
			{Name: "temp", Dims: []string{"lat", "lon"}, Values: fuzzFloats(f, 12)},
			{Name: "mask", Dims: []string{"lon"}, Values: fuzzInts(f, 4)},
		}
		return SaveNetCDF(p, dims, vars, map[string]string{"title": "seed"})
	})
	fuzzWrite(f, "record.nc", func(p string) error {
		dims := []NetCDFDim{{Name: "t", Length: 0}, {Name: "x", Length: 2}}
		vars := []NetCDFVar{
			{Name: "series", Dims: []string{"t", "x"}, Values: fuzzFloats(f, 6)},
		}
		return SaveNetCDF(p, dims, vars, nil)
	})
	path := filepath.Join(f.TempDir(), "fuzz.nc")
	f.Fuzz(func(t *testing.T, data []byte) {
		if err := os.WriteFile(path, data, 0o644); err != nil {
			t.Fatal(err)
		}
		dims, vars, _, err := LoadNetCDF(path)
		if err != nil {
			return
		}
		lengths := make(map[string]int, len(dims))
		for _, d := range dims {
			lengths[d.Name] = d.Length
		}
		for _, v := range vars {
			// A rank-0 (scalar) variable comes back as a one-element
			// vector with no dimension names, the documented shape.
			if len(v.Values.Shape()) != len(v.Dims) && !(len(v.Dims) == 0 && v.Values.Len() == 1) {
				t.Fatalf("variable %s: rank %d != %d named dimensions", v.Name, len(v.Values.Shape()), len(v.Dims))
			}
			record := false
			n := 1
			for _, name := range v.Dims {
				length, ok := lengths[name]
				if !ok {
					t.Fatalf("variable %s: unknown dimension %s", v.Name, name)
				}
				if length == 0 {
					record = true
					continue
				}
				n *= length
			}
			if !record && n != v.Values.Len() {
				t.Fatalf("variable %s: dimension product %d != payload %d", v.Name, n, v.Values.Len())
			}
		}
	})
}

func FuzzLoadFITS(f *testing.F) {
	fuzzWrite(f, "image.fits", func(p string) error {
		a, err := core.FromFloats([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, 4, 4)
		if err != nil {
			return err
		}
		return SaveFITS(p, a, map[string]string{"OBJECT": "seed", "TELESCOP": "TENSOR"})
	})
	path := filepath.Join(f.TempDir(), "fuzz.fits")
	f.Fuzz(func(t *testing.T, data []byte) {
		if err := os.WriteFile(path, data, 0o644); err != nil {
			t.Fatal(err)
		}
		a, _, err := LoadFITS(path)
		if err != nil {
			return
		}
		if n := fuzzShapeProduct(t, a.Shape()); n != a.Len() {
			t.Fatalf("image: shape product %d != payload %d", n, a.Len())
		}
	})
}

func FuzzLoadFITSTable(f *testing.F) {
	fuzzWrite(f, "binary.fits", func(p string) error {
		cols := []FITSTableColumn{
			{Name: "flux", Form: "D", Data: fuzzFloats(f, 5)},
			{Name: "id", Form: "K", Data: fuzzInts(f, 5)},
			{Name: "name", Form: "8A", Text: []string{"alpha", "beta", "gamma", "delta", "epsilon"}},
		}
		return SaveFITSTable(p, false, cols, nil)
	})
	fuzzWrite(f, "ascii.fits", func(p string) error {
		cols := []FITSTableColumn{
			{Name: "flux", Form: "E12.5", Data: fuzzFloats(f, 3)},
			{Name: "name", Form: "6A", Text: []string{"one", "two", "three"}},
		}
		return SaveFITSTable(p, true, cols, nil)
	})
	path := filepath.Join(f.TempDir(), "fuzztable.fits")
	f.Fuzz(func(t *testing.T, data []byte) {
		if err := os.WriteFile(path, data, 0o644); err != nil {
			t.Fatal(err)
		}
		table, err := LoadFITSTable(path)
		if err != nil {
			return
		}
		if table.Rows < 0 {
			t.Fatalf("negative row count %d", table.Rows)
		}
		n := len(table.Names)
		if len(table.Columns) != n || len(table.Text) != n {
			t.Fatalf("column lists disagree: %d names, %d value columns, %d text columns", n, len(table.Columns), len(table.Text))
		}
		for i := range table.Names {
			switch {
			case table.Columns[i] != nil && table.Text[i] != nil:
				t.Fatalf("column %d carries both values and text", i)
			case table.Columns[i] != nil:
				if table.Columns[i].Len() != table.Rows {
					t.Fatalf("column %d: length %d != %d rows", i, table.Columns[i].Len(), table.Rows)
				}
			case table.Text[i] != nil:
				if len(table.Text[i]) != table.Rows {
					t.Fatalf("text column %d: length %d != %d rows", i, len(table.Text[i]), table.Rows)
				}
			default:
				t.Fatalf("column %d carries neither values nor text", i)
			}
		}
	})
}

// FuzzHDF5WriteRead fuzzes the writer's round-trip contract: whatever
// parameters and values the mutator picks, SaveHDF5 must produce a
// file the reader decodes back to the very same shapes, bit patterns
// and attributes. A write that fails, a read that fails or a value
// that moves is a bug, not a skipped input: unlike the readers there
// is no untrusted bytes here, the writer owns every byte it emits.
func FuzzHDF5WriteRead(f *testing.F) {
	for _, seed := range [][]byte{
		{0, 0, 7, 3, 0, 0, 0, 0},
		{0, 0, 7, 3, 0, 0, 1, 0},
		{1, 1, 19, 5, 1, 1, 0, 1},
		{2, 0, 33, 1, 0, 1, 1, 1},
		{0, 1, 5, 7, 1, 0, 0, 2},
		{2, 1, 11, 4, 1, 1, 1, 1},
		{3, 0, 5, 4, 0, 0, 0, 0},
		{4, 1, 9, 3, 0, 1, 0, 1},
		{5, 0, 6, 2, 0, 0, 1, 0},
		{6, 1, 12, 5, 0, 0, 0, 1},
		{7, 0, 8, 3, 1, 1, 0, 0},
		{8, 1, 15, 4, 0, 0, 0, 1},
		{9, 0, 10, 2, 1, 0, 0, 0},
	} {
		f.Add(seed)
	}
	f.Fuzz(func(t *testing.T, data []byte) {
		if len(data) < 8 {
			t.Skip()
		}
		dtype := int(data[0] % 10)
		rank := 1 + int(data[1]%2)
		d0 := 1 + int(data[2])%31
		d1 := 1 + int(data[3])%9
		gzip := 0
		if data[4]%2 == 1 {
			gzip = 6
		}
		shuffle := data[5]%2 == 1
		// A filtered dataset in a latest-version file is refused by
		// design; the fuzz contract covers the accepted combinations.
		latest := data[6]%2 == 1 && gzip == 0 && !shuffle
		chunk := 0
		if data[7]%2 == 1 {
			chunk = 32
		}
		n := d0
		if rank == 2 {
			n *= d1
		}
		shape := []int{d0}
		if rank == 2 {
			shape = append(shape, d1)
		}
		var values *core.Array
		var err error
		switch dtype {
		case 0:
			v := make([]float64, n)
			for i := range v {
				v[i] = float64(data[(8+i)%len(data)]) - 128 + float64(i%5)*0.25
			}
			values, err = core.FromFloats(v, shape...)
		case 1:
			v := make([]float32, n)
			for i := range v {
				v[i] = float32(data[(8+i)%len(data)]) - 128 + float32(i%5)*0.25
			}
			values, err = core.FromFloat32s(v, shape...)
		case 2:
			v := make([]int64, n)
			for i := range v {
				v[i] = int64(data[(8+i)%len(data)])*1000 - 128000 + int64(i)
			}
			values, err = core.FromInts(v, shape...)
		case 3:
			v := make([]bool, n)
			for i := range v {
				v[i] = data[(8+i)%len(data)]%2 == 1
			}
			values, err = core.FromBools(v, shape...)
		case 4:
			v := make([]int8, n)
			for i := range v {
				v[i] = int8(int(data[(8+i)%len(data)]) - 128)
			}
			values, err = core.FromInt8s(v, shape...)
		case 5:
			v := make([]uint8, n)
			for i := range v {
				v[i] = data[(8+i)%len(data)]
			}
			values, err = core.FromUint8s(v, shape...)
		case 6:
			v := make([]int16, n)
			for i := range v {
				v[i] = int16(int(data[(8+i)%len(data)])*257 - 32768)
			}
			values, err = core.FromInt16s(v, shape...)
		case 7:
			v := make([]uint16, n)
			for i := range v {
				v[i] = uint16(int(data[(8+i)%len(data)]) * 257)
			}
			values, err = core.FromUint16s(v, shape...)
		case 8:
			v := make([]int32, n)
			for i := range v {
				v[i] = int32(int64(data[(8+i)%len(data)])*1000000 - 128000000 + int64(i))
			}
			values, err = core.FromInt32s(v, shape...)
		default:
			v := make([]uint32, n)
			for i := range v {
				v[i] = uint32(uint64(data[(8+i)%len(data)])*1000000 + uint64(i))
			}
			values, err = core.FromUint32s(v, shape...)
		}
		if err != nil {
			t.Fatalf("build input: %v", err)
		}
		sets := []HDF5Dataset{{Path: "/d", Shape: shape, Values: values, Attrs: map[string]string{"k": "1"}}}
		attrs := map[string]map[string]string{"/": {"title": "fuzz"}}
		path := filepath.Join(t.TempDir(), "fuzz.h5")
		if err := SaveHDF5(path, sets, attrs, HDF5WriteOptions{Gzip: gzip, Shuffle: shuffle, Latest: latest, ChunkBytes: chunk}); err != nil {
			t.Fatalf("SaveHDF5 (%v): %v", shape, err)
		}
		back, err := LoadHDF5(path)
		if err != nil {
			t.Fatalf("LoadHDF5 round trip (%v): %v", shape, err)
		}
		if len(back) != 1 {
			t.Fatalf("round trip returned %d datasets, want 1", len(back))
		}
		d := back[0]
		if d.Path != "/d" {
			t.Fatalf("path = %q, want /d", d.Path)
		}
		if !slices.Equal(d.Shape, shape) {
			t.Fatalf("shape = %v, want %v", d.Shape, shape)
		}
		if d.Values.Dtype() != values.Dtype() {
			t.Fatalf("dtype = %s, want %s", d.Values.Dtype(), values.Dtype())
		}
		switch dtype {
		case 0:
			got, want := d.Values.RawFloats()[:n], values.RawFloats()[:n]
			for i := range want {
				if math.Float64bits(got[i]) != math.Float64bits(want[i]) {
					t.Fatalf("[%d] = %v, want %v (bit-exact)", i, got[i], want[i])
				}
			}
		case 1:
			got, want := d.Values.RawFloat32s()[:n], values.RawFloat32s()[:n]
			for i := range want {
				if math.Float32bits(got[i]) != math.Float32bits(want[i]) {
					t.Fatalf("[%d] = %v, want %v (bit-exact)", i, got[i], want[i])
				}
			}
		case 2:
			got, want := d.Values.RawInts()[:n], values.RawInts()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		case 3:
			got, want := d.Values.RawBools()[:n], values.RawBools()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %v, want %v", i, got[i], want[i])
				}
			}
		case 4:
			got, want := d.Values.RawInt8s()[:n], values.RawInt8s()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		case 5:
			got, want := d.Values.RawUint8s()[:n], values.RawUint8s()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		case 6:
			got, want := d.Values.RawInt16s()[:n], values.RawInt16s()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		case 7:
			got, want := d.Values.RawUint16s()[:n], values.RawUint16s()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		case 8:
			got, want := d.Values.RawInt32s()[:n], values.RawInt32s()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		default:
			got, want := d.Values.RawUint32s()[:n], values.RawUint32s()[:n]
			for i := range want {
				if got[i] != want[i] {
					t.Fatalf("[%d] = %d, want %d", i, got[i], want[i])
				}
			}
		}
		if d.Attrs["k"] != "1" {
			t.Fatalf("dataset attr k = %q, want 1", d.Attrs["k"])
		}
		if d.Attrs["title"] != "fuzz" {
			t.Fatalf("merged root attr title = %q, want fuzz", d.Attrs["title"])
		}
	})
}