package learning

import (
	"math"
	"testing"

	"github.com/dtoro/oikos/internal/db/sqlcgen"
)

func TestWilsonLowerBound(t *testing.T) {
	tests := []struct {
		name    string
		success float64
		total   float64
		z       float64
		want    float64
		// For approximate checks
		approx bool
		epsilon float64
	}{
		{"zero total", 5, 0, 1.96, 0, false, 0},
		{"zero success", 0, 10, 1.96, 0, true, 1e-10},
		{"all success small n", 3, 3, 1.96, 0, true, 0.5},
		{"all success large n", 100, 100, 1.96, 0, true, 0.05},
		{"half success large n", 50, 100, 1.96, 0.39, true, 0.02},
		{"higher z gives lower bound", 8, 10, 3.0, 0, true, 0.5},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := wilsonLowerBound(tt.success, tt.total, tt.z)
			if tt.approx {
				if tt.want > 0 && math.Abs(got-tt.want) > tt.epsilon {
					t.Errorf("wilsonLowerBound(%v, %v, %v) = %v, want ~%v (±%v)", tt.success, tt.total, tt.z, got, tt.want, tt.epsilon)
				}
				if got < 0 {
					t.Errorf("wilsonLowerBound returned negative: %v", got)
				}
				if got > 1 {
					t.Errorf("wilsonLowerBound returned >1: %v", got)
				}
			} else {
				if got != tt.want {
					t.Errorf("wilsonLowerBound(%v, %v, %v) = %v, want %v", tt.success, tt.total, tt.z, got, tt.want)
				}
			}
		})
	}
}

func TestWilsonLowerBoundMonotonic(t *testing.T) {
	prev := 0.0
	for s := 0.0; s <= 20; s++ {
		got := wilsonLowerBound(s, 20, 1.96)
		if got < prev-1e-9 {
			t.Errorf("not monotonically increasing: s=%v got=%v prev=%v", s, got, prev)
		}
		prev = got
	}
}

func TestCountOutcomes(t *testing.T) {
	tests := []struct {
		name        string
		items       []sqlcgen.GetFeedbackAfterWatermarkRow
		wantSuccess int
		wantFailure int
	}{
		{
			"empty",
			nil,
			0, 0,
		},
		{
			"all success",
			[]sqlcgen.GetFeedbackAfterWatermarkRow{
				{Outcome: "success"},
				{Outcome: "success"},
			},
			2, 0,
		},
		{
			"all failure",
			[]sqlcgen.GetFeedbackAfterWatermarkRow{
				{Outcome: "failure"},
				{Outcome: "unexpected"},
			},
			0, 2,
		},
		{
			"mixed including partial",
			[]sqlcgen.GetFeedbackAfterWatermarkRow{
				{Outcome: "success"},
				{Outcome: "failure"},
				{Outcome: "partial"},
				{Outcome: "unexpected"},
			},
			2, 2,
		},
		{
			"unknown outcome ignored",
			[]sqlcgen.GetFeedbackAfterWatermarkRow{
				{Outcome: "success"},
				{Outcome: "bogus"},
			},
			1, 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s, f := countOutcomes(tt.items)
			if s != tt.wantSuccess || f != tt.wantFailure {
				t.Errorf("countOutcomes() = (%d, %d), want (%d, %d)", s, f, tt.wantSuccess, tt.wantFailure)
			}
		})
	}
}

func TestComputeConfidence(t *testing.T) {
	tests := []struct {
		name    string
		success int
		failure int
		wantMax float64
		wantMin float64
	}{
		{"zero total", 0, 0, 0, 0},
		{"one success no failures", 1, 0, 0.2, 0},
		{"five success no failures", 5, 0, 1.0, 0.3},
		{"ten success no failures", 10, 0, 1.0, 0.5},
		{"half success", 5, 5, 0.5, 0.2},
		{"all failures", 0, 10, 0.01, 0},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := computeConfidence(tt.success, tt.failure)
			if got < 0 || got > 1 {
				t.Errorf("confidence out of [0,1]: %v", got)
			}
			if got > tt.wantMax+0.01 {
				t.Errorf("confidence too high: got %v, max ~%v", got, tt.wantMax)
			}
			if tt.wantMin > 0 && got < tt.wantMin-0.1 {
				t.Errorf("confidence too low: got %v, min ~%v", got, tt.wantMin)
			}
		})
	}
}

func TestComputeConfidenceSampleSizeCap(t *testing.T) {
	got := computeConfidence(1, 0)
	if got > 0.2+1e-9 {
		t.Errorf("sample size cap not applied: 1 sample should cap at 1/5=0.2, got %v", got)
	}
	got = computeConfidence(4, 0)
	if got > 0.8+1e-9 {
		t.Errorf("sample size cap not applied: 4 samples should cap at 4/5=0.8, got %v", got)
	}
}

func TestShouldValidate(t *testing.T) {
	tests := []struct {
		name          string
		evidenceCount int
		confidence    float64
		quarantined   bool
		want          bool
	}{
		{"enough evidence and confidence", 5, 0.7, false, true},
		{"high evidence high confidence", 10, 0.9, false, true},
		{"not enough evidence", 4, 0.9, false, false},
		{"not enough confidence", 5, 0.69, false, false},
		{"quarantined blocks validation", 10, 0.9, true, false},
		{"exactly at threshold", 5, 0.7, false, true},
		{"zero evidence", 0, 0.9, false, false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := shouldValidate(tt.evidenceCount, tt.confidence, tt.quarantined)
			if got != tt.want {
				t.Errorf("shouldValidate(%d, %v, %v) = %v, want %v", tt.evidenceCount, tt.confidence, tt.quarantined, got, tt.want)
			}
		})
	}
}

func TestShouldQuarantine(t *testing.T) {
	tests := []struct {
		name  string
		total int
		want  bool
	}{
		{"zero", 0, false},
		{"small", 5, false},
		{"at threshold", 10, false},
		{"over threshold", 11, true},
		{"large burst", 50, true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := shouldQuarantine(tt.total)
			if got != tt.want {
				t.Errorf("shouldQuarantine(%d) = %v, want %v", tt.total, got, tt.want)
			}
		})
	}
}