package main

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ─── WebSocket client ─────────────────────────────────────────────────────────

type wsClient struct {
	hub  *Hub
	conn *websocket.Conn
	send chan []byte
}

func (c *wsClient) writePump() {
	pingTicker := time.NewTicker(30 * time.Second)
	defer func() {
		pingTicker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case msg, ok := <-c.send:
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
				return
			}
		case <-pingTicker.C:
			if err := c.conn.WriteControl(websocket.PingMessage, []byte{},
				time.Now().Add(10*time.Second)); err != nil {
				return
			}
		}
	}
}

func (c *wsClient) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(64 * 1024)
	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})
	for {
		_, msg, err := c.conn.ReadMessage()
		if err != nil {
			break
		}
		// Handle client messages (ping, setWindow, etc.) – currently just log.
		var env map[string]interface{}
		if json.Unmarshal(msg, &env) == nil {
			if t, ok := env["type"].(string); ok {
				switch t {
				case "ping":
					resp, _ := json.Marshal(map[string]string{"type": "pong"})
					select {
					case c.send <- resp:
					default:
					}
				}
			}
		}
		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	}
}

// ─── Hub ─────────────────────────────────────────────────────────────────────

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 64 * 1024,
	CheckOrigin:    func(r *http.Request) bool { return true },
}

// Hub is the central data broker between the UDP client and WebSocket clients.
type Hub struct {
	mu       sync.RWMutex
	signals  []SignalInfo
	configJS []byte // cached JSON config message
	configSeq uint64 // incremented on every UpdateConfig call

	clients     map[*wsClient]bool
	register    chan *wsClient
	unregister  chan *wsClient
	broadcastCh chan []byte // all sends go through Run() to avoid races on c.send

	dataCh chan DataSample // incoming samples from UDP goroutine

	// Time-signal calibration: only accessed from the Run() goroutine.
	// For temporal-array signals (TimeMode=FirstSample/LastSample), the
	// MARTe2 Time signal value (uint32 microseconds from start) is used as
	// the timing anchor.  We calibrate a per-signal wall-clock offset once
	// on the first data packet so that subsequent packets are stamped purely
	// from the embedded timer value → perfect continuity, no jitter gaps.
	timeSigCalib    map[string]float64 // key=time-signal name, value=wallTime-timerSecs offset
	configSeqAtCalib uint64             // configSeq value when timeSigCalib was last reset
}

// NewHub creates an initialised Hub.
func NewHub() *Hub {
	return &Hub{
		clients:      make(map[*wsClient]bool),
		register:     make(chan *wsClient, 8),
		unregister:   make(chan *wsClient, 8),
		broadcastCh:  make(chan []byte, 64),
		dataCh:       make(chan DataSample, 256),
		timeSigCalib: make(map[string]float64),
	}
}

// UpdateConfig stores a new signal config and broadcasts it to all WS clients.
func (h *Hub) UpdateConfig(sigs []SignalInfo) {
	msg, err := json.Marshal(map[string]interface{}{
		"type":    "config",
		"signals": sigs,
	})
	if err != nil {
		log.Printf("hub: marshal config: %v", err)
		return
	}
	h.mu.Lock()
	h.signals = sigs
	h.configJS = msg
	h.configSeq++
	h.mu.Unlock()

	h.broadcast(msg)
}

// PushData enqueues a data sample for broadcasting to WebSocket clients.
func (h *Hub) PushData(s DataSample) {
	select {
	case h.dataCh <- s:
	default:
		// Drop if buffer full to avoid blocking the UDP goroutine.
	}
}

// broadcast enqueues a message for delivery to all WebSocket clients.
// All actual sends happen inside Run() to avoid concurrent access to c.send.
func (h *Hub) broadcast(msg []byte) {
	select {
	case h.broadcastCh <- msg:
	default:
		// Drop if the broadcast queue is full.
	}
}

// HandleWebSocket upgrades an HTTP request to a WebSocket connection.
func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("ws upgrade: %v", err)
		return
	}
	c := &wsClient{
		hub:  h,
		conn: conn,
		send: make(chan []byte, 64),
	}
	h.register <- c

	go c.writePump()
	go c.readPump()
}

// Run is the hub's main goroutine. It must be started with go hub.Run().
func (h *Hub) Run() {
	// Batch data at ≤30 Hz.
	ticker := time.NewTicker(time.Second / 30)
	defer ticker.Stop()

	// Accumulate samples between ticks.
	pending := make([]DataSample, 0, 64)

	for {
		select {
		case c := <-h.register:
			h.mu.Lock()
			h.clients[c] = true
			cfg := h.configJS
			h.mu.Unlock()
			// Send current config immediately if we have one.
			if cfg != nil {
				select {
				case c.send <- cfg:
				default:
				}
			}

		case c := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[c]; ok {
				delete(h.clients, c)
				close(c.send)
			}
			h.mu.Unlock()

		case msg := <-h.broadcastCh:
			h.mu.RLock()
			for c := range h.clients {
				select {
				case c.send <- msg:
				default:
				}
			}
			h.mu.RUnlock()

		case s := <-h.dataCh:
			pending = append(pending, s)

		case <-ticker.C:
			if len(pending) == 0 {
				continue
			}
			h.mu.RLock()
			sigs := h.signals
			noClients := len(h.clients) == 0
			h.mu.RUnlock()

			if noClients || len(sigs) == 0 {
				pending = pending[:0]
				continue
			}

			msg := h.buildDataMessage(pending, sigs)
			pending = pending[:0]
			if msg != nil {
				h.broadcast(msg)
			}
		}
	}
}

// maxBatchPoints is the maximum number of points sent per signal per 30 Hz display tick.
// For temporal-array signals (e.g. 1 MSps packed as 1000 samples/packet), the expanded
// sample stream is decimated to this limit before transmission to WebSocket clients.
const maxBatchPoints = 2000

// sigData carries one signal's worth of time+value pairs in a single batch message.
type sigData struct {
	T []float64 `json:"t"` // unix seconds
	V []float64 `json:"v"` // physical values
}

// dataMsg is the JSON envelope for batched data sent to WebSocket clients.
// Each signal carries its own time axis so that temporal arrays (packed sample
// bursts with a known sampling rate) and scalar signals can coexist cleanly.
type dataMsg struct {
	Type    string             `json:"type"`
	Signals map[string]sigData `json:"signals"`
}

// buildDataMessage merges a batch of DataSamples into one JSON message.
//
// Three cases are handled:
//
//  1. Temporal array (NumElements > 1, TimeMode == FirstSample or LastSample):
//     The N samples in each packet represent a contiguous time burst at SamplingRate.
//     The embedded TimeSignal value (uint32 microseconds) is used as the timing
//     anchor for each burst.  A wall-clock offset is calibrated once on the first
//     packet so that abs-time stays consistent with other signals.
//     The full expanded stream is decimated to maxBatchPoints if needed.
//
//  2. Scalar signal (NumElements == 1):
//     One {t, v} pair per packet – wall arrival time used as timestamp.
//
//  3. Spatial / PacketTime array (NumElements > 1, TimeMode == 0):
//     Each element is tracked as a separate stream keyed "sig[i]", with wall
//     arrival time as the shared timestamp.
func (h *Hub) buildDataMessage(batch []DataSample, sigs []SignalInfo) []byte {
	if len(batch) == 0 {
		return nil
	}

	// Reset time-signal calibration whenever the config has changed.
	h.mu.RLock()
	seq := h.configSeq
	h.mu.RUnlock()
	if seq != h.configSeqAtCalib {
		h.configSeqAtCalib = seq
		h.timeSigCalib = make(map[string]float64)
	}

	out := make(map[string]sigData, len(sigs)*2)

	for _, sig := range sigs {
		n := sig.NumElements()
		isTemporal := n > 1 && (sig.TimeMode == TimeModeFirstSample || sig.TimeMode == TimeModeLastSample)

		switch {
		case isTemporal:
			// Resolve time signal (scalar that gives the anchor time in microseconds).
			hasTimeSig := sig.TimeSignalIdx != NoTimeSignal && int(sig.TimeSignalIdx) < len(sigs)
			var timeSigName string
			if hasTimeSig {
				timeSigName = sigs[sig.TimeSignalIdx].Name
			}

			dt := 0.0
			if sig.SamplingRate > 0 {
				dt = 1.0 / sig.SamplingRate
			}

			// Expand each packet's N samples into individual time-stamped points.
			allT := make([]float64, 0, len(batch)*n)
			allV := make([]float64, 0, len(batch)*n)

			for _, s := range batch {
				vals, ok := s.Values[sig.Name]
				if !ok || len(vals) < n {
					continue
				}

				// Compute the anchor timestamp for this burst.
				// Prefer the embedded time-signal value (microseconds) so that
				// consecutive bursts are perfectly contiguous regardless of jitter.
				var anchorTime float64
				anchorIsFirstSample := (sig.TimeMode == TimeModeFirstSample)

				if hasTimeSig {
					tVals, tOk := s.Values[timeSigName]
					if tOk && len(tVals) >= 1 {
						// Time signal is uint32 microseconds from system start.
						timerS := tVals[0] * 1e-6
						wallT := float64(s.WallTime.UnixNano()) / 1e9

						// Calibrate the wall-clock offset once per session so that
						// anchor times can be expressed as absolute Unix timestamps.
						if _, exists := h.timeSigCalib[timeSigName]; !exists {
							h.timeSigCalib[timeSigName] = wallT - timerS
						}
						anchorTime = h.timeSigCalib[timeSigName] + timerS
					} else {
						// Time signal missing in this packet – fall back to wall clock.
						anchorTime = float64(s.WallTime.UnixNano()) / 1e9
						anchorIsFirstSample = false // wallT = last sample
					}
				} else {
					// No time signal configured – use wall arrival as last-sample anchor.
					anchorTime = float64(s.WallTime.UnixNano()) / 1e9
					anchorIsFirstSample = false
				}

				for k := 0; k < n; k++ {
					var t float64
					if anchorIsFirstSample {
						t = anchorTime + float64(k)*dt
					} else {
						t = anchorTime - float64(n-1-k)*dt
					}
					allT = append(allT, t)
					allV = append(allV, vals[k])
				}
			}

			// Decimate to maxBatchPoints if necessary.
			step := 1
			if len(allT) > maxBatchPoints {
				step = len(allT) / maxBatchPoints
				if step < 1 {
					step = 1
				}
			}
			decimT := make([]float64, 0, len(allT)/step+1)
			decimV := make([]float64, 0, len(allV)/step+1)
			for i := 0; i < len(allT); i += step {
				decimT = append(decimT, allT[i])
				decimV = append(decimV, allV[i])
			}
			out[sig.Name] = sigData{T: decimT, V: decimV}

		case n == 1:
			// Scalar signal: one sample per packet.
			ts := make([]float64, 0, len(batch))
			vs := make([]float64, 0, len(batch))
			for _, s := range batch {
				vals, ok := s.Values[sig.Name]
				if !ok || len(vals) < 1 {
					continue
				}
				ts = append(ts, float64(s.WallTime.UnixNano())/1e9)
				vs = append(vs, vals[0])
			}
			out[sig.Name] = sigData{T: ts, V: vs}

		default:
			// Spatial / PacketTime array: one stream per element, keyed "sig[i]".
			for i := 0; i < n; i++ {
				key := arrayKey(sig.Name, i)
				ts := make([]float64, 0, len(batch))
				vs := make([]float64, 0, len(batch))
				for _, s := range batch {
					vals, ok := s.Values[sig.Name]
					if !ok || len(vals) <= i {
						continue
					}
					ts = append(ts, float64(s.WallTime.UnixNano())/1e9)
					vs = append(vs, vals[i])
				}
				out[key] = sigData{T: ts, V: vs}
			}
		}
	}

	result, err := json.Marshal(dataMsg{
		Type:    "data",
		Signals: out,
	})
	if err != nil {
		log.Printf("hub: marshal data: %v", err)
		return nil
	}
	return result
}

// arrayKey returns the buffer key for element i of an array signal.
func arrayKey(name string, i int) string {
	return name + "[" + itoa(i) + "]"
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	buf := [20]byte{}
	pos := len(buf)
	for n > 0 {
		pos--
		buf[pos] = byte('0' + n%10)
		n /= 10
	}
	return string(buf[pos:])
}