package main

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
)

// ─── 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
	}
	return r * c
}

// 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.
func ParseConfig(payload []byte) ([]SignalInfo, error) {
	if len(payload) < 4 {
		return nil, fmt.Errorf("config payload too short")
	}
	numSigs := binary.LittleEndian.Uint32(payload[0:4])
	offset := 4
	sigs := make([]SignalInfo, 0, numSigs)
	for i := uint32(0); i < numSigs; i++ {
		if offset+SigDescSize > len(payload) {
			return nil, 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
	}
	return sigs, 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
}

// ParseData decodes a fully-reassembled DATA payload using the provided signal config.
// arrivalTime is the wall-clock time at which the packet was received.
func ParseData(payload []byte, sigs []SignalInfo, arrivalTime time.Time) (DataSample, error) {
	if len(payload) < 8 {
		return DataSample{}, fmt.Errorf("data payload too short")
	}
	hrt := binary.LittleEndian.Uint64(payload[0:8])
	offset := 8

	vals := make(map[string][]float64, len(sigs))
	for _, sig := range sigs {
		n := sig.NumElements()
		elems := make([]float64, n)

		if sig.QuantType == QuantNone {
			sz := rawTypeSize(sig.TypeCode)
			needed := n * sz
			if offset+needed > len(payload) {
				return DataSample{}, 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 DataSample{}, 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
		}
		vals[sig.Name] = elems
	}
	return DataSample{HRTTimestamp: hrt, WallTime: arrivalTime, Values: vals}, nil
}