package synth

import (
	"math"
	"testing"

	"github.com/netsynth/netsynth/classify"
)

func TestNewBankHas14Layers(t *testing.T) {
	b := NewBank(1.0, ClassFreqConfigs)
	if len(b.layers) != 14 {
		t.Errorf("NewBank() has %d layers, want 14", len(b.layers))
	}
	// Verify each class has exactly one layer
	for _, class := range classify.AllClasses() {
		if _, ok := b.layers[class]; !ok {
			t.Errorf("NewBank() missing layer for class %q", class)
		}
	}
}

func TestRenderWindowOutputLength(t *testing.T) {
	b := NewBank(1.0, ClassFreqConfigs)
	snap := classify.WindowSnapshot{
		Counts:       make(map[classify.TrafficClass]int64),
		TotalPackets: 0,
		WindowIndex:  0,
	}
	frames := b.RenderWindow(snap)
	if len(frames) != SamplesPerWindow {
		t.Errorf("RenderWindow returned %d frames, want %d (SamplesPerWindow)", len(frames), SamplesPerWindow)
	}
}

func TestRenderWindowSilentWhenNoTraffic(t *testing.T) {
	b := NewBank(1.0, ClassFreqConfigs)
	// Empty counts — no class ever seen — all layers should stay at zero amplitude
	snap := classify.WindowSnapshot{
		Counts:       make(map[classify.TrafficClass]int64),
		TotalPackets: 0,
		WindowIndex:  0,
	}
	frames := b.RenderWindow(snap)
	for i, frame := range frames {
		if frame[0] != 0.0 || frame[1] != 0.0 {
			t.Errorf("frame[%d] = [%v, %v], want [0, 0] (silent when no traffic seen)", i, frame[0], frame[1])
			break
		}
	}
}

func TestRenderWindowNonZeroWithTraffic(t *testing.T) {
	b := NewBank(1.0, ClassFreqConfigs)
	counts := make(map[classify.TrafficClass]int64)
	counts[classify.ClassICMP] = 100
	snap := classify.WindowSnapshot{
		Counts:       counts,
		TotalPackets: 100,
		WindowIndex:  0,
	}
	frames := b.RenderWindow(snap)
	// Check that at least some frames are non-zero
	hasNonZero := false
	for _, frame := range frames {
		if frame[0] != 0.0 || frame[1] != 0.0 {
			hasNonZero = true
			break
		}
	}
	if !hasNonZero {
		t.Error("RenderWindow with ICMP count=100 should produce non-zero frames")
	}
}

func TestMixerNoClip(t *testing.T) {
	b := NewBank(0.01, ClassFreqConfigs) // fast EMA to quickly ramp up to near-max amplitude
	counts := make(map[classify.TrafficClass]int64)
	// All 14 classes at max count — worst-case mixing scenario
	for class := range ClassFreqConfigs {
		counts[class] = 1000
	}
	snap := classify.WindowSnapshot{
		Counts:       counts,
		TotalPackets: int64(len(ClassFreqConfigs)) * 1000,
		WindowIndex:  0,
	}
	// Render multiple windows to let EMA converge
	for i := 0; i < 10; i++ {
		frames := b.RenderWindow(snap)
		for _, frame := range frames {
			if frame[0] > 1.0 || frame[0] < -1.0 {
				t.Errorf("left channel clipped: %v (exceeds [-1, 1])", frame[0])
				return
			}
			if frame[1] > 1.0 || frame[1] < -1.0 {
				t.Errorf("right channel clipped: %v (exceeds [-1, 1])", frame[1])
				return
			}
		}
	}
}

func TestStereoPan(t *testing.T) {
	b := NewBank(0.01, ClassFreqConfigs) // fast EMA
	counts := make(map[classify.TrafficClass]int64)
	// ClassDHCP has pan=-0.75 (wide-left in config.go)
	counts[classify.ClassDHCP] = 1000
	snap := classify.WindowSnapshot{
		Counts:       counts,
		TotalPackets: 1000,
		WindowIndex:  0,
	}
	// Render multiple windows to allow EMA to build up amplitude
	var frames [][2]float64
	for i := 0; i < 5; i++ {
		frames = b.RenderWindow(snap)
	}
	// Compute RMS for L and R channels
	var sumL2, sumR2 float64
	for _, frame := range frames {
		sumL2 += frame[0] * frame[0]
		sumR2 += frame[1] * frame[1]
	}
	rmsL := math.Sqrt(sumL2 / float64(len(frames)))
	rmsR := math.Sqrt(sumR2 / float64(len(frames)))
	if rmsL <= rmsR {
		t.Errorf("ClassDHCP (pan=-0.75) should have rmsL > rmsR; got rmsL=%v, rmsR=%v", rmsL, rmsR)
	}
}

func TestMultipleWindowsEMAConvergence(t *testing.T) {
	b := NewBank(1.0, ClassFreqConfigs)
	counts := make(map[classify.TrafficClass]int64)
	counts[classify.ClassICMP] = 100
	snap := classify.WindowSnapshot{
		Counts:       counts,
		TotalPackets: 100,
		WindowIndex:  0,
	}
	// Compute RMS for first and last window render
	rmsFirst := windowRMS(b.RenderWindow(snap))
	// Render 4 more windows with the same snapshot
	var rmsLast float64
	for i := 0; i < 4; i++ {
		rmsLast = windowRMS(b.RenderWindow(snap))
	}
	if rmsLast <= rmsFirst {
		t.Errorf("EMA should converge upward: rmsFirst=%v, rmsLast=%v", rmsFirst, rmsLast)
	}
}

func TestNewBankDynamicGain(t *testing.T) {
	// Create a config map with only 3 classes
	cfgs := map[classify.TrafficClass]FreqConfig{
		classify.ClassICMP:  ClassFreqConfigs[classify.ClassICMP],
		classify.ClassDNS:   ClassFreqConfigs[classify.ClassDNS],
		classify.ClassHTTPS: ClassFreqConfigs[classify.ClassHTTPS],
	}
	b := NewBank(0.01, cfgs)
	if len(b.layers) != 3 {
		t.Errorf("NewBank with 3 configs has %d layers, want 3", len(b.layers))
	}
	// Verify gainPerLayer is 1/3
	expected := 1.0 / 3.0
	if b.gainPerLayer != expected {
		t.Errorf("gainPerLayer = %v, want %v", b.gainPerLayer, expected)
	}
}

func TestNewBankCustomConfigNoClip(t *testing.T) {
	cfgs := map[classify.TrafficClass]FreqConfig{
		classify.ClassICMP: ClassFreqConfigs[classify.ClassICMP],
		classify.ClassDNS:  ClassFreqConfigs[classify.ClassDNS],
	}
	b := NewBank(0.01, cfgs)
	counts := map[classify.TrafficClass]int64{
		classify.ClassICMP: 1000,
		classify.ClassDNS:  1000,
	}
	snap := classify.WindowSnapshot{Counts: counts, TotalPackets: 2000, WindowIndex: 0}
	for i := 0; i < 10; i++ {
		for _, frame := range b.RenderWindow(snap) {
			if frame[0] > 1.0 || frame[0] < -1.0 || frame[1] > 1.0 || frame[1] < -1.0 {
				t.Fatalf("clipped with 2-class config: L=%v R=%v", frame[0], frame[1])
			}
		}
	}
}

// windowRMS computes the root mean square amplitude across all stereo frames.
func windowRMS(frames [][2]float64) float64 {
	var sum float64
	for _, frame := range frames {
		sum += frame[0]*frame[0] + frame[1]*frame[1]
	}
	return math.Sqrt(sum / float64(len(frames)*2))
}