package concurrency

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestLimiterGetLimit(t *testing.T) {
	tests := []struct {
		name string
		cfg  Config
		key  string
		want int
	}{
		{
			name: "exact_match",
			cfg: Config{
				Default: 2,
				Limits: []ConcurrencyLimit{
					{Key: "copilot/claude-sonnet-4.6", MaxConcurrency: 1},
				},
			},
			key:  "copilot/claude-sonnet-4.6",
			want: 1,
		},
		{
			name: "default",
			cfg: Config{
				Default: 2,
				Limits: []ConcurrencyLimit{
					{Key: "copilot/claude-sonnet-4.6", MaxConcurrency: 1},
				},
			},
			key:  "copilot/gpt-5.4",
			want: 2,
		},
		{
			name: "glob_match",
			cfg: Config{
				Default: 3,
				Limits: []ConcurrencyLimit{
					{Key: "copilot/*", MaxConcurrency: 5},
				},
			},
			key:  "copilot/claude-sonnet-4.6",
			want: 5,
		},
		{
			name: "glob_fallback_default",
			cfg: Config{
				Default: 3,
				Limits: []ConcurrencyLimit{
					{Key: "copilot/*", MaxConcurrency: 5},
				},
			},
			key:  "zhipu/glm-4.7",
			want: 3,
		},
		{
			name: "exact_priority_over_glob",
			cfg: Config{
				Default: 2,
				Limits: []ConcurrencyLimit{
					{Key: "copilot/*", MaxConcurrency: 10},
					{Key: "copilot/claude-sonnet-4.6", MaxConcurrency: 1},
				},
			},
			key:  "copilot/claude-sonnet-4.6",
			want: 1,
		},
		{
			name: "glob_after_exact_miss",
			cfg: Config{
				Default: 2,
				Limits: []ConcurrencyLimit{
					{Key: "copilot/*", MaxConcurrency: 10},
					{Key: "copilot/claude-sonnet-4.6", MaxConcurrency: 1},
				},
			},
			key:  "copilot/gpt-5.4",
			want: 10,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			l := NewLimiter(tt.cfg)
			if got := l.getLimit(tt.key); got != tt.want {
				t.Errorf("getLimit(%q) = %d, want %d", tt.key, got, tt.want)
			}
		})
	}
}

func TestLimiter_AcquireRelease(t *testing.T) {
	l := NewLimiter(Config{Default: 2})

	wait, release, err := l.Acquire(context.Background(), "test/model")
	if err != nil {
		t.Fatal(err)
	}
	if wait > time.Millisecond {
		t.Errorf("first acquire should have near-zero wait, got %v", wait)
	}
	release()

	wait, release, err = l.Acquire(context.Background(), "test/model")
	if err != nil {
		t.Fatal(err)
	}
	release()
}

func TestLimiter_ConcurrencyEnforced(t *testing.T) {
	const maxConcurrency = 2
	const workerCount = 10
	l := NewLimiter(Config{Default: maxConcurrency})

	var concurrent atomic.Int32
	var maxConcurrent atomic.Int32
	var wg sync.WaitGroup
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	ready := make(chan struct{}, workerCount)
	start := make(chan struct{})
	entered := make(chan struct{}, workerCount)
	leave := make(chan struct{})
	workerErr := make(chan error, workerCount)
	started := false
	leaveClosed := false
	fail := func(format string, args ...any) {
		cancel()
		if !started {
			close(start)
			started = true
		}
		if !leaveClosed {
			close(leave)
			leaveClosed = true
		}
		wg.Wait()
		t.Fatalf(format, args...)
	}

	for i := 0; i < workerCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ready <- struct{}{}
			<-start

			_, release, err := l.Acquire(ctx, "test/model")
			if err != nil {
				workerErr <- err
				return
			}
			defer release()

			cur := concurrent.Add(1)
			for {
				old := maxConcurrent.Load()
				if cur <= old || maxConcurrent.CompareAndSwap(old, cur) {
					break
				}
			}
			entered <- struct{}{}
			<-leave
			concurrent.Add(-1)
		}()
	}

	for i := 0; i < workerCount; i++ {
		select {
		case <-ready:
		case err := <-workerErr:
			fail("worker failed before start: %v", err)
		case <-ctx.Done():
			fail("timed out waiting for workers to become ready: %v", ctx.Err())
		}
	}
	close(start)
	started = true
	for remaining := workerCount; remaining > 0; {
		batch := maxConcurrency
		if remaining < batch {
			batch = remaining
		}
		for i := 0; i < batch; i++ {
			select {
			case <-entered:
			case err := <-workerErr:
				fail("worker failed to acquire limiter slot: %v", err)
			case <-ctx.Done():
				fail("timed out waiting for %d workers to enter concurrency batch; max observed %d: %v", batch, maxConcurrent.Load(), ctx.Err())
			}
		}
		for i := 0; i < batch; i++ {
			leave <- struct{}{}
		}
		remaining -= batch
	}
	wg.Wait()

	if got := maxConcurrent.Load(); got > maxConcurrency {
		t.Errorf("max concurrency exceeded: got %d, limit %d", got, maxConcurrency)
	}
}

func TestLimiter_ContextCancellation(t *testing.T) {
	l := NewLimiter(Config{Default: 1})

	_, release, err := l.Acquire(context.Background(), "test/model")
	if err != nil {
		t.Fatal(err)
	}
	defer release()

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, _, err = l.Acquire(ctx, "test/model")
	if err == nil {
		t.Error("expected context cancellation error")
	}
}

func TestLimiter_Update(t *testing.T) {
	l := NewLimiter(Config{Default: 2})
	if got := l.getLimit("any/model"); got != 2 {
		t.Errorf("expected default 2, got %d", got)
	}

	l.Update(Config{
		Default: 5,
		Limits: []ConcurrencyLimit{
			{Key: "test/*", MaxConcurrency: 3},
		},
	})

	if got := l.getLimit("test/model"); got != 3 {
		t.Errorf("expected updated limit 3, got %d", got)
	}
	if got := l.getLimit("other/model"); got != 5 {
		t.Errorf("expected updated default 5, got %d", got)
	}
}