package dispatch

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

// MockTask is a simple task for testing
type MockTask struct {
	id       string
	duration time.Duration
	shouldErr bool
}

func (mt *MockTask) ID() string {
	return mt.id
}

func (mt *MockTask) Execute(ctx context.Context) (interface{}, error) {
	select {
	case <-time.After(mt.duration):
		if mt.shouldErr {
			return nil, fmt.Errorf("task %s failed", mt.id)
		}
		return fmt.Sprintf("result-%s", mt.id), nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func TestNewDispatcher(t *testing.T) {
	dispatcher := NewDispatcher(5)
	assert.NotNil(t, dispatcher)
	assert.Equal(t, 5, dispatcher.maxConcurrency)
}

func TestDispatchSingleTask(t *testing.T) {
	dispatcher := NewDispatcher(1)

	task := &MockTask{
		id:       "task-1",
		duration: 10 * time.Millisecond,
		shouldErr: false,
	}

	results, err := dispatcher.DispatchAll(context.Background(), []Task{task})
	assert.NoError(t, err)
	assert.Equal(t, 1, len(results))

	result, exists := dispatcher.GetResult("task-1")
	assert.True(t, exists)
	assert.NoError(t, result.Error)
	assert.Equal(t, "result-task-1", result.Result)
}

func TestDispatchMultipleTasks(t *testing.T) {
	dispatcher := NewDispatcher(10)

	tasks := make([]Task, 0)
	for i := 1; i <= 5; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  10 * time.Millisecond,
			shouldErr: false,
		})
	}

	results, err := dispatcher.DispatchAll(context.Background(), tasks)
	assert.NoError(t, err)
	assert.Equal(t, 5, len(results))

	for i := 1; i <= 5; i++ {
		taskID := fmt.Sprintf("task-%d", i)
		result, exists := dispatcher.GetResult(taskID)
		assert.True(t, exists)
		assert.NoError(t, result.Error)
	}
}

func TestDispatchWithErrors(t *testing.T) {
	dispatcher := NewDispatcher(10)

	tasks := []Task{
		&MockTask{id: "task-1", duration: 10 * time.Millisecond, shouldErr: false},
		&MockTask{id: "task-2", duration: 10 * time.Millisecond, shouldErr: true},
		&MockTask{id: "task-3", duration: 10 * time.Millisecond, shouldErr: false},
	}

	results, _ := dispatcher.DispatchAll(context.Background(), tasks)
	// Errors don't prevent all tasks from completing
	assert.Equal(t, 3, len(results))

	result2, _ := dispatcher.GetResult("task-2")
	assert.Error(t, result2.Error)
}

func TestParallelExecution(t *testing.T) {
	dispatcher := NewDispatcher(10)

	// Create 9 tasks, each taking 100ms
	tasks := make([]Task, 0)
	for i := 1; i <= 9; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  100 * time.Millisecond,
			shouldErr: false,
		})
	}

	start := time.Now()
	results, err := dispatcher.DispatchAll(context.Background(), tasks)
	elapsed := time.Since(start)

	assert.NoError(t, err)
	assert.Equal(t, 9, len(results))

	// With parallel execution, should take ~100ms (not 900ms)
	// Allow some margin (150ms)
	assert.Less(t, elapsed, 150*time.Millisecond)
}

func TestSpeedup(t *testing.T) {
	dispatcher := NewDispatcher(10)

	tasks := make([]Task, 0)
	for i := 1; i <= 9; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  50 * time.Millisecond,
			shouldErr: false,
		})
	}

	_, _ = dispatcher.DispatchAll(context.Background(), tasks)

	speedup := dispatcher.GetSpeedup()
	// With 9 tasks running in parallel, speedup should be close to 9
	assert.Greater(t, speedup, 8.0)
	assert.Less(t, speedup, 10.0)
}

func TestExecutionTime(t *testing.T) {
	dispatcher := NewDispatcher(10)

	tasks := make([]Task, 0)
	for i := 1; i <= 3; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  100 * time.Millisecond,
			shouldErr: false,
		})
	}

	_, _ = dispatcher.DispatchAll(context.Background(), tasks)

	executionTime := dispatcher.GetExecutionTime()
	// Should be roughly 100ms (parallel execution)
	assert.Greater(t, executionTime, 80*time.Millisecond)
	assert.Less(t, executionTime, 200*time.Millisecond)
}

func TestTotalTaskDuration(t *testing.T) {
	dispatcher := NewDispatcher(10)

	tasks := make([]Task, 0)
	for i := 1; i <= 3; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  100 * time.Millisecond,
			shouldErr: false,
		})
	}

	_, _ = dispatcher.DispatchAll(context.Background(), tasks)

	totalDuration := dispatcher.GetTotalTaskDuration()
	// Sum should be roughly 300ms
	assert.Greater(t, totalDuration, 290*time.Millisecond)
	assert.Less(t, totalDuration, 350*time.Millisecond)
}

func TestGetStats(t *testing.T) {
	dispatcher := NewDispatcher(5)

	tasks := make([]Task, 0)
	for i := 1; i <= 5; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  50 * time.Millisecond,
			shouldErr: false,
		})
	}

	_, _ = dispatcher.DispatchAll(context.Background(), tasks)

	stats := dispatcher.GetStats()
	assert.Equal(t, 5, stats["total_tasks"])
	assert.Equal(t, 5, stats["completed"])
	assert.Equal(t, 5, stats["concurrency"])
	assert.NotZero(t, stats["total_duration"])
}

func TestIsComplete(t *testing.T) {
	dispatcher := NewDispatcher(1)

	task := &MockTask{
		id:        "task-1",
		duration:  10 * time.Millisecond,
		shouldErr: false,
	}

	dispatcher.DispatchAll(context.Background(), []Task{task})

	assert.True(t, dispatcher.IsComplete("task-1"))
	assert.False(t, dispatcher.IsComplete("task-2"))
}

func TestAreAllComplete(t *testing.T) {
	dispatcher := NewDispatcher(5)

	tasks := make([]Task, 0)
	for i := 1; i <= 3; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  10 * time.Millisecond,
			shouldErr: false,
		})
	}

	dispatcher.DispatchAll(context.Background(), tasks)

	assert.True(t, dispatcher.AreAllComplete())
}

func TestGetCompletedCount(t *testing.T) {
	dispatcher := NewDispatcher(5)

	tasks := make([]Task, 0)
	for i := 1; i <= 5; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  10 * time.Millisecond,
			shouldErr: false,
		})
	}

	dispatcher.DispatchAll(context.Background(), tasks)

	assert.Equal(t, 5, dispatcher.GetCompletedCount())
}

func TestConcurrencyLimit(t *testing.T) {
	// Create dispatcher with low concurrency
	dispatcher := NewDispatcher(2)

	// All tasks should still complete
	tasks := make([]Task, 0)
	for i := 1; i <= 5; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  10 * time.Millisecond,
			shouldErr: false,
		})
	}

	results, err := dispatcher.DispatchAll(context.Background(), tasks)
	assert.NoError(t, err)
	assert.Equal(t, 5, len(results))
}

func TestContextCancellation(t *testing.T) {
	dispatcher := NewDispatcher(2) // Low concurrency

	tasks := make([]Task, 0)
	for i := 1; i <= 10; i++ {
		tasks = append(tasks, &MockTask{
			id:        fmt.Sprintf("task-%d", i),
			duration:  500 * time.Millisecond,
			shouldErr: false,
		})
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(50 * time.Millisecond)
		cancel()
	}()

	_, _ = dispatcher.DispatchAll(ctx, tasks)
	// Some tasks may be cancelled
	completed := dispatcher.GetCompletedCount()
	assert.Less(t, completed, 10)
}

func TestEmptyTaskList(t *testing.T) {
	dispatcher := NewDispatcher(5)

	results, err := dispatcher.DispatchAll(context.Background(), []Task{})
	assert.NoError(t, err)
	assert.Equal(t, 0, len(results))
}

func TestTaskResultFields(t *testing.T) {
	dispatcher := NewDispatcher(1)

	task := &MockTask{
		id:        "task-1",
		duration:  50 * time.Millisecond,
		shouldErr: false,
	}

	dispatcher.DispatchAll(context.Background(), []Task{task})

	result, _ := dispatcher.GetResult("task-1")
	assert.NotZero(t, result.StartTime)
	assert.NotZero(t, result.EndTime)
	assert.NotZero(t, result.Duration)
	assert.True(t, result.EndTime.After(result.StartTime))
}

func BenchmarkParallelDispatch(b *testing.B) {
	dispatcher := NewDispatcher(10)

	for i := 0; i < b.N; i++ {
		tasks := make([]Task, 0)
		for j := 0; j < 10; j++ {
			tasks = append(tasks, &MockTask{
				id:        fmt.Sprintf("task-%d", j),
				duration:  5 * time.Millisecond,
				shouldErr: false,
			})
		}
		dispatcher.DispatchAll(context.Background(), tasks)
	}
}

func BenchmarkDispatchSingleTask(b *testing.B) {
	dispatcher := NewDispatcher(1)

	for i := 0; i < b.N; i++ {
		task := &MockTask{
			id:        "task-1",
			duration:  5 * time.Millisecond,
			shouldErr: false,
		}
		dispatcher.DispatchAll(context.Background(), []Task{task})
	}
}