package main

import (
	"fmt"
	"math"
	"sort"
	"time"

	pb "powerpod/gotool/pb"

	"google.golang.org/protobuf/proto"
)

type latencyStats struct {
	N      int     `json:"n"`
	MinMs  float64 `json:"min_ms"`
	MedMs  float64 `json:"median_ms"`
	MeanMs float64 `json:"mean_ms"`
	MaxMs  float64 `json:"max_ms"`
}

type bleRttResult struct {
	Success bool         `json:"success"`
	Count   int          `json:"count"`
	Warmup  int          `json:"warmup"`
	Stats   latencyStats `json:"stats"`
	Error   string       `json:"error,omitempty"`
}

type bleStreamResult struct {
	Success     bool         `json:"success"`
	Received    int          `json:"received"`
	Requested   int          `json:"requested"`
	PeriodMs    uint32       `json:"period_ms"`
	TolMs       float64      `json:"tol_ms"`
	OutsideTol  int          `json:"outside_tol"`
	WindowOK    bool         `json:"window_ok"`
	DTEsp       latencyStats `json:"dt_esp_ms"`
	DRPc        latencyStats `json:"dr_pc_ms"`
	ExtraVsFirst latencyStats `json:"extra_vs_first_ms"`
	Error       string       `json:"error,omitempty"`
}

func measureBleRtt(h hostConn, count, warmup int, settle time.Duration) bleRttResult {
	out := bleRttResult{Count: count, Warmup: warmup}
	if count < 1 {
		out.Error = "count must be >= 1"
		return out
	}
	if warmup < 0 {
		out.Error = "warmup must be >= 0"
		return out
	}
	if settle > 0 {
		time.Sleep(settle)
	}

	samples := make([]float64, 0, count)
	total := warmup + count
	for i := 0; i < total; i++ {
		t0 := time.Now()
		if _, err := hostGetVersion(h); err != nil {
			out.Error = fmt.Sprintf("VERSION round-trip %d: %v", i, err)
			return out
		}
		ms := float64(time.Since(t0).Microseconds()) / 1000.0
		if i >= warmup {
			samples = append(samples, ms)
		}
	}
	out.Stats = makeLatencyStats(samples)
	out.Success = true
	return out
}

func measureBleStream(h hostConn, clientID, periodMs uint32, count int, tolMs, timeoutSec float64) bleStreamResult {
	out := bleStreamResult{
		Requested: count,
		PeriodMs:  periodMs,
		TolMs:     tolMs,
	}
	if periodMs < 10 || periodMs > 16 {
		out.Error = "period_ms must be 10..16"
		return out
	}
	if count < 2 {
		out.Error = "count must be >= 2"
		return out
	}
	sc, ok := h.(hostStreamConn)
	if !ok {
		out.Error = "BLE stream ticks require BLE host"
		return out
	}

	req := &pb.AccelStreamRequest{
		Write:    true,
		Enable:   true,
		PeriodMs: periodMs,
	}
	if clientID == 0 {
		req.AllClients = true
	} else {
		req.ClientId = clientID
	}
	resp, err := hostAccelStream(h, req)
	if err != nil {
		out.Error = err.Error()
		return out
	}
	if !resp.GetSuccess() {
		out.Error = "accel stream enable failed"
		return out
	}

	type sample struct {
		espUs uint64
		pcT  float64
	}
	samples := make([]sample, 0, count)
	deadline := time.Now().Add(time.Duration(timeoutSec * float64(time.Second)))
	if timeoutSec <= 0 {
		deadline = time.Now().Add(15 * time.Second)
	}

	for len(samples) < count {
		remain := time.Until(deadline)
		if remain <= 0 {
			break
		}
		payload, err := sc.NextStreamTick(remain)
		if err != nil {
			break
		}
		pcT := time.Now().UnixMicro()
		var msg pb.UartMessage
		if err := proto.Unmarshal(payload[1:], &msg); err != nil {
			continue
		}
		tick := msg.GetBleStreamTick()
		if tick == nil {
			continue
		}
		samples = append(samples, sample{espUs: tick.GetEspTimestampUs(), pcT: float64(pcT)})
	}

	off := &pb.AccelStreamRequest{Write: true, Enable: false, AllClients: req.AllClients, ClientId: req.ClientId}
	_, _ = hostAccelStream(h, off)

	out.Received = len(samples)
	if len(samples) < 2 {
		out.Error = fmt.Sprintf("need >=2 ticks, got %d (stream running / DIP8=BLE?)", len(samples))
		return out
	}

	var dT, dR, extra []float64
	esp0 := samples[0].espUs
	pc0 := samples[0].pcT
	for i := 1; i < len(samples); i++ {
		dT = append(dT, float64(samples[i].espUs-samples[i-1].espUs)/1000.0)
		dR = append(dR, (samples[i].pcT-samples[i-1].pcT)/1000.0)
		relEsp := float64(samples[i].espUs-esp0) / 1000.0
		relPc := (samples[i].pcT - pc0) / 1000.0
		extra = append(extra, relPc-relEsp)
	}

	period := float64(periodMs)
	outCount := 0
	for _, dt := range dT {
		if math.Abs(dt-period) > tolMs {
			outCount++
		}
	}
	out.OutsideTol = outCount
	out.WindowOK = outCount*2 <= len(dT)
	out.DTEsp = makeLatencyStats(dT)
	out.DRPc = makeLatencyStats(dR)
	out.ExtraVsFirst = makeLatencyStats(extra)
	out.Success = true
	return out
}

func makeLatencyStats(v []float64) latencyStats {
	if len(v) == 0 {
		return latencyStats{}
	}
	s := append([]float64(nil), v...)
	sort.Float64s(s)
	return latencyStats{
		N:      len(s),
		MinMs:  s[0],
		MedMs:  percentile(s, 0.5),
		MeanMs: mean(s),
		MaxMs:  s[len(s)-1],
	}
}

func mean(v []float64) float64 {
	if len(v) == 0 {
		return math.NaN()
	}
	var s float64
	for _, x := range v {
		s += x
	}
	return s / float64(len(v))
}

func percentile(sorted []float64, p float64) float64 {
	if len(sorted) == 0 {
		return math.NaN()
	}
	if len(sorted) == 1 {
		return sorted[0]
	}
	idx := p * float64(len(sorted)-1)
	lo := int(idx)
	hi := lo + 1
	if hi >= len(sorted) {
		return sorted[lo]
	}
	frac := idx - float64(lo)
	return sorted[lo]*(1-frac) + sorted[hi]*frac
}