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

package io

import (
	"encoding/hex"
	"os"
	"path/filepath"
	"testing"

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

// The record (unlimited) dimension: written and read back, and read
// from a file another implementation wrote.

// TestNetCDFRecordRoundTrip writes a file with a record dimension and a
// fixed variable, then reads it back: the record count, the shapes and
// the values must survive, and the record axis must stay first.
func TestNetCDFRecordRoundTrip(t *testing.T) {
	path := filepath.Join(t.TempDir(), "rec.nc")
	tp := mustFloats(t, []float64{1.5, 2.5, 3.5}, 3)
	sp := mustInts(t, []int64{1, 2, 3, 4, 5, 6, 7, 8, 9}, 3, 3)
	vp := mustFloats32(t, []float32{7, 8, 9}, 3)
	dims := []NetCDFDim{{Name: "time", Length: 0}, {Name: "x", Length: 3}}
	vars := []NetCDFVar{
		{Name: "t", Dims: []string{"time"}, Values: tp},
		{Name: "s", Dims: []string{"time", "x"}, Values: sp},
		{Name: "v", Dims: []string{"x"}, Values: vp},
	}
	if err := SaveNetCDF(path, dims, vars, map[string]string{"title": "record round trip"}); err != nil {
		t.Fatalf("SaveNetCDF: %v", err)
	}
	gotDims, gotVars, attrs, err := LoadNetCDF(path)
	if err != nil {
		t.Fatalf("LoadNetCDF: %v", err)
	}
	if len(gotDims) != 2 || gotDims[0].Name != "time" || gotDims[0].Length != 0 || gotDims[1].Length != 3 {
		t.Fatalf("dims = %+v, want the record dimension first with length 0", gotDims)
	}
	if attrs["title"] != "record round trip" {
		t.Fatalf("attrs = %v", attrs)
	}
	if len(gotVars) != 3 {
		t.Fatalf("vars = %d, want 3", len(gotVars))
	}
	byName := map[string]NetCDFVar{}
	for _, v := range gotVars {
		byName[v.Name] = v
	}
	if s := byName["t"].Values.Shape(); s[0] != 3 {
		t.Fatalf("t shape = %v, want [3]", s)
	}
	if got := byName["s"].Values.Shape(); got[0] != 3 || got[1] != 3 {
		t.Fatalf("s shape = %v, want [3 3]", got)
	}
	for i, want := range []float64{1.5, 2.5, 3.5} {
		if got := byName["t"].Values.FloatAt(i); got != want {
			t.Fatalf("t[%d] = %v, want %v", i, got, want)
		}
	}
	for i, want := range []float64{1, 2, 3, 4, 5, 6, 7, 8, 9} {
		if got := byName["s"].Values.FloatAt(i); got != want {
			t.Fatalf("s[%d] = %v, want %v", i, got, want)
		}
	}
	for i, want := range []float64{7, 8, 9} {
		if got := byName["v"].Values.FloatAt(i); got != want {
			t.Fatalf("v[%d] = %v, want %v", i, got, want)
		}
	}
}

// TestNetCDFRecordFromNetCDF4 reads a file the netCDF reference
// implementation wrote: an unlimited time axis, a 1-D record variable
// of NC_DOUBLE, a 2-D record variable of NC_SHORT whose six-byte slab
// is padded to eight on disk, and a fixed NC_FLOAT variable. It pins
// the layout a foreign writer produces, padding included.
func TestNetCDFRecordFromNetCDF4(t *testing.T) {
	const dump = "" +
		"43444601000000020000000a000000020000000474696d6500000000000000017800000000000003" +
		"00000000000000000000000b00000003000000017400000000000001000000000000000000000000" +
		"0000000600000008000000b400000001730000000000000200000000000000010000000000000000" +
		"0000000300000008000000bc00000001760000000000000100000001000000000000000000000005" +
		"0000000c000000a840e0000041000000411000003ff8000000000000000100020003800140040000" +
		"000000000004000500068001"
	raw, err := hex.DecodeString(dump)
	if err != nil {
		t.Fatalf("hex: %v", err)
	}
	path := filepath.Join(t.TempDir(), "foreign.nc")
	if err := os.WriteFile(path, raw, 0o644); err != nil {
		t.Fatalf("write: %v", err)
	}
	dims, vars, _, err := LoadNetCDF(path)
	if err != nil {
		t.Fatalf("LoadNetCDF: %v", err)
	}
	if len(dims) != 2 || dims[0].Name != "time" || dims[0].Length != 0 || dims[1].Length != 3 {
		t.Fatalf("dims = %+v", dims)
	}
	if len(vars) != 3 {
		t.Fatalf("vars = %d, want 3", len(vars))
	}
	byName := map[string]NetCDFVar{}
	for _, v := range vars {
		byName[v.Name] = v
	}
	// Two records of one double.
	if got := byName["t"].Values.RawFloats(); len(got) != 2 || got[0] != 1.5 || got[1] != 2.5 {
		t.Fatalf("t = %v, want [1.5 2.5]", got)
	}
	// Two records of three shorts, each slab padded on disk.
	if got := byName["s"].Values.Shape(); got[0] != 2 || got[1] != 3 {
		t.Fatalf("s shape = %v, want [2 3]", got)
	}
	for i, want := range []float64{1, 2, 3, 4, 5, 6} {
		if got := byName["s"].Values.FloatAt(i); got != want {
			t.Fatalf("s[%d] = %v, want %v", i, got, want)
		}
	}
	// The fixed variable, written before the records. Every variable
	// comes back widened to float64, whatever the file stores.
	for i, want := range []float64{7, 8, 9} {
		if got := byName["v"].Values.FloatAt(i); got != want {
			t.Fatalf("v[%d] = %v, want %v", i, got, want)
		}
	}
}

// TestNetCDFRecordZeroRecords pins the empty unlimited dimension: a
// file may declare the record dimension and hold no records at all.
func TestNetCDFRecordZeroRecords(t *testing.T) {
	path := filepath.Join(t.TempDir(), "empty.nc")
	empty, err := core.FromFloats([]float64{}, 0)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	err = SaveNetCDF(path, []NetCDFDim{{Name: "time", Length: 0}},
		[]NetCDFVar{{Name: "t", Dims: []string{"time"}, Values: empty}}, nil)
	if err != nil {
		t.Fatalf("SaveNetCDF: %v", err)
	}
	dims, vars, _, err := LoadNetCDF(path)
	if err != nil {
		t.Fatalf("LoadNetCDF: %v", err)
	}
	if len(dims) != 1 || dims[0].Length != 0 {
		t.Fatalf("dims = %+v", dims)
	}
	if len(vars) != 1 || vars[0].Values.Len() != 0 {
		t.Fatalf("vars = %+v, want one empty variable", vars)
	}
	if s := vars[0].Values.Shape(); s[0] != 0 {
		t.Fatalf("shape = %v, want [0]", s)
	}
}