package staffing

import (
	"context"
	"errors"
	"testing"
)

func solve(t *testing.T, in PortfolioInput) PortfolioSolution {
	t.Helper()
	sol, err := SolvePortfolio(context.Background(), in, 1_000_000, 100_000)
	if err != nil {
		t.Fatalf("solve: %v", err)
	}
	return sol
}

func TestSolveEmpty(t *testing.T) {
	if sol := solve(t, PortfolioInput{}); len(sol.Assignments) != 0 || sol.TotalValue != 0 {
		t.Fatalf("empty problem must yield an empty allocation, got %+v", sol)
	}
	// candidates but no edges -> empty (everyone is a source).
	if sol := solve(t, PortfolioInput{NumCandidates: 3, Slots: []int{1, 1}}); len(sol.Assignments) != 0 {
		t.Fatalf("no edges means no assignment, got %+v", sol)
	}
}

func TestSolveSingle(t *testing.T) {
	sol := solve(t, PortfolioInput{
		NumCandidates: 1,
		Slots:         []int{1},
		Edges:         []AssignEdge{{0, 0, 100}},
	})
	if len(sol.Assignments) != 1 || sol.TotalValue != 100 {
		t.Fatalf("single profitable pair must be assigned, got %+v", sol)
	}
}

// TestSolveTrueOptimumNotGreedy is the proof that this is a real optimizer and not a greedy heuristic.
// Greedy gives c0 its best requirement r0 (10), leaving c1 only r1 (1), total 11. The true optimum routes
// c0 to r1 (8) and c1 to r0 (9), total 17. The successive shortest path solve must find 17.
func TestSolveTrueOptimumNotGreedy(t *testing.T) {
	sol := solve(t, PortfolioInput{
		NumCandidates: 2,
		Slots:         []int{1, 1},
		Edges: []AssignEdge{
			{0, 0, 10}, {0, 1, 8},
			{1, 0, 9}, {1, 1, 1},
		},
	})
	if sol.TotalValue != 17 {
		t.Fatalf("max profit must be 17 (greedy would give 11), got %d", sol.TotalValue)
	}
	got := map[int]int{}
	for _, a := range sol.Assignments {
		got[a.CandidateIdx] = a.RequirementIdx
	}
	if got[0] != 1 || got[1] != 0 {
		t.Fatalf("optimal allocation must be c0->r1 and c1->r0, got %+v", got)
	}
}

func TestSolveUnfilledSlotsAreImplicitSource(t *testing.T) {
	// One candidate, two single slot requirements they both fit; the candidate takes the higher value one,
	// the other slot is left unfilled (a source in the domain).
	sol := solve(t, PortfolioInput{
		NumCandidates: 1,
		Slots:         []int{1, 1},
		Edges:         []AssignEdge{{0, 0, 50}, {0, 1, 90}},
	})
	if len(sol.Assignments) != 1 || sol.Assignments[0].RequirementIdx != 1 || sol.TotalValue != 90 {
		t.Fatalf("the candidate must take the higher value slot, leaving the other unfilled, got %+v", sol)
	}
}

func TestSolveRespectsSlotCapacity(t *testing.T) {
	// One requirement with two slots, three qualified candidates -> exactly two assigned (the two highest).
	sol := solve(t, PortfolioInput{
		NumCandidates: 3,
		Slots:         []int{2},
		Edges:         []AssignEdge{{0, 0, 70}, {1, 0, 95}, {2, 0, 80}},
	})
	if len(sol.Assignments) != 2 {
		t.Fatalf("a two slot requirement must take exactly two candidates, got %d", len(sol.Assignments))
	}
	if sol.TotalValue != 175 { // 95 + 80
		t.Fatalf("the two highest value candidates must fill the slots, total 175, got %d", sol.TotalValue)
	}
}

func TestSolveDeterministic(t *testing.T) {
	in := PortfolioInput{
		NumCandidates: 4,
		Slots:         []int{1, 1, 1},
		Edges: []AssignEdge{
			{0, 0, 50}, {0, 1, 50}, {1, 0, 50}, {1, 1, 50},
			{2, 1, 50}, {2, 2, 50}, {3, 0, 50}, {3, 2, 50},
		},
	}
	first := solve(t, in)
	for i := 0; i < 20; i++ {
		again := solve(t, in)
		if again.TotalValue != first.TotalValue || len(again.Assignments) != len(first.Assignments) {
			t.Fatalf("non deterministic total: %+v vs %+v", first, again)
		}
		for k := range again.Assignments {
			if again.Assignments[k] != first.Assignments[k] {
				t.Fatalf("non deterministic allocation at %d: %+v vs %+v", k, first.Assignments, again.Assignments)
			}
		}
	}
}

func TestSolveProblemTooLarge(t *testing.T) {
	in := PortfolioInput{NumCandidates: 2, Slots: []int{1}, Edges: []AssignEdge{{0, 0, 1}, {1, 0, 1}}}
	if _, err := SolvePortfolio(context.Background(), in, 1, 100); !errors.Is(err, ErrProblemTooLarge) {
		t.Fatalf("an edge count over the cap must be refused, got %v", err)
	}
	if _, err := SolvePortfolio(context.Background(), in, 100, 2); !errors.Is(err, ErrProblemTooLarge) {
		t.Fatalf("a node count over the cap must be refused, got %v", err)
	}
}

func TestSolveBudgetExceeded(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // already cancelled: the solve must refuse rather than run
	in := PortfolioInput{NumCandidates: 1, Slots: []int{1}, Edges: []AssignEdge{{0, 0, 100}}}
	if _, err := SolvePortfolio(ctx, in, 1000, 1000); !errors.Is(err, ErrSolveBudget) {
		t.Fatalf("a cancelled context must surface ErrSolveBudget, got %v", err)
	}
}

func TestCounterfactualsKnown(t *testing.T) {
	in := PortfolioInput{
		NumCandidates: 2,
		Slots:         []int{1, 1},
		Edges: []AssignEdge{
			{0, 0, 10}, {0, 1, 8},
			{1, 0, 9}, {1, 1, 1},
		},
	}
	sol := solve(t, in) // c0->r1 (8), c1->r0 (9)
	cf := Counterfactuals(in, sol)
	// c0 assigned to r1 (8); next best requirement is r0 (10); delta = 8 - 10 = -2 (giving c0 r1 instead of
	// its own best protected the portfolio).
	if cf[0].CandidateNextReqIdx != 0 || cf[0].CandidateValueDelta != -2 {
		t.Fatalf("c0 counterfactual wrong: %+v", cf[0])
	}
	// requirement r1's next best candidate (other than c0) is c1 at value 1; assigned value 8; delta 7.
	if cf[0].RequirementNextCandIdx != 1 || cf[0].RequirementValueDelta != 7 {
		t.Fatalf("r1 next best candidate counterfactual wrong: %+v", cf[0])
	}
}