package router

import (
	"encoding/json"
	"fmt"
	"orchestra/internal/authz"
	"orchestra/internal/domain"
	"orchestra/internal/registry"
	"orchestra/internal/store"
	"sort"
	"testing"
	"time"
)

// These are p95 budgets for a complete scheduling pass on an otherwise idle
// coordinator. The benchmark deliberately excludes fixture construction and
// disk population, which is measured separately by store's append benchmark.
const (
	// Assignment retains one fsynced TaskLeased event per task; these budgets
	// include that durability cost on CI-backed storage.
	assignmentP95Budget1K  = 6 * time.Second
	assignmentP95Budget10K = time.Minute
)

func BenchmarkAssignPending1K(b *testing.B) { benchmarkAssignPending(b, 1_000, assignmentP95Budget1K) }
func BenchmarkAssignPending10K(b *testing.B) {
	benchmarkAssignPending(b, 10_000, assignmentP95Budget10K)
}

func benchmarkAssignPending(b *testing.B, tasks int, budget time.Duration) {
	b.Helper()
	samples := make([]time.Duration, 0, b.N)
	b.ResetTimer()
	for sample := 0; sample < b.N; sample++ {
		b.StopTimer()
		s, err := store.Open(b.TempDir())
		if err != nil {
			b.Fatal(err)
		}
		for i := 0; i < tasks; i++ {
			id := fmt.Sprintf("task-%d", i)
			payload, _ := json.Marshal(map[string]any{"source": "benchmark", "external_id": id, "project": "p"})
			if err := s.Append(domain.Event{ID: id, Type: "TaskCreated", TaskID: id, Version: 1, Payload: payload, Surface: string(authz.System)}); err != nil {
				b.Fatal(err)
			}
		}
		r, err := registry.New(registry.Config{
			Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
			Machines: []registry.Machine{{ID: "m", Address: "unused"}},
			Herdrs:   []registry.Herdr{{ID: "h", MachineID: "m"}}, // unlimited concurrency
		})
		if err != nil {
			b.Fatal(err)
		}
		router := Router{Store: s, Registry: r}
		b.StartTimer()
		started := time.Now()
		assigned, err := router.AssignPending()
		samples = append(samples, time.Since(started))
		b.StopTimer()
		if err != nil {
			b.Fatal(err)
		}
		if len(assigned) != tasks {
			b.Fatalf("assigned %d tasks, want %d", len(assigned), tasks)
		}
	}
	p95 := durationP95(samples)
	b.ReportMetric(float64(p95.Microseconds()), "assign_p95_us")
	if p95 > budget {
		b.Fatalf("assignment p95 %s exceeds budget %s for %d tasks", p95, budget, tasks)
	}
}

func durationP95(samples []time.Duration) time.Duration {
	if len(samples) == 0 {
		return 0
	}
	sorted := append([]time.Duration(nil), samples...)
	sort.Slice(sorted, func(i, j int) bool { return sorted[i] < sorted[j] })
	return sorted[(len(sorted)*95+99)/100-1]
}