// Command streamhub-e2e is an end-to-end test client for the C++ StreamHub.
//
// It connects to a running StreamHub WebSocket endpoint (with at least one
// connected UDPStreamer source, e.g. the stack launched by run_e2e_test.sh)
// and verifies the full protocol:
//
//  1. "sources" event with at least one connected source
//  2. "config" event per source with at least one signal
//  3. binary v1 data pushes: parseable, per-signal monotonic time,
//     timestamps within a few seconds of wall clock (Unix time base)
//  4. "stats" event with a positive receive rate
//  5. WS zoom round-trip: reqId echoed, points returned in [t0,t1]
//  6. hub-side trigger: setTrigger+arm → triggerState(armed) → binary v2
//     capture frame with the latched pre/post window
//
// Exit code 0 on success; 1 with a FAIL message otherwise.
package main

import (
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math"
	"os"
	"time"

	"github.com/gorilla/websocket"
)

var hub = flag.String("hub", "127.0.0.1:8090", "StreamHub host:port")
var timeout = flag.Duration("timeout", 30*time.Second, "overall test timeout")
var verbose = flag.Bool("v", false, "log every received event")

// ---------------------------------------------------------------------------
// Wire types (subset of the StreamHub JSON protocol)
// ---------------------------------------------------------------------------

type sourceInfo struct {
	ID    string `json:"id"`
	Label string `json:"label"`
	Addr  string `json:"addr"`
	State string `json:"state"`
}

type signalInfo struct {
	Name     string  `json:"name"`
	TypeCode uint32  `json:"typeCode"`
	NumRows  uint32  `json:"numRows"`
	NumCols  uint32  `json:"numCols"`
	TimeMode int     `json:"timeMode"`
	Rate     float64 `json:"samplingRate"`
}

type statInfo struct {
	State         string  `json:"state"`
	TotalReceived uint64  `json:"totalReceived"`
	RateHz        float64 `json:"rateHz"`
	CycleHist     []f64   `json:"cycleHist"`
}

type f64 = float64

type zoomPoints struct {
	T []float64 `json:"t"`
	V []float64 `json:"v"`
}

type historyInfoMsg struct {
	Enabled       bool               `json:"enabled"`
	DurationHours float64            `json:"durationHours"`
	Decimation    uint32             `json:"decimation"`
	Signals       map[string]struct {
		T0       float64 `json:"t0"`
		T1       float64 `json:"t1"`
		Count    uint32  `json:"count"`
		Capacity uint32  `json:"capacity"`
	} `json:"signals"`
}

type event struct {
	Type     string                `json:"type"`
	Sources  json.RawMessage       `json:"sources"`
	SourceID string                `json:"sourceId"`
	Signals  json.RawMessage       `json:"signals"`
	ReqID    uint32                `json:"reqId"`
	State    string                `json:"state"`
	TrigTime float64               `json:"trigTime"`
}

// Parsed binary v1 push frame: sourceId → signal → samples.
type pushFrame struct {
	sourceID string
	signals  map[string]zoomPoints
}

// Parsed binary v2 capture frame.
type captureFrame struct {
	trigTime, preSec, postSec float64
	signals                   map[string]zoomPoints
}

// ---------------------------------------------------------------------------
// Binary parsers
// ---------------------------------------------------------------------------

func parsePush(b []byte) (*pushFrame, error) {
	if len(b) < 2 || b[0] != 1 {
		return nil, fmt.Errorf("not a v1 frame")
	}
	idLen := int(b[1])
	off := 2
	if len(b) < off+idLen+4 {
		return nil, fmt.Errorf("truncated header")
	}
	f := &pushFrame{sourceID: string(b[off : off+idLen]),
		signals: map[string]zoomPoints{}}
	off += idLen
	nSig := int(binary.LittleEndian.Uint32(b[off:]))
	off += 4
	for s := 0; s < nSig; s++ {
		if len(b) < off+2 {
			return nil, fmt.Errorf("truncated keyLen (sig %d)", s)
		}
		keyLen := int(binary.LittleEndian.Uint16(b[off:]))
		off += 2
		if len(b) < off+keyLen+4 {
			return nil, fmt.Errorf("truncated key (sig %d)", s)
		}
		key := string(b[off : off+keyLen])
		off += keyLen
		n := int(binary.LittleEndian.Uint32(b[off:]))
		off += 4
		if len(b) < off+16*n {
			return nil, fmt.Errorf("truncated data (sig %s n=%d)", key, n)
		}
		pts := zoomPoints{T: make([]float64, n), V: make([]float64, n)}
		for i := 0; i < n; i++ {
			pts.T[i] = math.Float64frombits(binary.LittleEndian.Uint64(b[off+8*i:]))
		}
		off += 8 * n
		for i := 0; i < n; i++ {
			pts.V[i] = math.Float64frombits(binary.LittleEndian.Uint64(b[off+8*i:]))
		}
		off += 8 * n
		f.signals[key] = pts
	}
	return f, nil
}

func parseCapture(b []byte) (*captureFrame, error) {
	if len(b) < 1+24+4 || b[0] != 2 {
		return nil, fmt.Errorf("not a v2 frame")
	}
	rdF64 := func(off int) float64 {
		return math.Float64frombits(binary.LittleEndian.Uint64(b[off:]))
	}
	f := &captureFrame{
		trigTime: rdF64(1), preSec: rdF64(9), postSec: rdF64(17),
		signals: map[string]zoomPoints{},
	}
	off := 25
	nSig := int(binary.LittleEndian.Uint32(b[off:]))
	off += 4
	for s := 0; s < nSig; s++ {
		keyLen := int(binary.LittleEndian.Uint16(b[off:]))
		off += 2
		key := string(b[off : off+keyLen])
		off += keyLen
		n := int(binary.LittleEndian.Uint32(b[off:]))
		off += 4
		if len(b) < off+16*n {
			return nil, fmt.Errorf("truncated capture (sig %s n=%d)", key, n)
		}
		pts := zoomPoints{T: make([]float64, n), V: make([]float64, n)}
		for i := 0; i < n; i++ {
			pts.T[i] = math.Float64frombits(binary.LittleEndian.Uint64(b[off+8*i:]))
		}
		off += 8 * n
		for i := 0; i < n; i++ {
			pts.V[i] = math.Float64frombits(binary.LittleEndian.Uint64(b[off+8*i:]))
		}
		off += 8 * n
		f.signals[key] = pts
	}
	return f, nil
}

// ---------------------------------------------------------------------------
// Test driver
// ---------------------------------------------------------------------------

type client struct {
	ws       *websocket.Conn
	deadline time.Time

	sources  []sourceInfo
	configs  map[string][]signalInfo // sourceId → signals
	pushes   []*pushFrame
	stats    map[string]statInfo
	zooms        map[uint32]map[string]zoomPoints
	histZooms    map[uint32]map[string]zoomPoints
	historyInfo  *historyInfoMsg
	trigSt       []string // observed triggerState sequence
	captures     []*captureFrame
}

func (c *client) send(v interface{}) {
	b, _ := json.Marshal(v)
	if err := c.ws.WriteMessage(websocket.TextMessage, b); err != nil {
		fail("ws write: %v", err)
	}
}

// pump reads one WS message (with a short read deadline) and dispatches it.
func (c *client) pump() {
	c.ws.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	mt, data, err := c.ws.ReadMessage()
	if err != nil {
		if websocket.IsUnexpectedCloseError(err) {
			fail("ws closed: %v", err)
		}
		return // read timeout — fine
	}
	switch mt {
	case websocket.BinaryMessage:
		if len(data) == 0 {
			return
		}
		switch data[0] {
		case 1:
			if f, err := parsePush(data); err == nil {
				c.pushes = append(c.pushes, f)
			} else {
				fail("bad v1 frame: %v", err)
			}
		case 2:
			if f, err := parseCapture(data); err == nil {
				c.captures = append(c.captures, f)
			} else {
				fail("bad v2 frame: %v", err)
			}
		default:
			fail("unknown binary frame version %d", data[0])
		}
	case websocket.TextMessage:
		var ev event
		if err := json.Unmarshal(data, &ev); err != nil {
			fail("bad JSON event: %v (%.120s)", err, data)
		}
		if *verbose {
			log.Printf("event %-12s %.160s", ev.Type, data)
		}
		switch ev.Type {
		case "sources":
			var srcs []sourceInfo
			if err := json.Unmarshal(ev.Sources, &srcs); err == nil {
				c.sources = srcs
			}
		case "config":
			var sigs []signalInfo
			if err := json.Unmarshal(ev.Signals, &sigs); err == nil {
				c.configs[ev.SourceID] = sigs
			} else {
				log.Printf("config parse error: %v (%.200s)", err, data)
			}
		case "stats":
			var st map[string]statInfo
			if err := json.Unmarshal(ev.Sources, &st); err == nil {
				c.stats = st
			}
		case "zoom":
			var body struct {
				Signals map[string]zoomPoints `json:"signals"`
			}
			if err := json.Unmarshal(data, &body); err == nil {
				c.zooms[ev.ReqID] = body.Signals
			}
		case "historyZoom":
			var body struct {
				Signals map[string]zoomPoints `json:"signals"`
			}
			if err := json.Unmarshal(data, &body); err == nil {
				c.histZooms[ev.ReqID] = body.Signals
			}
		case "historyInfo":
			var hi historyInfoMsg
			if err := json.Unmarshal(data, &hi); err == nil {
				c.historyInfo = &hi
			}
		case "triggerState":
			c.trigSt = append(c.trigSt, ev.State)
		}
	}
}

// waitFor pumps messages until cond() or the step deadline expires.
func (c *client) waitFor(what string, d time.Duration, cond func() bool) {
	end := time.Now().Add(d)
	if end.After(c.deadline) {
		end = c.deadline
	}
	for time.Now().Before(end) {
		if cond() {
			log.Printf("OK   %s", what)
			return
		}
		c.pump()
	}
	fail("timeout waiting for %s", what)
}

func fail(format string, args ...interface{}) {
	fmt.Printf("FAIL "+format+"\n", args...)
	os.Exit(1)
}

func main() {
	flag.Parse()
	url := "ws://" + *hub + "/ws"
	log.Printf("connecting to %s", url)

	ws, _, err := websocket.DefaultDialer.Dial(url, nil)
	if err != nil {
		fail("dial %s: %v", url, err)
	}
	defer ws.Close()

	c := &client{
		ws:        ws,
		deadline:  time.Now().Add(*timeout),
		configs:   map[string][]signalInfo{},
		zooms:     map[uint32]map[string]zoomPoints{},
		histZooms: map[uint32]map[string]zoomPoints{},
	}

	// ── 1. sources ────────────────────────────────────────────────────────
	c.send(map[string]interface{}{"type": "getSources"})
	c.waitFor("sources event with a connected source", 10*time.Second, func() bool {
		for _, s := range c.sources {
			if s.State == "connected" {
				return true
			}
		}
		return false
	})

	// ── 2. config per connected source ───────────────────────────────────
	for _, s := range c.sources {
		log.Printf("source %s (%s): state=%s", s.ID, s.Label, s.State)
		c.send(map[string]interface{}{"type": "getConfig", "sourceId": s.ID})
	}
	c.waitFor("config with signals for every connected source", 10*time.Second, func() bool {
		for _, s := range c.sources {
			if s.State != "connected" {
				continue
			}
			if len(c.configs[s.ID]) == 0 {
				return false
			}
		}
		return len(c.sources) > 0
	})

	// ── 3. binary pushes: wall-clock time base + monotonicity ────────────
	c.waitFor("binary v1 data pushes (>=10 frames)", 10*time.Second, func() bool {
		return len(c.pushes) >= 10
	})
	now := float64(time.Now().UnixNano()) / 1e9
	seen := map[string][]float64{} // last times per src:sig
	for _, f := range c.pushes {
		for key, pts := range f.signals {
			full := f.sourceID + ":" + key
			for i, t := range pts.T {
				if math.Abs(t-now) > 30.0 {
					fail("timestamp not wall-clock: %s t=%.3f now=%.3f", full, t, now)
				}
				prev := seen[full]
				if len(prev) > 0 && t < prev[len(prev)-1]-1e-9 {
					fail("non-monotonic time on %s: %.9f after %.9f (i=%d)",
						full, t, prev[len(prev)-1], i)
				}
				seen[full] = append(seen[full], t)
			}
		}
	}
	if len(seen) == 0 {
		fail("pushes contained no signal data")
	}
	log.Printf("OK   wall-clock & monotonic time on %d signal streams", len(seen))

	// ── 4. stats ──────────────────────────────────────────────────────────
	c.send(map[string]interface{}{"type": "getStats"})
	c.waitFor("stats with positive rate", 10*time.Second, func() bool {
		for _, st := range c.stats {
			if st.State == "connected" && st.RateHz > 0 && st.TotalReceived > 0 {
				return true
			}
		}
		return false
	})

	// ── 5. zoom round-trip ───────────────────────────────────────────────
	// Use the busiest streamed signal and the time range we actually saw.
	var zoomKey string
	var zMax int
	for k, ts := range seen {
		if len(ts) > zMax {
			zMax, zoomKey = len(ts), k
		}
	}
	ts := seen[zoomKey]
	t1 := ts[len(ts)-1]
	t0 := t1 - 0.5
	const reqID = 4242
	c.send(map[string]interface{}{
		"type": "zoom", "reqId": reqID, "t0": t0, "t1": t1, "n": 200,
		"signals": zoomKey,
	})
	c.waitFor(fmt.Sprintf("zoom reply (reqId=%d, %s)", reqID, zoomKey),
		10*time.Second, func() bool {
			sigs, ok := c.zooms[reqID]
			if !ok {
				return false
			}
			pts, ok := sigs[zoomKey]
			if !ok || len(pts.T) < 2 {
				fail("zoom reply missing %s (got %d signals)", zoomKey, len(sigs))
			}
			for _, t := range pts.T {
				if t < t0-1e-6 || t > t1+1e-6 {
					fail("zoom point outside range: t=%.9f not in [%.9f,%.9f]", t, t0, t1)
				}
			}
			return true
		})

	// ── 5b. historyInfo — check the hub broadcast it on connect ─────────
	if c.historyInfo != nil && c.historyInfo.Enabled {
		log.Printf("OK   historyInfo: enabled, %.1fh, decimation=%d, %d signals",
			c.historyInfo.DurationHours, c.historyInfo.Decimation,
			len(c.historyInfo.Signals))

		// ── 5c. historyZoom round-trip ──────────────────────────────────
		const hReqID = 4243
		c.send(map[string]interface{}{
			"type": "historyZoom", "reqId": hReqID,
			"t0": t0, "t1": t1, "n": 200,
			"signals": zoomKey,
		})
		c.waitFor(fmt.Sprintf("historyZoom reply (reqId=%d, %s)", hReqID, zoomKey),
			10*time.Second, func() bool {
				sigs, ok := c.histZooms[hReqID]
				if !ok {
					return false
				}
				pts, ok := sigs[zoomKey]
				if !ok || len(pts.T) < 1 {
					// History data may still be sparse right after startup
					return true
				}
				for _, ht := range pts.T {
					if ht < t0-1e-6 || ht > t1+1e-6 {
						fail("historyZoom point outside range: t=%.9f not in [%.9f,%.9f]", ht, t0, t1)
					}
				}
				return true
			})
	} else {
		log.Println("     (history not enabled — skipping historyZoom test)")
	}

	// ── 6. trigger: arm → capture ────────────────────────────────────────
	// Trigger on an *oscillating* signal at its mean observed value: a
	// monotonic ramp (counter, time array) crosses its past mean only once,
	// before arming, so a rising edge would never fire on it.  Pick the
	// busiest signal whose last push frame is non-monotonic (a sine).
	lastVals := map[string][]float64{}
	for _, f := range c.pushes {
		for name, pts := range f.signals {
			if len(pts.V) >= 4 {
				lastVals[f.sourceID+":"+name] = pts.V
			}
		}
	}
	trigKey := ""
	tMaxPts := 0
	for k, vs := range lastVals {
		monotonic := true
		for i := 1; i < len(vs); i++ {
			if vs[i] < vs[i-1] {
				monotonic = false
				break
			}
		}
		if !monotonic && len(seen[k]) > tMaxPts {
			tMaxPts, trigKey = len(seen[k]), k
		}
	}
	if trigKey == "" {
		fail("no oscillating signal found for trigger test")
	}
	vals := lastVals[trigKey]
	mean := 0.0
	for _, v := range vals {
		mean += v
	}
	mean /= float64(len(vals))
	log.Printf("     trigger signal %s, threshold %.6g", trigKey, mean)

	c.send(map[string]interface{}{
		"type": "setTrigger", "signal": trigKey, "edge": "rising",
		"threshold": mean, "windowSec": 0.1, "prePercent": 20.0,
		"mode": "single",
	})
	c.send(map[string]interface{}{"type": "arm"})
	// The trigger can fire within microseconds of arming (5 MS/s sine), so
	// the broadcast emitted by the arm command may already say "collecting"
	// or even "triggered" — any of these proves the arm was accepted.
	c.waitFor("triggerState: armed/collecting/triggered", 5*time.Second, func() bool {
		for _, s := range c.trigSt {
			if s == "armed" || s == "collecting" || s == "triggered" {
				return true
			}
		}
		return false
	})
	c.waitFor("binary v2 capture frame", 15*time.Second, func() bool {
		return len(c.captures) > 0
	})

	cap0 := c.captures[0]
	if math.Abs(cap0.preSec-0.02) > 1e-9 || math.Abs(cap0.postSec-0.08) > 1e-9 {
		fail("capture window mismatch: pre=%.6f post=%.6f (want 0.02/0.08)",
			cap0.preSec, cap0.postSec)
	}
	pts, ok := cap0.signals[trigKey]
	if !ok || len(pts.T) == 0 {
		fail("capture missing trigger signal %s (%d signals)", trigKey, len(cap0.signals))
	}
	for _, t := range pts.T {
		if t < cap0.trigTime-cap0.preSec-1e-3 || t > cap0.trigTime+cap0.postSec+1e-3 {
			fail("capture point outside window: t=%.9f trig=%.9f", t, cap0.trigTime)
		}
	}
	log.Printf("OK   capture: trig=%.6f pre=%.3fs post=%.3fs %d signals",
		cap0.trigTime, cap0.preSec, cap0.postSec, len(cap0.signals))

	c.waitFor("triggerState: triggered", 5*time.Second, func() bool {
		for _, s := range c.trigSt {
			if s == "triggered" {
				return true
			}
		}
		return false
	})
	c.send(map[string]interface{}{"type": "disarm"})

	fmt.Println("PASS streamhub-e2e: all checks passed")
}