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

package io

import (
	"encoding/binary"
	"math"
	"os"
	"path/filepath"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"strings"
	"testing"
)

func fitsTempPath(t *testing.T) string {
	t.Helper()
	return filepath.Join(t.TempDir(), "image.fits")
}

// padCard right-pads a card body with spaces to the full 80 bytes.
func padCard(text string) []byte {
	return []byte(text + strings.Repeat(" ", 80-len(text)))
}

// TestFITSRoundTripFloat64 moves a rank-2 float64 image with negative
// and non-round values through the file format and back, values and
// header strings included.
func TestFITSRoundTripFloat64(t *testing.T) {
	a := mustFloats(t, []float64{
		1.5, -2.25, 3.125, 4,
		-5.5, 6.75, -7.875, 8,
		9.25, -10.5, 11.125, -12,
	}, 3, 4)
	path := fitsTempPath(t)
	headers := map[string]string{
		"OBJECT":   "M31",
		"OBSERVER": "petr's dome",
		"EXPTIME":  "600",
	}
	if err := SaveFITS(path, a, headers); err != nil {
		t.Fatalf("SaveFITS: %v", err)
	}
	back, hdr, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS: %v", err)
	}
	if back.Dtype() != core.Float || back.NDim() != 2 || back.Shape()[0] != 3 || back.Shape()[1] != 4 {
		t.Fatalf("shape/dtype mismatch: %v %s", back.Shape(), back.Dtype())
	}
	for i := range a.Len() {
		if back.FloatAt(i) != a.FloatAt(i) {
			t.Fatalf("value[%d] = %g, want %g", i, back.FloatAt(i), a.FloatAt(i))
		}
	}
	for key, want := range map[string]string{
		"OBJECT":   "M31",
		"OBSERVER": "petr's dome",
		"EXPTIME":  "600",
	} {
		if hdr[key] != want {
			t.Fatalf("header %q = %q, want %q", key, hdr[key], want)
		}
	}
}

// TestFITSAxisConvention pins the wire format against the raw bytes:
// NAXIS1 must carry the fastest (last Go) axis and the payload must
// be big-endian in flat row-major order.
func TestFITSAxisConvention(t *testing.T) {
	a := mustFloats(t, []float64{1, 2, 3, 4, 5, 6}, 2, 3)
	path := fitsTempPath(t)
	if err := SaveFITS(path, a, nil); err != nil {
		t.Fatalf("SaveFITS: %v", err)
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	if len(raw)%2880 != 0 {
		t.Fatalf("file length %d is not a multiple of 2880", len(raw))
	}
	header := raw[:2880]
	for _, want := range []string{
		"SIMPLE  =                    T",
		"BITPIX  =                  -64",
		"NAXIS   =                    2",
		"NAXIS1  =                    3",
		"NAXIS2  =                    2",
	} {
		if !strings.Contains(string(header), want) {
			t.Fatalf("header misses %q", want)
		}
	}
	endAt := strings.Index(string(header), "END")
	if endAt < 0 {
		t.Fatal("header has no END card")
	}
	var word [8]byte
	for i := range 6 {
		binary.BigEndian.PutUint64(word[:], math.Float64bits(float64(i+1)))
		got := raw[2880+i*8 : 2880+i*8+8]
		if string(got) != string(word[:]) {
			t.Fatalf("payload word %d = % x, want % x", i, got, word)
		}
	}
}

// TestFITSRoundTripFloat32 keeps the float32 element type through the
// round trip, which is what BITPIX −32 stores.
func TestFITSRoundTripFloat32(t *testing.T) {
	a := core.New(core.Float32, 5)
	for i := range 5 {
		a.RawFloat32s()[i] = float32(i) * 1.25
	}
	path := fitsTempPath(t)
	if err := SaveFITS(path, a, nil); err != nil {
		t.Fatalf("SaveFITS: %v", err)
	}
	back, _, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS: %v", err)
	}
	if back.Dtype() != core.Float32 {
		t.Fatalf("dtype = %s, want core.Float32", back.Dtype())
	}
	for i := range 5 {
		if back.RawFloat32s()[i] != a.RawFloat32s()[i] {
			t.Fatalf("value[%d] = %g, want %g", i, back.RawFloat32s()[i], a.RawFloat32s()[i])
		}
	}
}

// TestFITSSkipsValuelessCards checks COMMENT and HISTORY cards are
// tolerated and skipped rather than parsed as values.
func TestFITSSkipsValuelessCards(t *testing.T) {
	a := mustFloats(t, []float64{1, 2, 3, 4}, 2, 2)
	path := fitsTempPath(t)
	if err := SaveFITS(path, a, map[string]string{"OBJECT": "TEST"}); err != nil {
		t.Fatalf("SaveFITS: %v", err)
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	// Splice two valueless cards in before the END card and re-pad the
	// header to the block boundary. The END card is matched in full
	// ("END" plus its padding) because EXTEND contains the same three
	// letters.
	endCard := strings.Index(string(raw), "END"+strings.Repeat(" ", 77))
	if endCard < 0 {
		t.Fatal("no END card")
	}
	spliced := append([]byte{}, raw[:endCard]...)
	spliced = append(spliced, padCard("COMMENT   a note without a value")...)
	spliced = append(spliced, padCard("HISTORY   an audit trail entry")...)
	spliced = append(spliced, padCard("END")...)
	for len(spliced)%2880 != 0 {
		spliced = append(spliced, padCard("")...)
	}
	spliced = append(spliced, raw[2880:]...)
	if err := os.WriteFile(path, spliced, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	back, hdr, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS: %v", err)
	}
	if back.Len() != 4 || back.FloatAt(3) != 4 {
		t.Fatalf("payload damaged: %v %v", back.Shape(), back.RawFloats())
	}
	if hdr["OBJECT"] != "TEST" {
		t.Fatalf("OBJECT = %q, want %q", hdr["OBJECT"], "TEST")
	}
	if _, ok := hdr["COMMENT"]; ok {
		t.Fatal("COMMENT card must not enter the header map")
	}
}

// TestFITSErrors covers every refusal path in the format layer.
func TestFITSErrors(t *testing.T) {
	path := fitsTempPath(t)

	ints, ierr := core.FromInts([]int64{1, 2, 3, 4}, 2, 2)
	if ierr != nil {
		t.Fatalf("FromInts: %v", ierr)
	}
	if err := SaveFITS(path, ints, nil); err == nil {
		t.Fatal("int64 input: want an error")
	}
	if err := SaveFITS(path, mustComplexes(t, []complex128{1, 2, 3, 4}, 2, 2), nil); err == nil {
		t.Fatal("complex input: want an error")
	}
	ok := mustFloats(t, []float64{1, 2}, 2)
	if err := SaveFITS(path, ok, map[string]string{"TOOLONGKEYWORD": "x"}); err == nil {
		t.Fatal("long keyword: want an error")
	}
	if err := SaveFITS(path, ok, map[string]string{"BITPIX": "x"}); err == nil {
		t.Fatal("reserved keyword: want an error")
	}
	if err := SaveFITS(path, ok, map[string]string{"naxis2": "x"}); err == nil {
		t.Fatal("NAXISn keyword: want an error")
	}
	if err := SaveFITS(path, ok, map[string]string{"BAD KEY": "x"}); err == nil {
		t.Fatal("space in keyword: want an error")
	}
	if err := SaveFITS(path, ok, map[string]string{"NOTE": strings.Repeat("x", 69)}); err == nil {
		t.Fatal("overlong value: want an error")
	}

	// A valid file to damage in every way the parser must catch.
	if err := SaveFITS(path, ok, nil); err != nil {
		t.Fatalf("SaveFITS: %v", err)
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	with := func(mutate func([]byte) []byte) {
		t.Helper()
		if err := os.WriteFile(path, mutate(append([]byte{}, raw...)), 0o644); err != nil {
			t.Fatalf("WriteFile: %v", err)
		}
		if _, _, err := LoadFITS(path); err == nil {
			t.Fatal("damaged file: want an error")
		}
	}
	with(func(b []byte) []byte { return b[:100] })                       // no END card
	with(func(b []byte) []byte { return b[:2880] })                      // data truncated
	with(func(b []byte) []byte { b[30] = 'F'; return b })                // SIMPLE = F
	with(func(b []byte) []byte { copy(b[11:20], "XTENSION"); return b }) // wrong first card

	// BITPIX 8 (unsigned bytes) is outside the supported image types.
	unsupported := append(append(append(append([]byte{},
		padCard("SIMPLE  =                    T")...),
		padCard("BITPIX  =                    8")...),
		padCard("NAXIS   =                    1")...),
		padCard("NAXIS1  =                    4")...)
	unsupported = append(unsupported, padCard("END")...)
	unsupported = append(unsupported, make([]byte, 2880-5*80)...)
	unsupported = append(unsupported, make([]byte, 2880)...)
	if err := os.WriteFile(path, unsupported, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	if _, _, err := LoadFITS(path); err == nil {
		t.Fatal("BITPIX 8: want an error")
	}

	// A hostile header whose axis product overflows int must be
	// refused as truncated data, never panic the process.
	hostile := append(append(append(append([]byte{},
		padCard("SIMPLE  =                    T")...),
		padCard("BITPIX  =                  -64")...),
		padCard("NAXIS   =                    2")...),
		padCard("NAXIS1  =          1099511627776")...)
	hostile = append(hostile, padCard("NAXIS2  =          1099511627776")...)
	hostile = append(hostile, padCard("END")...)
	hostile = append(hostile, make([]byte, 2880-6*80)...)
	hostile = append(hostile, make([]byte, 2880)...)
	if err := os.WriteFile(path, hostile, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	if _, _, err := LoadFITS(path); err == nil {
		t.Fatal("overflowing axis product: want an error")
	}

	// An XTENSION-first file is an extension, not a primary image.
	ext := append([]byte{}, padCard("XTENSION= 'IMAGE   '")...)
	ext = append(ext, raw[80:]...)
	if err := os.WriteFile(path, ext, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	if _, _, err := LoadFITS(path); err == nil {
		t.Fatal("extension header: want an error")
	}

	// The intact file still loads after all the damage around it.
	if err := os.WriteFile(path, raw, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	if _, _, err := LoadFITS(path); err != nil {
		t.Fatalf("intact file: %v", err)
	}
}

// TestCommentaryCardsAreNotValueCards pins the commentary rule: a
// COMMENT or HISTORY card may legitimately carry an "= " sequence in
// columns 9-10, and such a card must be skipped as commentary, not
// parsed as a keyword with a value.
func TestCommentaryCardsAreNotValueCards(t *testing.T) {
	cards := []string{
		fitsBoolCard("SIMPLE", true),
		fitsIntCard("BITPIX", -64),
		fitsIntCard("NAXIS", 1),
		fitsIntCard("NAXIS1", 2),
		fitsPadCard("COMMENT   = this looks like a value card"),
		fitsPadCard("HISTORY   = so does this one"),
		fitsStringCardRaw("OBSERVER", "tester"),
		fitsEndCard(),
	}
	data := fitsAppendCards(nil, cards)
	payload := []byte{0x3f, 0xf0, 0, 0, 0, 0, 0, 0, 0x40, 0, 0, 0, 0, 0, 0, 0} // 1.0, 2.0
	data = append(data, payload...)
	data = fitsAppendZeroPad(data)

	path := filepath.Join(t.TempDir(), "commentary.fits")
	if err := osWriteFile(path, data); err != nil {
		t.Fatalf("os.WriteFile: %v", err)
	}
	img, headers, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS: %v", err)
	}
	if img.Len() != 2 || img.FloatAt(0) != 1 || img.FloatAt(1) != 2 {
		t.Fatalf("image = %s, want [1, 2]", img)
	}
	if _, ok := headers["COMMENT"]; ok {
		t.Error("COMMENT was parsed as a value keyword")
	}
	if _, ok := headers["HISTORY"]; ok {
		t.Error("HISTORY was parsed as a value keyword")
	}
	if headers["OBSERVER"] != "tester" {
		t.Errorf("OBSERVER = %q, want tester", headers["OBSERVER"])
	}
}

// TestFitsCheckKeywordRefusesCommentary pins that COMMENT and HISTORY
// are refused as user keywords: they carry no value in the format.
func TestFitsCheckKeywordRefusesCommentary(t *testing.T) {
	dir := t.TempDir()
	img := mustFloats(t, []float64{1, 2}, 2)
	for _, kw := range []string{"COMMENT", "HISTORY"} {
		if err := SaveFITS(filepath.Join(dir, strings.ToLower(kw)+".fits"), img, map[string]string{kw: "x"}); err == nil {
			t.Errorf("SaveFITS accepted the reserved keyword %q", kw)
		}
	}
}