Pad/internal/io/pool/context_test.go

package pool

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

	"pad/internal/io/pool/mock"
)

// cancelTask is a task that respects context cancellation.
type cancelTask struct {
	id       string
	taskType TaskType
	priority Priority
	ctx      context.Context
	cancel   context.CancelFunc
	delay    time.Duration
	called   atomic.Bool
}

func newCancelTask(id string, delay time.Duration) *cancelTask {
	ctx, cancel := context.WithCancel(context.Background())
	return &cancelTask{
		id:       id,
		taskType: TypeReadFile,
		priority: HighPriority,
		ctx:      ctx,
		cancel:   cancel,
		delay:    delay,
	}
}

func (t *cancelTask) Execute() Result {
	t.called.Store(true)
	select {
	case <-t.ctx.Done():
		return Result{
			TaskID:   t.id,
			TaskType: t.taskType,
			Success:  false,
			Error:    t.ctx.Err(),
		}
	case <-time.After(t.delay):
		return Result{
			TaskID:   t.id,
			TaskType: t.taskType,
			Success:  true,
			Data:     "completed",
		}
	}
}

func (t *cancelTask) Priority() Priority       { return t.priority }
func (t *cancelTask) TaskID() string           { return t.id }
func (t *cancelTask) TaskType() TaskType       { return t.taskType }
func (t *cancelTask) DirPath() string          { return "/test" }
func (t *cancelTask) Context() context.Context { return t.ctx }
func (t *cancelTask) Cancel()                  { t.cancel() }
func (t *cancelTask) Timeout() time.Duration   { return 0 }

// timeoutTask is a task that respects timeout.
type timeoutTask struct {
	id       string
	taskType TaskType
	priority Priority
	delay    time.Duration
}

func (t *timeoutTask) Execute() Result {
	time.Sleep(t.delay)
	return Result{
		TaskID:   t.id,
		TaskType: t.taskType,
		Success:  true,
		Data:     "completed",
	}
}

func (t *timeoutTask) Priority() Priority       { return t.priority }
func (t *timeoutTask) TaskID() string           { return t.id }
func (t *timeoutTask) TaskType() TaskType       { return t.taskType }
func (t *timeoutTask) DirPath() string          { return "/test" }
func (t *timeoutTask) Context() context.Context { return context.Background() }
func (t *timeoutTask) Cancel()                  {}
func (t *timeoutTask) Timeout() time.Duration   { return 50 * time.Millisecond }

func TestTaskContextCancellation(t *testing.T) {
	task := newCancelTask("test-cancel", 1*time.Second)

	// Cancel the task
	task.Cancel()

	// Execute should return error
	result := task.Execute()
	if result.IsSuccess() {
		t.Error("Task should not succeed after cancellation")
	}
	if result.Error == nil {
		t.Error("Task should have error after cancellation")
	}
}

func TestTaskTimeoutMethod(t *testing.T) {
	task := &timeoutTask{id: "test-timeout"}
	if task.Timeout() != 50*time.Millisecond {
		t.Errorf("Timeout = %v, want %v", task.Timeout(), 50*time.Millisecond)
	}
}

func TestWorkerPool_ContextCancellation(t *testing.T) {
	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Create a task that will be cancelled
	task := newCancelTask("cancel-test", 500*time.Millisecond)

	// Dispatch the task
	pool.Dispatch(task)

	// Cancel after a short delay
	time.Sleep(50 * time.Millisecond)
	task.Cancel()

	// Wait for result
	select {
	case result := <-pool.resultChan:
		// Should complete with error due to cancellation
		if result.IsSuccess() {
			t.Log("Task completed before cancellation - acceptable")
		} else {
			t.Logf("Task cancelled as expected: %v", result.Error)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("Timed out waiting for result")
	}
}

func TestWorkerPool_TaskTimeout(t *testing.T) {
	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Create a task with a short timeout that will sleep longer
	task := &timeoutTask{
		id:       "timeout-test",
		taskType: TypeReadFile,
		priority: HighPriority,
		delay:    5 * time.Second, // Much longer than timeout
	}

	pool.Dispatch(task)

	// Should complete with timeout error
	select {
	case result := <-pool.resultChan:
		if result.IsSuccess() {
			t.Error("Task should have timed out")
		}
		if result.Error == nil {
			t.Error("Task should have error after timeout")
		}
		t.Logf("Task timed out as expected: %v", result.Error)
	case <-time.After(2 * time.Second):
		t.Fatal("Timed out waiting for result")
	}
}

func TestWorkerPool_CancelByDirPath(t *testing.T) {
	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Create tasks for different directories
	task1 := newCancelTask("dir1-task", 500*time.Millisecond)
	task2 := newCancelTask("dir2-task", 500*time.Millisecond)

	// Set DirPath for task1
	// Note: We need to modify the task to return specific DirPath
	// For now, we'll test the CancelPendingTasks method

	pool.Dispatch(task1)
	pool.Dispatch(task2)

	// Cancel tasks for "/test" directory (both tasks use this)
	pool.CancelPendingTasks("/test")

	// Wait for results
	for i := 0; i < 2; i++ {
		select {
		case result := <-pool.resultChan:
			t.Logf("Got result for task %s: success=%v, error=%v", result.TaskID, result.IsSuccess(), result.Error)
		case <-time.After(1 * time.Second):
			t.Fatalf("Timed out waiting for result %d", i)
		}
	}
}

func TestWorkerPool_CancelPendingTasks_NoEffect(t *testing.T) {
	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Cancel tasks for non-existent directory
	pool.CancelPendingTasks("/nonexistent")

	// Should not panic or cause issues
	if !pool.IsRunning() {
		t.Error("Pool should still be running")
	}
}

func TestWorkerPool_ContextPropagation(t *testing.T) {
	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Create a task that uses context
	task := newCancelTask("ctx-prop", 100*time.Millisecond)

	pool.Dispatch(task)

	// Wait for result
	select {
	case result := <-pool.resultChan:
		if !result.IsSuccess() {
			t.Errorf("Task failed: %v", result.Error)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("Timed out waiting for result")
	}
}

// TestWorkerPool_DirPathInResult verifies that the task's DirPath is correctly reflected in the result.
func TestWorkerPool_DirPathInResult(t *testing.T) {
	fs := mock.NewFileSystem()
	fs.CreateDir("/test/dir1")
	fs.CreateFile("/test/dir1/file.txt", []byte("content"))

	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Create a read dir task
	task := NewReadDirTask("/test/dir1", fs)
	pool.Dispatch(task)

	// Wait for result
	select {
	case result := <-pool.resultChan:
		if !result.IsSuccess() {
			t.Errorf("Task failed: %v", result.Error)
		}
		// DirPath should be in the result
		if result.DirPath != "/test/dir1" {
			t.Errorf("Result DirPath = %q, want %q", result.DirPath, "/test/dir1")
		}
	case <-time.After(1 * time.Second):
		t.Fatal("Timed out waiting for result")
	}
}

func TestWorkerPool_ConcurrentCancellation(t *testing.T) {
	pool := NewWorkerPool(4)
	pool.Start()
	defer pool.Stop()

	// Create multiple tasks
	var tasks []*cancelTask
	for i := 0; i < 10; i++ {
		task := newCancelTask(fmt.Sprintf("concurrent-%d", i), 1*time.Second)
		tasks = append(tasks, task)
		pool.Dispatch(task)
	}

	// Cancel half of them
	for i := 0; i < 5; i++ {
		tasks[i].Cancel()
	}

	// Wait for all results
	for i := 0; i < 10; i++ {
		select {
		case result := <-pool.resultChan:
			// Results may be success or error depending on timing
			t.Logf("Result for %s: success=%v", result.TaskID, result.IsSuccess())
		case <-time.After(2 * time.Second):
			t.Fatalf("Timed out waiting for result %d", i)
		}
	}
}

func TestWorkerPool_TimeoutWithRealTasks(t *testing.T) {
	fs := mock.NewFileSystem()
	fs.SetDelay(100 * time.Millisecond)
	fs.CreateFile("/test/timeout.txt", []byte("content"))

	pool := NewWorkerPool(2)
	pool.Start()
	defer pool.Stop()

	// Create a task with short timeout
	task := NewReadFileTask("/test/timeout.txt", fs)

	// Note: Real tasks don't have custom timeout yet
	// This test verifies that the timeout mechanism doesn't break existing functionality
	pool.Dispatch(task)

	select {
	case result := <-pool.resultChan:
		if !result.IsSuccess() {
			t.Errorf("Task failed: %v", result.Error)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Timed out waiting for result")
	}
}