package shell

import (
	"context"
	"io"
	"testing"
)

// stubShell implements Shell for testing.
type stubShell struct {
	name string
}

func (s *stubShell) Name() string        { return s.name }
func (s *stubShell) Description() string { return "stub" }
func (s *stubShell) Handle(_ context.Context, _ *SessionContext, _ io.ReadWriteCloser) error {
	return nil
}

func TestRegistryRegisterAndGet(t *testing.T) {
	r := NewRegistry()
	sh := &stubShell{name: "test"}

	if err := r.Register(sh, 1); err != nil {
		t.Fatalf("Register: %v", err)
	}

	got, ok := r.Get("test")
	if !ok {
		t.Fatal("Get returned false")
	}
	if got.Name() != "test" {
		t.Errorf("Name = %q, want %q", got.Name(), "test")
	}
}

func TestRegistryGetMissing(t *testing.T) {
	r := NewRegistry()
	_, ok := r.Get("nope")
	if ok {
		t.Fatal("Get returned true for missing shell")
	}
}

func TestRegistryDuplicateName(t *testing.T) {
	r := NewRegistry()
	r.Register(&stubShell{name: "dup"}, 1)
	err := r.Register(&stubShell{name: "dup"}, 1)
	if err == nil {
		t.Fatal("expected error for duplicate name")
	}
}

func TestRegistryInvalidWeight(t *testing.T) {
	r := NewRegistry()
	err := r.Register(&stubShell{name: "a"}, 0)
	if err == nil {
		t.Fatal("expected error for weight 0")
	}
	err = r.Register(&stubShell{name: "b"}, -1)
	if err == nil {
		t.Fatal("expected error for negative weight")
	}
}

func TestRegistrySelectEmpty(t *testing.T) {
	r := NewRegistry()
	_, err := r.Select()
	if err == nil {
		t.Fatal("expected error from empty registry")
	}
}

func TestRegistrySelectSingle(t *testing.T) {
	r := NewRegistry()
	r.Register(&stubShell{name: "only"}, 1)

	for range 10 {
		sh, err := r.Select()
		if err != nil {
			t.Fatalf("Select: %v", err)
		}
		if sh.Name() != "only" {
			t.Errorf("Name = %q, want %q", sh.Name(), "only")
		}
	}
}

func TestRegistrySelectWeighted(t *testing.T) {
	r := NewRegistry()
	r.Register(&stubShell{name: "heavy"}, 100)
	r.Register(&stubShell{name: "light"}, 1)

	counts := map[string]int{}
	for range 1000 {
		sh, err := r.Select()
		if err != nil {
			t.Fatalf("Select: %v", err)
		}
		counts[sh.Name()]++
	}

	// "heavy" has weight 100 vs "light" weight 1, so heavy should get ~99%.
	if counts["heavy"] < 900 {
		t.Errorf("heavy selected %d/1000 times, expected >900", counts["heavy"])
	}
}