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

package io

import (
	"encoding/binary"
	"strings"
	"testing"
)

// Regression pins for hostile FITS input: negative and missing axis
// cards, TBCOL and row-byte arithmetic that wraps, repeat counts past
// the int range, unbounded HDU skips and garbage PCOUNT, all refused by
// name before any payload is read.

// tableHostile builds a FITS table file from card bodies plus one
// 2880-byte data block of the given row content.
func tableHostile(cards []string, rows int) []byte {
	var b []byte
	for _, body := range cards {
		b = append(b, card(body)...)
	}
	b = append(b, card("END")...)
	if pad := len(b) % 2880; pad != 0 {
		b = append(b, make([]byte, 2880-pad)...)
	}
	b = append(b, make([]byte, 2880)...)
	_ = rows
	return b
}

// TestFITSNegativeAxisCard pins the refusal of a negative or
// missing NAXIS before the axis slice is allocated, and the image skip
// past a zero axis in the table reader.
func TestFITSNegativeAxisCard(t *testing.T) {
	build := func(bodies []string) []byte {
		var b []byte
		for _, body := range bodies {
			b = append(b, card(body)...)
		}
		b = append(b, card("END")...)
		if pad := len(b) % 2880; pad != 0 {
			b = append(b, make([]byte, 2880-pad)...)
		}
		return b
	}
	noNaxis := writeHostile(t, "no-naxis.fits", build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
	}))
	if _, _, err := LoadFITS(noNaxis); err == nil {
		t.Fatal("expected an error for a missing NAXIS card")
	}
	negative := writeHostile(t, "neg-naxis.fits", build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                   -3",
	}))
	if _, _, err := LoadFITS(negative); err == nil {
		t.Fatal("expected an error for a negative NAXIS")
	}
	// A zero axis followed by another divides the running product in
	// the table reader's image skip; the skip must pass it without
	// dividing by the zeroed product and report that no table follows.
	zeroAxis := writeHostile(t, "zero-axis.fits", build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                    2",
		"NAXIS1  =                    0",
		"NAXIS2  =                    1",
	}))
	if _, err := LoadFITSTable(zeroAxis); err == nil || !strings.Contains(err.Error(), "no table extension") {
		t.Fatalf("LoadFITSTable past a zero axis: err = %v, want the no-table report", err)
	}
	if _, _, err := LoadFITS(zeroAxis); err == nil {
		t.Fatal("LoadFITS: expected an error for a zero axis under positive NAXIS")
	}
}

// TestFITSHugeNaxisRefused pins that a NAXIS beyond the NAXISn cards
// the file carries is refused before the axis slice is allocated: a
// NAXIS of MaxInt64 used to reach make([]int, naxis) and panic with
// makeslice instead of answering the card-count error.
func TestFITSHugeNaxisRefused(t *testing.T) {
	build := func(bodies []string) []byte {
		var b []byte
		for _, body := range bodies {
			b = append(b, card(body)...)
		}
		b = append(b, card("END")...)
		if pad := len(b) % 2880; pad != 0 {
			b = append(b, make([]byte, 2880-pad)...)
		}
		return b
	}
	huge := writeHostile(t, "huge-naxis.fits", build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =        9223372036854775807",
	}))
	if _, _, err := LoadFITS(huge); err == nil || !strings.Contains(err.Error(), "NAXISn cards") {
		t.Fatalf("LoadFITS with NAXIS = MaxInt64: err = %v, want the card-count refusal", err)
	}
}

// TestASCIITableTBCOLWrap pins the overflow-free bound on
// TBCOL + width: a TBCOL near MaxInt64 used to wrap the sum negative,
// pass the guard and panic on the slice.
func TestASCIITableTBCOLWrap(t *testing.T) {
	data := tableHostile([]string{
		"XTENSION= 'TABLE   '",
		"BITPIX  =                    8",
		"NAXIS   =                    2",
		"NAXIS1  =                   20",
		"NAXIS2  =                    1",
		"TFIELDS =                    1",
		"TFORM1  = 'I10               '",
		"TBCOL1  = '9223372036854775807'",
	}, 1)
	path := writeHostile(t, "tbcol-wrap.fits", data)
	if _, err := LoadFITSTable(path); err == nil {
		t.Fatal("expected an error for a TBCOL whose sum with the width wraps")
	}
}

// TestBINTABLERowBytesWrap pins the per-column bound on the
// row prefix sum: two A columns of 2^62 repeat each used to wrap
// rowBytes negative, skip the truncation guard and read past the file.
func TestBINTABLERowBytesWrap(t *testing.T) {
	data := tableHostile([]string{
		"XTENSION= 'BINTABLE'",
		"BITPIX  =                    8",
		"NAXIS   =                    2",
		"NAXIS1  =                   16",
		"NAXIS2  =                    1",
		"TFIELDS =                    2",
		"TFORM1  = '4611686018427387904A'",
		"TFORM2  = '4611686018427387904A'",
	}, 1)
	path := writeHostile(t, "rowbytes-wrap.fits", data)
	if _, err := LoadFITSTable(path); err == nil {
		t.Fatal("expected an error for column widths whose sum wraps")
	}
}

// TestTFORMRepeatOverflowRefused pins the named refusal for a
// repeat count that does not fit an int, which used to fall back to a
// silent scalar of width 1.
func TestTFORMRepeatOverflowRefused(t *testing.T) {
	if _, _, err := parseTFORM("99999999999999999999E"); err == nil {
		t.Fatal("parseTFORM: expected an error for an overflowing repeat")
	}
	if r, code, err := parseTFORM("16A"); err != nil || r != 16 || code != "A" {
		t.Fatalf("parseTFORM(16A) = %d, %q, %v", r, code, err)
	}
}

// TestImageHDUSkipBounded pins the image-HDU skip arithmetic:
// axis products that would wrap the int cannot skip into attacker
// bytes; the file is refused as truncated instead.
func TestImageHDUSkipBounded(t *testing.T) {
	var b []byte
	for _, body := range []string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                    2",
		"NAXIS1  =           2147483647",
		"NAXIS2  =                    4",
		"PCOUNT  =                    0",
		"GCOUNT  =                    2",
	} {
		b = append(b, card(body)...)
	}
	b = append(b, card("END")...)
	if pad := len(b) % 2880; pad != 0 {
		b = append(b, make([]byte, 2880-pad)...)
	}
	b = append(b, make([]byte, 2880)...)
	path := writeHostile(t, "skip-wrap.fits", b)
	// The table reader skips the image HDU to look for a table after
	// it; the wrapped product used to land the skip inside the image
	// data. Now the skip is bounded and the file has no table.
	if _, err := LoadFITSTable(path); err == nil {
		t.Fatal("expected an error: no table HDU follows the bounded skip")
	}
}

// TestPCOUNTGarbageRefused pins that a non-integer PCOUNT is
// an error, not a silent zero.
func TestPCOUNTGarbageRefused(t *testing.T) {
	var b []byte
	for _, body := range []string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                    2",
		"NAXIS1  =                    2",
		"NAXIS2  =                    2",
		"PCOUNT  = '1e3'",
	} {
		b = append(b, card(body)...)
	}
	b = append(b, card("END")...)
	if pad := len(b) % 2880; pad != 0 {
		b = append(b, make([]byte, 2880-pad)...)
	}
	b = append(b, make([]byte, 2880)...)
	path := writeHostile(t, "pcount.fits", b)
	if _, err := LoadFITSTable(path); err == nil {
		t.Fatal("expected an error for a non-integer PCOUNT")
	}
}

// TestFITSNAXISnOrder pins that NAXISn cards are bound by
// their axis number: a header writing NAXIS2 before NAXIS1 under
// NAXIS = 2 with matching extents stays the image it declares, and a
// missing axis is an error.
func TestFITSNAXISnOrder(t *testing.T) {
	build := func(bodies []string) []byte {
		var b []byte
		for _, body := range bodies {
			b = append(b, card(body)...)
		}
		b = append(b, card("END")...)
		if pad := len(b) % 2880; pad != 0 {
			b = append(b, make([]byte, 2880-pad)...)
		}
		// 2x2 float64 image data.
		for range 4 {
			b = binary.BigEndian.AppendUint64(b, 1)
		}
		return b
	}
	// Reversed card order: the image is still 3 by 2.
	path := writeHostile(t, "order.fits", build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                    2",
		"NAXIS2  =                    2",
		"NAXIS1  =                    3",
	}))
	// 3x2 = 6 doubles but only 4 present: truncated either way, so the
	// claim is checked by the value the reader reports. Give it enough
	// data instead: rebuild with full payload.
	full := build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                    2",
		"NAXIS2  =                    2",
		"NAXIS1  =                    3",
	})
	full = append(full, make([]byte, 2880)...)
	path = writeHostile(t, "order-full.fits", full)
	a, _, err := LoadFITS(path)
	if err != nil {
		t.Fatalf("LoadFITS with reversed NAXISn cards: %v", err)
	}
	if a.Shape()[0] != 2 || a.Shape()[1] != 3 {
		t.Fatalf("shape = %v, want [2 3] (Fortran order, NAXIS1 = 3)", a.Shape())
	}
	// A missing axis is refused by name.
	missing := build([]string{
		"SIMPLE  =                    T",
		"BITPIX  =                   -64",
		"NAXIS   =                    2",
		"NAXIS1  =                    2",
	})
	mpath := writeHostile(t, "missing.fits", missing)
	if _, _, err := LoadFITS(mpath); err == nil || !strings.Contains(err.Error(), "NAXIS") {
		t.Fatalf("err = %v, want the missing NAXISn card named", err)
	}
}