// Command chain-client is the authoritative mock StreamHub client for the
// streaming-chain E2E suite. It connects to a running StreamHub, records the
// live binary stream to disk (received_<id>.bin), and runs behavioural checks
// (live / zoom / window / trigger), writing the results to checks_<id>.json.
//
// Unlike the streamhub smoke test, individual check failures are *recorded*
// (not fatal): only connection/protocol corruption exits non-zero, so one
// scenario's quirk never aborts the matrix. The waveform validator and report
// decide overall pass/fail from checks_<id>.json + the recorded stream.
//
// received_<id>.bin format ("RCV1"):
//
//	magic[4]="RCV1"; [u32 nSig]; per signal:
//	  [u16 keyLen][key][u32 N][N×f64 t][N×f64 v]
//
// where each signal's samples are merged across all v1 pushes, sorted by time
// and de-duplicated (same timestamp kept once).
package main

import (
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

var (
	hubFlag   = flag.String("hub", "127.0.0.1:8090", "StreamHub host:port")
	scenario  = flag.String("scenario", "", "scenario id (artifact basename)")
	trigsig   = flag.String("trigsig", "", "trigger signal key src:sig (empty = skip)")
	trigthr   = flag.Float64("trigthr", math.NaN(), "trigger threshold (NaN = use mean)")
	checksCSV = flag.String("checks", "live,zoom,window,trigger", "checks to run")
	outDir    = flag.String("out", "/tmp/chain_e2e", "artifact output dir")
	durSec    = flag.Float64("dur", 4.0, "live recording duration (s)")
	timeout   = flag.Duration("timeout", 90*time.Second, "overall timeout")
	verbose   = flag.Bool("v", false, "log every event")
)

// ── wire types ────────────────────────────────────────────────────────────

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 points struct {
	T []float64 `json:"t"`
	V []float64 `json:"v"`
}

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"`
}

type pushFrame struct {
	sourceID string
	signals  map[string]points
}

type captureFrame struct {
	trigTime, preSec, postSec float64
	signals                   map[string]points
}

// ── binary parsers (from Test/E2E/streamhub/main.go) ────────────────────────

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]points{}}
	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 := points{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")
	}
	rd := func(off int) float64 { return math.Float64frombits(binary.LittleEndian.Uint64(b[off:])) }
	f := &captureFrame{trigTime: rd(1), preSec: rd(9), postSec: rd(17), signals: map[string]points{}}
	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 := points{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
}

// ── client ──────────────────────────────────────────────────────────────────

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

	sources  []sourceInfo
	configs  map[string][]signalInfo
	pushes   []*pushFrame
	zooms    map[uint32]map[string]points
	trigSt   []string
	captures []*captureFrame
}

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

func (c *client) pump() {
	c.ws.SetReadDeadline(time.Now().Add(300 * time.Millisecond))
	mt, data, err := c.ws.ReadMessage()
	if err != nil {
		if websocket.IsUnexpectedCloseError(err) {
			fatal("ws closed: %v", err)
		}
		return
	}
	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 {
				fatal("bad v1 frame: %v", err)
			}
		case 2:
			if f, err := parseCapture(data); err == nil {
				c.captures = append(c.captures, f)
			} else {
				fatal("bad v2 frame: %v", err)
			}
		default:
			fatal("unknown binary frame version %d", data[0])
		}
	case websocket.TextMessage:
		var ev event
		if err := json.Unmarshal(data, &ev); err != nil {
			return
		}
		if *verbose {
			log.Printf("event %-12s %.140s", ev.Type, data)
		}
		switch ev.Type {
		case "sources":
			var s []sourceInfo
			if json.Unmarshal(ev.Sources, &s) == nil {
				c.sources = s
			}
		case "config":
			var s []signalInfo
			if json.Unmarshal(ev.Signals, &s) == nil {
				c.configs[ev.SourceID] = s
			}
		case "zoom":
			var body struct {
				Signals map[string]points `json:"signals"`
			}
			if json.Unmarshal(data, &body) == nil {
				c.zooms[ev.ReqID] = body.Signals
			}
		case "triggerState":
			c.trigSt = append(c.trigSt, ev.State)
		}
	}
}

func (c *client) waitFor(d time.Duration, cond func() bool) bool {
	end := time.Now().Add(d)
	if end.After(c.deadline) {
		end = c.deadline
	}
	for time.Now().Before(end) {
		if cond() {
			return true
		}
		c.pump()
	}
	return false
}

func fatal(format string, a ...interface{}) {
	fmt.Printf("FATAL "+format+"\n", a...)
	os.Exit(1)
}

// merged returns time-sorted, de-duplicated samples per "src:sig" key.
func (c *client) merged() map[string]points {
	tmp := map[string]points{}
	for _, f := range c.pushes {
		for k, p := range f.signals {
			full := f.sourceID + ":" + k
			cur := tmp[full]
			cur.T = append(cur.T, p.T...)
			cur.V = append(cur.V, p.V...)
			tmp[full] = cur
		}
	}
	out := map[string]points{}
	for k, p := range tmp {
		idx := make([]int, len(p.T))
		for i := range idx {
			idx[i] = i
		}
		sort.Slice(idx, func(a, b int) bool { return p.T[idx[a]] < p.T[idx[b]] })
		var op points
		last := math.NaN()
		for _, i := range idx {
			if !math.IsNaN(last) && p.T[i] == last {
				continue
			}
			op.T = append(op.T, p.T[i])
			op.V = append(op.V, p.V[i])
			last = p.T[i]
		}
		out[k] = op
	}
	return out
}

// ── received_<id>.bin writer ─────────────────────────────────────────────────

func writeReceived(path string, m map[string]points) error {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	var buf []byte
	put32 := func(v uint32) { var t [4]byte; binary.LittleEndian.PutUint32(t[:], v); buf = append(buf, t[:]...) }
	put16 := func(v uint16) { var t [2]byte; binary.LittleEndian.PutUint16(t[:], v); buf = append(buf, t[:]...) }
	putf := func(v float64) { var t [8]byte; binary.LittleEndian.PutUint64(t[:], math.Float64bits(v)); buf = append(buf, t[:]...) }
	buf = append(buf, []byte("RCV1")...)
	put32(uint32(len(keys)))
	for _, k := range keys {
		put16(uint16(len(k)))
		buf = append(buf, []byte(k)...)
		p := m[k]
		put32(uint32(len(p.T)))
		for _, t := range p.T {
			putf(t)
		}
		for _, v := range p.V {
			putf(v)
		}
	}
	_, err = f.Write(buf)
	return err
}

// ── checks structures ────────────────────────────────────────────────────────

type zoomCheck struct {
	Range    [2]float64 `json:"range"`
	N        int        `json:"n"`
	Returned int        `json:"returned"`
	InRange  bool       `json:"inrange"`
	Key      string     `json:"key"`
}

type windowCheck struct {
	WindowSec float64 `json:"windowSec"`
	Span      float64 `json:"span"`
	Returned  int     `json:"returned"`
	OK        bool    `json:"ok"`
	Key       string  `json:"key"`
}

type trigCheck struct {
	Edge       string  `json:"edge"`
	Mode       string  `json:"mode"`
	Fired      bool    `json:"fired"`
	TrigTime   float64 `json:"trigTime"`
	PreSec     float64 `json:"preSec"`
	PostSec    float64 `json:"postSec"`
	CapturePts int     `json:"capturePts"`
	EdgeOK     bool    `json:"edgeOk"`
	WindowOK   bool    `json:"windowOk"`
	Rearmed    bool    `json:"rearmed"`
	Key        string  `json:"key"`
}

type liveCheck struct {
	OK         bool    `json:"ok"`
	Frames     int     `json:"frames"`
	Signals    int     `json:"signals"`
	Monotonic  bool    `json:"monotonic"`
	WallClock  bool    `json:"wallclock"`
	DurationS  float64 `json:"duration"`
}

type checksOut struct {
	Scenario string        `json:"scenario"`
	Live     liveCheck     `json:"live"`
	Zoom     []zoomCheck   `json:"zoom"`
	Window   windowCheck   `json:"window"`
	Trigger  []trigCheck   `json:"trigger"`
}

func has(set, name string) bool {
	for _, s := range strings.Split(set, ",") {
		if strings.TrimSpace(s) == name {
			return true
		}
	}
	return false
}

func main() {
	flag.Parse()
	if *scenario == "" {
		fatal("missing -scenario")
	}
	if err := os.MkdirAll(*outDir, 0o755); err != nil {
		fatal("mkdir %s: %v", *outDir, err)
	}
	url := "ws://" + *hubFlag + "/ws"
	log.Printf("connecting to %s (scenario %s)", url, *scenario)
	ws, _, err := websocket.DefaultDialer.Dial(url, nil)
	if err != nil {
		fatal("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]points{},
	}
	out := checksOut{Scenario: *scenario}

	// 1. sources connected
	c.send(map[string]interface{}{"type": "getSources"})
	if !c.waitFor(20*time.Second, func() bool {
		for _, s := range c.sources {
			if s.State == "connected" {
				return true
			}
		}
		return false
	}) {
		fatal("no connected source within timeout")
	}
	for _, s := range c.sources {
		c.send(map[string]interface{}{"type": "getConfig", "sourceId": s.ID})
	}
	c.waitFor(5*time.Second, func() bool {
		for _, s := range c.sources {
			if s.State == "connected" && len(c.configs[s.ID]) == 0 {
				return false
			}
		}
		return true
	})

	// 2. live recording
	recEnd := time.Now().Add(time.Duration(*durSec * float64(time.Second)))
	for time.Now().Before(recEnd) {
		c.pump()
	}
	m := c.merged()
	recvPath := filepath.Join(*outDir, "received_"+*scenario+".bin")
	if err := writeReceived(recvPath, m); err != nil {
		fatal("write received: %v", err)
	}
	now := float64(time.Now().UnixNano()) / 1e9
	mono, wall := true, true
	for _, p := range m {
		for i, t := range p.T {
			if math.Abs(t-now) > 60.0 {
				wall = false
			}
			if i > 0 && t < p.T[i-1]-1e-9 {
				mono = false
			}
		}
	}
	out.Live = liveCheck{
		OK: len(c.pushes) >= 5 && len(m) > 0 && mono && wall,
		Frames: len(c.pushes), Signals: len(m), Monotonic: mono,
		WallClock: wall, DurationS: *durSec,
	}
	log.Printf("live: %d frames, %d signals, mono=%v wall=%v", len(c.pushes), len(m), mono, wall)

	// busiest signal for zoom/window
	var busy string
	bn := 0
	for k, p := range m {
		if len(p.T) > bn {
			bn, busy = len(p.T), k
		}
	}

	// 3. zoom: narrow + wide
	if has(*checksCSV, "zoom") && busy != "" {
		ts := m[busy].T
		t1 := ts[len(ts)-1]
		t0full := ts[0]
		ranges := [][2]float64{{t1 - 0.05, t1}, {t0full, t1}} // narrow, wide
		for ri, rg := range ranges {
			reqID := uint32(1000 + ri)
			c.send(map[string]interface{}{
				"type": "zoom", "reqId": reqID, "t0": rg[0], "t1": rg[1],
				"n": 300, "signals": busy,
			})
			ok := c.waitFor(8*time.Second, func() bool { _, ok := c.zooms[reqID]; return ok })
			zc := zoomCheck{Range: rg, N: 300, Key: busy, InRange: true}
			if ok {
				pts := c.zooms[reqID][busy]
				zc.Returned = len(pts.T)
				for _, t := range pts.T {
					if t < rg[0]-1e-6 || t > rg[1]+1e-6 {
						zc.InRange = false
					}
				}
			} else {
				zc.InRange = false
			}
			out.Zoom = append(out.Zoom, zc)
			log.Printf("zoom [%.4f,%.4f] returned=%d inrange=%v", rg[0], rg[1], zc.Returned, zc.InRange)
		}
	}

	// 4. window: zoom over [now-windowSec, now] and check span ≤ window
	if has(*checksCSV, "window") && busy != "" {
		ts := m[busy].T
		t1 := ts[len(ts)-1]
		const winSec = 1.0
		reqID := uint32(2000)
		c.send(map[string]interface{}{
			"type": "zoom", "reqId": reqID, "t0": t1 - winSec, "t1": t1,
			"n": 600, "signals": busy,
		})
		ok := c.waitFor(8*time.Second, func() bool { _, ok := c.zooms[reqID]; return ok })
		wc := windowCheck{WindowSec: winSec, Key: busy}
		if ok {
			pts := c.zooms[reqID][busy]
			wc.Returned = len(pts.T)
			if len(pts.T) >= 2 {
				wc.Span = pts.T[len(pts.T)-1] - pts.T[0]
				wc.OK = wc.Span <= winSec+1e-3
			}
		}
		out.Window = wc
		log.Printf("window %.2fs: span=%.4f returned=%d ok=%v", winSec, wc.Span, wc.Returned, wc.OK)
	}

	// 5. trigger matrix
	if has(*checksCSV, "trigger") && *trigsig != "" {
		thr := *trigthr
		if math.IsNaN(thr) {
			if p, ok := m[*trigsig]; ok && len(p.V) > 0 {
				s := 0.0
				for _, v := range p.V {
					s += v
				}
				thr = s / float64(len(p.V))
			} else {
				thr = 0.0
			}
		}
		log.Printf("trigger signal %s threshold %.6g", *trigsig, thr)
		for _, edge := range []string{"rising", "falling", "both"} {
			for _, mode := range []string{"normal", "single"} {
				out.Trigger = append(out.Trigger, c.runTrigger(*trigsig, edge, mode, thr))
			}
		}
	}

	// write checks json
	cj := filepath.Join(*outDir, "checks_"+*scenario+".json")
	b, _ := json.MarshalIndent(out, "", "  ")
	if err := os.WriteFile(cj, b, 0o644); err != nil {
		fatal("write checks: %v", err)
	}
	fmt.Printf("OK chain-client %s: %s + %s\n", *scenario, filepath.Base(recvPath), filepath.Base(cj))
}

// runTrigger configures one edge/mode trigger, arms it, and records the result.
func (c *client) runTrigger(key, edge, mode string, thr float64) trigCheck {
	tc := trigCheck{Edge: edge, Mode: mode, Key: key}
	beforeCaps := len(c.captures)
	c.send(map[string]interface{}{
		"type": "setTrigger", "signal": key, "edge": edge,
		"threshold": thr, "windowSec": 0.1, "prePercent": 20.0, "mode": mode,
	})
	c.send(map[string]interface{}{"type": "arm"})
	fired := c.waitFor(8*time.Second, func() bool { return len(c.captures) > beforeCaps })
	tc.Fired = fired
	if !fired {
		c.send(map[string]interface{}{"type": "disarm"})
		return tc
	}
	cap0 := c.captures[len(c.captures)-1]
	tc.TrigTime, tc.PreSec, tc.PostSec = cap0.trigTime, cap0.preSec, cap0.postSec
	tc.WindowOK = math.Abs(cap0.preSec-0.02) < 1e-6 && math.Abs(cap0.postSec-0.08) < 1e-6
	if pts, ok := cap0.signals[key]; ok {
		tc.CapturePts = len(pts.T)
		tc.EdgeOK = edgeCrosses(pts, cap0.trigTime, edge, thr)
	}
	// re-arm behaviour
	if mode == "normal" {
		tc.Rearmed = c.waitFor(4*time.Second, func() bool { return len(c.captures) > beforeCaps+1 })
	} else { // single: must NOT fire again until rearm
		again := c.waitFor(1500*time.Millisecond, func() bool { return len(c.captures) > beforeCaps+1 })
		c.send(map[string]interface{}{"type": "rearm"})
		tc.Rearmed = !again && c.waitFor(4*time.Second, func() bool { return len(c.captures) > beforeCaps+1 })
	}
	c.send(map[string]interface{}{"type": "disarm"})
	log.Printf("trigger %s/%s fired=%v edgeOk=%v winOk=%v rearm=%v",
		edge, mode, tc.Fired, tc.EdgeOK, tc.WindowOK, tc.Rearmed)
	return tc
}

// edgeCrosses verifies the captured waveform crosses thr in the edge direction
// near trigTime.
func edgeCrosses(p points, trigTime float64, edge string, thr float64) bool {
	// find the sample pair straddling trigTime
	for i := 1; i < len(p.T); i++ {
		if p.T[i-1] <= trigTime && p.T[i] >= trigTime {
			a, b := p.V[i-1], p.V[i]
			switch edge {
			case "rising":
				return a <= thr && b >= thr
			case "falling":
				return a >= thr && b <= thr
			case "both":
				return (a <= thr && b >= thr) || (a >= thr && b <= thr)
			}
		}
	}
	return false
}