package orchestrator

import (
	"math/rand"
	"testing"

	"git.hangman-lab.top/hzhang/Dialectic.Backend/internal/models"
)

func sig(agentID string, camps ...models.Camp) models.SignupView {
	return models.SignupView{AgentID: agentID, WillingCamps: camps}
}

// Helper: assert no duplicate agents across the 3 camps.
func assertDistinct(t *testing.T, alloc map[models.Camp]string) {
	t.Helper()
	seen := map[string]models.Camp{}
	for c, a := range alloc {
		if prev, ok := seen[a]; ok {
			t.Fatalf("agent %q allocated to both %s and %s", a, prev, c)
		}
		seen[a] = c
	}
}

func TestAllocate_EmptyPoolCancels(t *testing.T) {
	r := Allocate(nil, rand.New(rand.NewSource(1)))
	if r.CancelReason == "" {
		t.Fatal("expected cancel reason, got allocation")
	}
}

func TestAllocate_TwoSignupsCancels(t *testing.T) {
	r := Allocate([]models.SignupView{
		sig("a", models.CampPro),
		sig("b", models.CampCon),
	}, rand.New(rand.NewSource(1)))
	if r.CancelReason == "" {
		t.Fatalf("expected cancel reason (pool<3), got %v", r.Allocation)
	}
}

func TestAllocate_OneVolunteerPerCampFills(t *testing.T) {
	signups := []models.SignupView{
		sig("a", models.CampPro),
		sig("b", models.CampCon),
		sig("c", models.CampJudge),
	}
	r := Allocate(signups, rand.New(rand.NewSource(1)))
	if r.CancelReason != "" {
		t.Fatalf("unexpected cancel: %s", r.CancelReason)
	}
	if r.Allocation[models.CampPro] != "a" || r.Allocation[models.CampCon] != "b" || r.Allocation[models.CampJudge] != "c" {
		t.Fatalf("wrong allocation: %v", r.Allocation)
	}
	assertDistinct(t, r.Allocation)
}

func TestAllocate_AgentMultiVolunteerPicksOnlyOnce(t *testing.T) {
	// 'a' volunteers for all 3 camps. Should only be allocated to one
	// (pro, since it's first in iteration order); other camps need
	// other volunteers or get filled via backfill.
	signups := []models.SignupView{
		sig("a", models.CampPro, models.CampCon, models.CampJudge),
		sig("b", models.CampCon),
		sig("c", models.CampJudge),
	}
	r := Allocate(signups, rand.New(rand.NewSource(1)))
	if r.CancelReason != "" {
		t.Fatalf("unexpected cancel: %s", r.CancelReason)
	}
	if r.Allocation[models.CampPro] != "a" {
		t.Fatalf("expected 'a' in pro, got %v", r.Allocation)
	}
	if r.Allocation[models.CampCon] != "b" {
		t.Fatalf("expected 'b' in con, got %v", r.Allocation)
	}
	if r.Allocation[models.CampJudge] != "c" {
		t.Fatalf("expected 'c' in judge, got %v", r.Allocation)
	}
	assertDistinct(t, r.Allocation)
}

func TestAllocate_BackfillFromUnallocated(t *testing.T) {
	// pro has 2 volunteers ('a','c'), con has 1 ('b'), judge has 0.
	// Allocator picks one of {a,c} for pro, then b for con, then
	// backfills judge from whichever of {a,c} is unallocated.
	signups := []models.SignupView{
		sig("a", models.CampPro),
		sig("b", models.CampCon),
		sig("c", models.CampPro),
	}
	r := Allocate(signups, rand.New(rand.NewSource(1)))
	if r.CancelReason != "" {
		t.Fatalf("unexpected cancel: %s; alloc=%v", r.CancelReason, r.Allocation)
	}
	assertDistinct(t, r.Allocation)
	if len(r.Allocation) != 3 {
		t.Fatalf("expected all 3 camps filled; got %d (%v)", len(r.Allocation), r.Allocation)
	}
	// Con must be 'b' (only volunteer); pro and judge must be {a, c} in some order.
	if r.Allocation[models.CampCon] != "b" {
		t.Fatalf("expected con=b, got %v", r.Allocation)
	}
	pro := r.Allocation[models.CampPro]
	judge := r.Allocation[models.CampJudge]
	if !(pro == "a" && judge == "c") && !(pro == "c" && judge == "a") {
		t.Fatalf("expected pro/judge to be {a,c} permutation, got pro=%s judge=%s", pro, judge)
	}
}

func TestAllocate_BackfillInsufficientCancels(t *testing.T) {
	// pro filled by 'a'; con filled by 'b'; judge has no volunteer
	// AND no remaining unallocated signups → cancel.
	signups := []models.SignupView{
		sig("a", models.CampPro),
		sig("b", models.CampCon),
	}
	r := Allocate(signups, rand.New(rand.NewSource(1)))
	if r.CancelReason == "" {
		t.Fatalf("expected cancel; got allocation %v", r.Allocation)
	}
}

func TestAllocate_LargePoolDistinctness(t *testing.T) {
	// Many signups, all willing for all camps. Allocation should pick 3
	// distinct agents, randomly.
	signups := []models.SignupView{}
	for i := 0; i < 20; i++ {
		signups = append(signups, sig(string(rune('a'+i)), models.CampPro, models.CampCon, models.CampJudge))
	}
	r := Allocate(signups, rand.New(rand.NewSource(42)))
	if r.CancelReason != "" {
		t.Fatalf("unexpected cancel: %s", r.CancelReason)
	}
	assertDistinct(t, r.Allocation)
}