package agenttask

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

func TestSchedulerProviderCapacityAndParallelRelease(t *testing.T) {
	units := []WorkUnit{
		testUnit("one", WriteSetDisjoint),
		testUnit("two", WriteSetOverlap),
		testUnit("three", WriteSetUnknown),
	}
	snapshot := testSnapshot("project", "workspace", units...)
	harness := newHarness(t, map[ProjectID]ProjectWorkflowSnapshot{"project": snapshot}, 2)
	for _, unit := range units {
		harness.invoker.delays[unit.ID] = 25 * time.Millisecond
	}
	harness.start("project", "workspace", nil)
	if err := harness.manager.Reconcile(context.Background()); err != nil {
		t.Fatalf("Reconcile: %v", err)
	}
	if maximum := harness.invoker.maxConcurrency(); maximum != 2 {
		t.Fatalf("max provider concurrency = %d, want 2", maximum)
	}
	if harness.invoker.callCount() != 3 {
		t.Fatalf("provider calls = %d, want 3", harness.invoker.callCount())
	}
}

func TestSchedulerCancelReleasesTicket(t *testing.T) {
	scheduler := NewScheduler()
	target := ExecutionTarget{
		ProviderID: "provider", ProfileID: "profile", Capacity: 1,
	}
	first, err := scheduler.Acquire(context.Background(), DispatchCandidate{
		ProjectID: "p1", WorkspaceID: "w", WorkUnitID: "one", AttemptID: "one#1",
		Target: target,
	})
	if err != nil {
		t.Fatalf("first Acquire: %v", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err = scheduler.Acquire(ctx, DispatchCandidate{
		ProjectID: "p2", WorkspaceID: "w", WorkUnitID: "two", AttemptID: "two#1",
		Target: target,
	})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("cancelled Acquire error = %v", err)
	}
	first.Release()
	second, err := scheduler.Acquire(context.Background(), DispatchCandidate{
		ProjectID: "p2", WorkspaceID: "w", WorkUnitID: "two", AttemptID: "two#1",
		Target: target,
	})
	if err != nil {
		t.Fatalf("Acquire after release: %v", err)
	}
	second.Release()
	if scheduler.Active(target.PoolKey()) != 0 {
		t.Fatalf("scheduler leaked capacity")
	}
}

func TestIsolatedDispatchUsesDistinctTaskRoots(t *testing.T) {
	snapshot := testSnapshot(
		"project", "workspace",
		testUnit("overlap-a", WriteSetOverlap),
		testUnit("overlap-b", WriteSetOverlap),
	)
	harness := newHarness(t, map[ProjectID]ProjectWorkflowSnapshot{"project": snapshot}, 2)
	harness.start("project", "workspace", nil)
	if err := harness.manager.Reconcile(context.Background()); err != nil {
		t.Fatalf("Reconcile: %v", err)
	}
	roots := harness.invoker.roots()
	if len(roots) != 2 || roots[0] == roots[1] {
		t.Fatalf("isolated task roots = %v", roots)
	}
}