package engine

import (
	"sync"
	"testing"
	"time"
)

// newTestSubAgent creates a minimal SubAgent for registry tests.
// done channel is open by default; close it to simulate completion.
func newTestSubAgent(id string, status SubAgentStatus) *SubAgent {
	return &SubAgent{
		ID:       id,
		Status:   status,
		Progress: &SubAgentProgress{},
		done:     make(chan struct{}),
	}
}

func TestNewSubAgentRegistry(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	if r == nil {
		t.Fatal("NewSubAgentRegistry returned nil")
	}
	if r.Count() != 0 {
		t.Errorf("new registry count = %d, want 0", r.Count())
	}
	if list := r.List(); len(list) != 0 {
		t.Errorf("new registry list len = %d, want 0", len(list))
	}
}

func TestSubAgentRegistry_RegisterAndGet(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	sa := newTestSubAgent("sa-1", SubAgentStatusRunning)

	if err := r.Register(sa); err != nil {
		t.Fatalf("Register: %v", err)
	}
	if r.Count() != 1 {
		t.Errorf("count after register = %d, want 1", r.Count())
	}

	got, ok := r.Get("sa-1")
	if !ok {
		t.Fatal("Get sa-1: not found")
	}
	if got.ID != "sa-1" {
		t.Errorf("Get sa-1: ID = %q", got.ID)
	}

	// non-existent
	_, ok = r.Get("sa-999")
	if ok {
		t.Error("Get sa-999: should not exist")
	}
}

func TestSubAgentRegistry_RegisterDuplicate(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	sa := newTestSubAgent("dup", SubAgentStatusRunning)
	if err := r.Register(sa); err != nil {
		t.Fatalf("first Register: %v", err)
	}

	sa2 := newTestSubAgent("dup", SubAgentStatusPending)
	if err := r.Register(sa2); err == nil {
		t.Error("duplicate Register should return error")
	}
	if r.Count() != 1 {
		t.Errorf("count after duplicate = %d, want 1", r.Count())
	}
}

func TestSubAgentRegistry_List(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	ids := []string{"a", "b", "c"}
	for _, id := range ids {
		sa := newTestSubAgent(id, SubAgentStatusRunning)
		if err := r.Register(sa); err != nil {
			t.Fatalf("Register %s: %v", id, err)
		}
	}

	list := r.List()
	if len(list) != 3 {
		t.Fatalf("List len = %d, want 3", len(list))
	}

	// check all IDs present (order not guaranteed due to map iteration)
	found := make(map[string]bool)
	for _, sa := range list {
		found[sa.ID] = true
	}
	for _, id := range ids {
		if !found[id] {
			t.Errorf("List missing ID %q", id)
		}
	}
}

func TestSubAgentRegistry_ListByStatus(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	running := newTestSubAgent("r1", SubAgentStatusRunning)
	completed := newTestSubAgent("c1", SubAgentStatusCompleted)
	failed := newTestSubAgent("f1", SubAgentStatusFailed)

	for _, sa := range []*SubAgent{running, completed, failed} {
		if err := r.Register(sa); err != nil {
			t.Fatalf("Register %s: %v", sa.ID, err)
		}
	}

	runningList := r.ListByStatus(SubAgentStatusRunning)
	if len(runningList) != 1 || runningList[0].ID != "r1" {
		t.Errorf("ListByStatus(Running) = %v, want [r1]", runningList)
	}

	pendingList := r.ListByStatus(SubAgentStatusPending)
	if len(pendingList) != 0 {
		t.Errorf("ListByStatus(Pending) = %v, want []", pendingList)
	}
}

func TestSubAgentRegistry_Remove(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	sa := newTestSubAgent("rm-1", SubAgentStatusCompleted)
	_ = r.Register(sa)

	if !r.Remove("rm-1") {
		t.Error("Remove rm-1: should return true")
	}
	if r.Count() != 0 {
		t.Errorf("count after remove = %d, want 0", r.Count())
	}

	// remove non-existent
	if r.Remove("rm-1") {
		t.Error("Remove rm-1 again: should return false")
	}
}

func TestSubAgentRegistry_CancelAll(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	// create a cancellable SubAgent
	cancelled := false
	sa := newTestSubAgent("cancel-me", SubAgentStatusRunning)
	sa.cancel = func() { cancelled = true }

	_ = r.Register(sa)

	// also register a non-running one (should not be cancelled)
	sa2 := newTestSubAgent("already-done", SubAgentStatusCompleted)
	_ = r.Register(sa2)

	r.CancelAll()

	if !cancelled {
		t.Error("CancelAll did not invoke cancel on running SubAgent")
	}
}

func TestSubAgentRegistry_OnComplete(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	completedIDs := make(chan string, 2)
	r.OnComplete(func(sa *SubAgent) {
		completedIDs <- sa.ID
	})

	sa := newTestSubAgent("watch-1", SubAgentStatusRunning)
	_ = r.Register(sa)

	// simulate completion by closing the done channel
	close(sa.done)

	select {
	case id := <-completedIDs:
		if id != "watch-1" {
			t.Errorf("OnComplete callback got ID = %q, want watch-1", id)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("OnComplete callback not fired within timeout")
	}
}

func TestSubAgentRegistry_OnComplete_MultipleCallbacks(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	var mu sync.Mutex
	callCount := 0

	r.OnComplete(func(sa *SubAgent) {
		mu.Lock()
		callCount++
		mu.Unlock()
	})
	r.OnComplete(func(sa *SubAgent) {
		mu.Lock()
		callCount++
		mu.Unlock()
	})

	sa := newTestSubAgent("multi-cb", SubAgentStatusRunning)
	_ = r.Register(sa)
	close(sa.done)

	// wait for callbacks
	time.Sleep(100 * time.Millisecond)

	mu.Lock()
	if callCount != 2 {
		t.Errorf("expected 2 callbacks, got %d", callCount)
	}
	mu.Unlock()
}

func TestSubAgentRegistry_Close_StopsWatchers(t *testing.T) {
	r := NewSubAgentRegistry()

	sa := newTestSubAgent("hang", SubAgentStatusRunning)
	_ = r.Register(sa)

	// sa.done is never closed -- watchCompletion would hang forever
	// without r.stopCh. Close() should unblock it.
	r.Close()

	// double close is safe (idempotent)
	r.Close()
}

func TestSubAgentRegistry_Summary(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	agents := []*SubAgent{
		newTestSubAgent("r1", SubAgentStatusRunning),
		newTestSubAgent("r2", SubAgentStatusRunning),
		newTestSubAgent("c1", SubAgentStatusCompleted),
		newTestSubAgent("f1", SubAgentStatusFailed),
	}
	for _, sa := range agents {
		_ = r.Register(sa)
	}

	summary := r.Summary()
	if summary[SubAgentStatusRunning] != 2 {
		t.Errorf("running count = %d, want 2", summary[SubAgentStatusRunning])
	}
	if summary[SubAgentStatusCompleted] != 1 {
		t.Errorf("completed count = %d, want 1", summary[SubAgentStatusCompleted])
	}
	if summary[SubAgentStatusFailed] != 1 {
		t.Errorf("failed count = %d, want 1", summary[SubAgentStatusFailed])
	}
	if summary[SubAgentStatusPending] != 0 {
		t.Errorf("pending count = %d, want 0", summary[SubAgentStatusPending])
	}
}

func TestSubAgentRegistry_ConcurrentAccess(t *testing.T) {
	r := NewSubAgentRegistry()
	defer r.Close()

	// pre-register some agents
	for i := 0; i < 10; i++ {
		sa := newTestSubAgent(
			"concurrent-"+string(rune('a'+i)),
			SubAgentStatusRunning,
		)
		_ = r.Register(sa)
	}

	// concurrent reads and writes should not race
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func(n int) {
			defer wg.Done()
			_ = r.List()
			_ = r.Count()
			_ = r.Summary()
			r.ListByStatus(SubAgentStatusRunning)
			r.Get("concurrent-a")
		}(i)
	}
	wg.Wait()
}