package udpsprotocol

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"time"
)

// ─── Constants ───────────────────────────────────────────────────────────────

const (
	MagicUDPS uint32 = 0x53504455 // 'UDPS' little-endian

	PktData       uint8 = 0
	PktConfig     uint8 = 1
	PktACK        uint8 = 2
	PktConnect    uint8 = 3
	PktDisconnect uint8 = 4

	HeaderSize   = 17
	SigDescSize  = 136
	NoTimeSignal = uint32(0xFFFFFFFF)

	QuantNone   uint8 = 0
	QuantUint8  uint8 = 1
	QuantInt8   uint8 = 2
	QuantUint16 uint8 = 3
	QuantInt16  uint8 = 4

	// TimeMode values – must match UDPStreamerTimeMode enum in UDPStreamer.h
	TimeModePacket      uint8 = 0 // use wall-clock packet arrival time
	TimeModeFullArray   uint8 = 1 // TimeSignal has same N elements; not expanded here
	TimeModeFirstSample uint8 = 2 // TimeSignal scalar = time of element [0]
	TimeModeLastSample  uint8 = 3 // TimeSignal scalar = time of element [N-1]

	// PublishMode values – must match UDPStreamerPublishMode enum in UDPStreamer.h
	PublishModeStrict     uint8 = 0 // one packet per Synchronise() call
	PublishModeAccumulate uint8 = 1 // variable batch; DATA has [8 HRT][4 numSamples][signals...]
	PublishModeDecimate   uint8 = 2 // one packet every Ratio calls
)

// ─── Packet header (17 bytes, little-endian, packed) ─────────────────────────

type PacketHeader struct {
	Magic          uint32
	Type           uint8
	Counter        uint32
	FragmentIdx    uint16
	TotalFragments uint16
	PayloadBytes   uint32
}

// ParseHeader decodes exactly HeaderSize bytes into a PacketHeader.
func ParseHeader(b []byte) (PacketHeader, error) {
	if len(b) < HeaderSize {
		return PacketHeader{}, fmt.Errorf("header too short: %d bytes", len(b))
	}
	var h PacketHeader
	r := bytes.NewReader(b[:HeaderSize])
	if err := binary.Read(r, binary.LittleEndian, &h); err != nil {
		return PacketHeader{}, err
	}
	if h.Magic != MagicUDPS {
		return PacketHeader{}, fmt.Errorf("bad magic: 0x%08X", h.Magic)
	}
	return h, nil
}

// buildHeader serialises a PacketHeader to a 17-byte slice.
func buildHeader(h PacketHeader) []byte {
	buf := new(bytes.Buffer)
	_ = binary.Write(buf, binary.LittleEndian, h)
	return buf.Bytes()
}

// BuildConnectPacket returns a 17-byte CONNECT datagram.
func BuildConnectPacket() []byte {
	return buildHeader(PacketHeader{
		Magic:          MagicUDPS,
		Type:           PktConnect,
		Counter:        0,
		FragmentIdx:    0,
		TotalFragments: 1,
		PayloadBytes:   0,
	})
}

// BuildDisconnectPacket returns a 17-byte DISCONNECT datagram.
func BuildDisconnectPacket() []byte {
	return buildHeader(PacketHeader{
		Magic:          MagicUDPS,
		Type:           PktDisconnect,
		Counter:        0,
		FragmentIdx:    0,
		TotalFragments: 1,
		PayloadBytes:   0,
	})
}

// ─── Signal descriptor (136 bytes) ───────────────────────────────────────────

// SignalInfo holds the parsed metadata for one signal.
type SignalInfo struct {
	Name          string  `json:"name"`
	TypeCode      uint8   `json:"typeCode"`
	QuantType     uint8   `json:"quantType"`
	NumDimensions uint8   `json:"numDimensions"`
	NumRows       uint32  `json:"numRows"`
	NumCols       uint32  `json:"numCols"`
	RangeMin      float64 `json:"rangeMin"`
	RangeMax      float64 `json:"rangeMax"`
	TimeMode      uint8   `json:"timeMode"`
	SamplingRate  float64 `json:"samplingRate"`
	TimeSignalIdx uint32  `json:"timeSignalIdx"`
	Unit          string  `json:"unit"`
}

// NumElements returns the total number of scalar values in one sample of this signal.
func (s SignalInfo) NumElements() int {
	r := int(s.NumRows)
	c := int(s.NumCols)
	if r == 0 {
		r = 1
	}
	if c == 0 {
		c = 1
	}
	/* HI-2: cap at 1M to prevent integer overflow / OOM from crafted packets */
	n := r * c
	if n < 0 || n > 1024*1024 {
		return 1024 * 1024
	}
	return n
}

// rawTypeSize returns the byte size for one element of the raw (unquantised) type.
func rawTypeSize(typeCode uint8) int {
	switch typeCode {
	case 0, 1: // uint8, int8
		return 1
	case 2, 3: // uint16, int16
		return 2
	case 4, 5: // uint32, int32
		return 4
	case 6, 7: // uint64, int64
		return 8
	case 8: // float32
		return 4
	case 9: // float64
		return 8
	default:
		return 1
	}
}

// quantSize returns the byte size of one quantised element.
func quantSize(qt uint8) int {
	switch qt {
	case QuantUint8, QuantInt8:
		return 1
	case QuantUint16, QuantInt16:
		return 2
	default:
		return 0
	}
}

// readRawElement reads one element at offset and converts it to float64.
func readRawElement(b []byte, offset int, typeCode uint8) float64 {
	switch typeCode {
	case 0:
		return float64(b[offset])
	case 1:
		return float64(int8(b[offset]))
	case 2:
		return float64(binary.LittleEndian.Uint16(b[offset:]))
	case 3:
		return float64(int16(binary.LittleEndian.Uint16(b[offset:])))
	case 4:
		return float64(binary.LittleEndian.Uint32(b[offset:]))
	case 5:
		return float64(int32(binary.LittleEndian.Uint32(b[offset:])))
	case 6:
		return float64(binary.LittleEndian.Uint64(b[offset:]))
	case 7:
		return float64(int64(binary.LittleEndian.Uint64(b[offset:])))
	case 8:
		bits := binary.LittleEndian.Uint32(b[offset:])
		return float64(math.Float32frombits(bits))
	case 9:
		bits := binary.LittleEndian.Uint64(b[offset:])
		return math.Float64frombits(bits)
	default:
		return 0
	}
}

// dequantise converts a raw quantised integer to a physical float64.
func dequantise(qt uint8, raw uint16, rangeMin, rangeMax float64) float64 {
	span := rangeMax - rangeMin
	switch qt {
	case QuantUint8:
		return rangeMin + (float64(uint8(raw))/255.0)*span
	case QuantInt8:
		return rangeMin + (float64(int8(raw)+127)/254.0)*span
	case QuantUint16:
		return rangeMin + (float64(raw)/65535.0)*span
	case QuantInt16:
		return rangeMin + (float64(int16(raw)+32767)/65534.0)*span
	default:
		return 0
	}
}

// nullTermString converts a zero-padded byte slice to a Go string.
func nullTermString(b []byte) string {
	n := bytes.IndexByte(b, 0)
	if n < 0 {
		return string(b)
	}
	return string(b[:n])
}

// ─── CONFIG payload parser ────────────────────────────────────────────────────

// ParseConfig decodes a fully-reassembled CONFIG payload.
// Returns the signal list, the publishing mode byte (PublishMode*), and any error.
func ParseConfig(payload []byte) ([]SignalInfo, uint8, error) {
	if len(payload) < 4 {
		return nil, 0, fmt.Errorf("config payload too short")
	}
	numSigs := binary.LittleEndian.Uint32(payload[0:4])
	/* HI-2: validate numSigs against payload length before allocating */
	maxSigs := uint32(len(payload) / SigDescSize)
	if numSigs > maxSigs {
		return nil, 0, fmt.Errorf("config claims %d signals but payload can hold at most %d", numSigs, maxSigs)
	}
	offset := 4
	sigs := make([]SignalInfo, 0, numSigs)
	for i := uint32(0); i < numSigs; i++ {
		if offset+SigDescSize > len(payload) {
			return nil, 0, fmt.Errorf("config payload truncated at signal %d", i)
		}
		raw := payload[offset : offset+SigDescSize]
		si := SignalInfo{
			Name:          nullTermString(raw[0:64]),
			TypeCode:      raw[64],
			QuantType:     raw[65],
			NumDimensions: raw[66],
			NumRows:       binary.LittleEndian.Uint32(raw[67:71]),
			NumCols:       binary.LittleEndian.Uint32(raw[71:75]),
			RangeMin:      math.Float64frombits(binary.LittleEndian.Uint64(raw[75:83])),
			RangeMax:      math.Float64frombits(binary.LittleEndian.Uint64(raw[83:91])),
			TimeMode:      raw[91],
			SamplingRate:  math.Float64frombits(binary.LittleEndian.Uint64(raw[92:100])),
			TimeSignalIdx: binary.LittleEndian.Uint32(raw[100:104]),
			Unit:          nullTermString(raw[104:136]),
		}
		sigs = append(sigs, si)
		offset += SigDescSize
	}
	// Trailing publish-mode byte (added after signal descriptors).
	publishMode := PublishModeStrict
	if offset < len(payload) {
		publishMode = payload[offset]
	}
	return sigs, publishMode, nil
}

// ─── DATA payload parser ──────────────────────────────────────────────────────

// DataSample holds the decoded values from one DATA packet.
type DataSample struct {
	HRTTimestamp uint64
	WallTime     time.Time            // wall-clock time at UDP arrival; used as x-axis
	Values       map[string][]float64 // key = signal name, value = []float64 with NumElements entries
}

// parseElems reads n elements for sig from payload at offset, advancing offset.
// Returns the slice of float64 values and the new offset.
func parseElems(payload []byte, offset, n int, sig SignalInfo) ([]float64, int, error) {
	elems := make([]float64, n)
	if sig.QuantType == QuantNone {
		sz := rawTypeSize(sig.TypeCode)
		needed := n * sz
		if offset+needed > len(payload) {
			return nil, offset, fmt.Errorf("data payload truncated for signal %q", sig.Name)
		}
		for i := 0; i < n; i++ {
			elems[i] = readRawElement(payload, offset+i*sz, sig.TypeCode)
		}
		offset += needed
	} else {
		sz := quantSize(sig.QuantType)
		needed := n * sz
		if offset+needed > len(payload) {
			return nil, offset, fmt.Errorf("data payload truncated (quant) for signal %q", sig.Name)
		}
		for i := 0; i < n; i++ {
			var raw uint16
			if sz == 1 {
				raw = uint16(payload[offset+i])
			} else {
				raw = binary.LittleEndian.Uint16(payload[offset+i*2:])
			}
			elems[i] = dequantise(sig.QuantType, raw, sig.RangeMin, sig.RangeMax)
		}
		offset += needed
	}
	return elems, offset, nil
}

// ParseData decodes a fully-reassembled DATA payload using the provided signal config
// and publishing mode. arrivalTime is the wall-clock time at which the packet arrived.
//
// For PublishModeAccumulate the payload format is:
//
//	[8 HRT][4 numSamples][for each signal: accumulated scalars → numSamples elems; arrays → NumElements elems]
//
// The function returns one DataSample per accumulated snapshot so the hub can
// process each slot independently with its own timestamp.
func ParseData(payload []byte, sigs []SignalInfo, publishMode uint8, arrivalTime time.Time) ([]DataSample, error) {
	if len(payload) < 8 {
		return nil, fmt.Errorf("data payload too short")
	}
	hrt := binary.LittleEndian.Uint64(payload[0:8])
	offset := 8

	if publishMode == PublishModeAccumulate {
		if len(payload) < 12 {
			return nil, fmt.Errorf("accumulate data payload too short (missing numSamples)")
		}
		numSamples := int(binary.LittleEndian.Uint32(payload[8:12]))
		offset = 12
		if numSamples == 0 {
			return []DataSample{}, nil
		}
		/* HI-2: sanity-cap numSamples to prevent OOM from crafted packets */
		if numSamples < 0 || numSamples > 1024*1024 {
			return nil, fmt.Errorf("accumulate numSamples %d out of range", numSamples)
		}

		// Parse per-signal data blocks (all slots for a signal are contiguous).
		accumVals := make(map[string][]float64, len(sigs)) // scalars: numSamples values
		fixedVals := make(map[string][]float64, len(sigs)) // arrays: NumElements values

		for _, sig := range sigs {
			n := sig.NumElements()
			if n == 1 {
				// Accumulated scalar: read numSamples back-to-back elements.
				elems, newOff, err := parseElems(payload, offset, numSamples, sig)
				if err != nil {
					return nil, err
				}
				offset = newOff
				accumVals[sig.Name] = elems
			} else {
				// Fixed array (non-accumulated): one set of NumElements values.
				elems, newOff, err := parseElems(payload, offset, n, sig)
				if err != nil {
					return nil, err
				}
				offset = newOff
				fixedVals[sig.Name] = elems
			}
		}

		// Build one DataSample per slot.
		samples := make([]DataSample, numSamples)
		for k := 0; k < numSamples; k++ {
			vals := make(map[string][]float64, len(sigs))
			for sigName, av := range accumVals {
				vals[sigName] = []float64{av[k]}
			}
			for sigName, fv := range fixedVals {
				vals[sigName] = fv // shared read-only reference; hub does not modify
			}
			samples[k] = DataSample{HRTTimestamp: hrt, WallTime: arrivalTime, Values: vals}
		}
		return samples, nil
	}

	// Strict / Decimate: single snapshot, one element set per signal.
	vals := make(map[string][]float64, len(sigs))
	for _, sig := range sigs {
		n := sig.NumElements()
		elems, newOff, err := parseElems(payload, offset, n, sig)
		if err != nil {
			return nil, err
		}
		offset = newOff
		vals[sig.Name] = elems
	}
	return []DataSample{{HRTTimestamp: hrt, WallTime: arrivalTime, Values: vals}}, nil
}