package scheduler_test

import (
	"testing"
	"time"

	"github.com/jfraeys/fetch_ml/internal/scheduler"
	fixtures "github.com/jfraeys/fetch_ml/tests/fixtures"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestSchedulerRestartRecovery validates state persistence across restarts
func TestSchedulerRestartRecovery(t *testing.T) {
	tempDir := t.TempDir()

	// Phase 1: Create scheduler, submit jobs, assign to workers
	cfg1 := fixtures.DefaultHubConfig()
	cfg1.StateDir = tempDir

	hub1, err := scheduler.NewHub(cfg1, nil)
	require.NoError(t, err)

	err = hub1.Start()
	require.NoError(t, err)

	// Create worker and register
	worker1 := fixtures.NewMockWorker(t, hub1, "worker-restart-1")
	worker1.Register(scheduler.WorkerCapabilities{GPUCount: 0})

	// Submit a job
	jobID := "persistent-job-001"
	err = hub1.SubmitJob(scheduler.JobSpec{
		ID:      jobID,
		Type:    scheduler.JobTypeBatch,
		Command: []string{"python", "train.py"},
	})
	require.NoError(t, err)

	// Worker signals ready
	worker1.SignalReady(scheduler.SlotStatus{BatchTotal: 4, BatchInUse: 0}, "polling")

	// Wait for job assignment
	msg := worker1.RecvTimeout(2 * time.Second)
	require.Equal(t, scheduler.MsgJobAssign, msg.Type)

	// Accept the job
	worker1.AcceptJob(jobID)

	// Give time for state to be written
	time.Sleep(100 * time.Millisecond)

	// Stop scheduler (simulating restart)
	hub1.Stop()
	worker1.Close()

	// Phase 2: Create new scheduler instance, restore state
	hub2, err := scheduler.NewHub(cfg1, nil)
	require.NoError(t, err)
	defer hub2.Stop()

	// Start should replay state
	err = hub2.Start()
	require.NoError(t, err)

	// Verify job state was restored
	task := hub2.GetTask(jobID)
	require.NotNil(t, task, "task should be restored after restart")
	assert.Equal(t, jobID, task.ID)

	t.Log("Scheduler restart recovery test passed - state persisted correctly")
}

// TestSchedulerRestartWithPendingJobs validates jobs survive restart
func TestSchedulerRestartWithPendingJobs(t *testing.T) {
	tempDir := t.TempDir()

	cfg := fixtures.DefaultHubConfig()
	cfg.StateDir = tempDir

	// Phase 1: Create scheduler and submit jobs (no workers)
	hub1, err := scheduler.NewHub(cfg, nil)
	require.NoError(t, err)

	err = hub1.Start()
	require.NoError(t, err)

	// Submit multiple jobs
	jobIDs := []string{"pending-1", "pending-2", "pending-3"}
	for _, id := range jobIDs {
		err = hub1.SubmitJob(scheduler.JobSpec{
			ID:      id,
			Type:    scheduler.JobTypeBatch,
			Command: []string{"python", "train.py"},
		})
		require.NoError(t, err)
	}

	// Stop scheduler
	hub1.Stop()

	// Phase 2: Restart and verify jobs are still pending
	hub2, err := scheduler.NewHub(cfg, nil)
	require.NoError(t, err)
	defer hub2.Stop()

	err = hub2.Start()
	require.NoError(t, err)

	// Verify all jobs are still in the queue
	for _, id := range jobIDs {
		task := hub2.GetTask(id)
		require.NotNil(t, task, "pending job %s should survive restart", id)
		assert.Equal(t, id, task.ID)
	}

	t.Log("Pending jobs survived restart correctly")
}

// TestModeSwitchBatchToService validates switching worker mode
func TestModeSwitchBatchToService(t *testing.T) {
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create worker in batch mode
	worker := fixture.CreateWorker("mode-switch-worker", scheduler.WorkerCapabilities{
		GPUCount: 0,
	})

	// Submit batch job
	batchJobID := "batch-job-001"
	fixture.SubmitJob(scheduler.JobSpec{
		ID:       batchJobID,
		Type:     scheduler.JobTypeBatch,
		SlotPool: "batch",
		Command:  []string{"python", "train.py"},
	})

	// Signal ready for batch
	worker.SignalReady(scheduler.SlotStatus{BatchTotal: 4, BatchInUse: 0}, "polling")

	// Should receive batch job
	msg := worker.RecvTimeout(2 * time.Second)
	require.Equal(t, scheduler.MsgJobAssign, msg.Type)
	worker.AcceptJob(batchJobID)

	// Complete batch job
	worker.CompleteJob(batchJobID, 0, "")

	// Submit service job
	serviceJobID := "service-job-001"
	fixture.SubmitJob(scheduler.JobSpec{
		ID:       serviceJobID,
		Type:     scheduler.JobTypeService,
		SlotPool: "service",
		Command:  []string{"python", "-m", "http.server", "8080"},
	})

	// Switch to service mode
	worker.SignalReady(scheduler.SlotStatus{ServiceTotal: 4, ServiceInUse: 0}, "mode_switch")

	// Should receive service job
	msg2 := worker.RecvTimeout(2 * time.Second)
	require.Equal(t, scheduler.MsgJobAssign, msg2.Type)

	t.Log("Worker successfully switched from batch to service mode")
}

// TestModeSwitchServiceToBatch validates switching worker mode back
func TestModeSwitchServiceToBatch(t *testing.T) {
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create worker
	worker := fixture.CreateWorker("mode-switch-worker-2", scheduler.WorkerCapabilities{
		GPUCount: 0,
	})

	// Submit service job first
	serviceJobID := "service-job-002"
	fixture.SubmitJob(scheduler.JobSpec{
		ID:       serviceJobID,
		Type:     scheduler.JobTypeService,
		SlotPool: "service",
		Command:  []string{"python", "-m", "http.server", "8080"},
	})

	// Start in service mode
	worker.SignalReady(scheduler.SlotStatus{ServiceTotal: 4, ServiceInUse: 0}, "polling")

	// Should receive service job
	msg := worker.RecvTimeout(2 * time.Second)
	require.Equal(t, scheduler.MsgJobAssign, msg.Type)
	worker.AcceptJob(serviceJobID)

	// Submit batch job
	batchJobID := "batch-job-002"
	fixture.SubmitJob(scheduler.JobSpec{
		ID:       batchJobID,
		Type:     scheduler.JobTypeBatch,
		SlotPool: "batch",
		Command:  []string{"python", "train.py"},
	})

	// Switch to batch mode
	worker.SignalReady(scheduler.SlotStatus{BatchTotal: 4, BatchInUse: 0}, "mode_switch")

	// Should receive batch job
	msg2 := worker.RecvTimeout(2 * time.Second)
	require.Equal(t, scheduler.MsgJobAssign, msg2.Type)

	t.Log("Worker successfully switched from service to batch mode")
}

// TestEndToEndJobLifecycle full workflow from submission to completion
func TestEndToEndJobLifecycle(t *testing.T) {
	fixture := fixtures.NewSchedulerTestFixture(t, fixtures.DefaultHubConfig())
	defer fixture.Cleanup()

	// Create multiple workers
	worker1 := fixture.CreateWorker("e2e-worker-1", scheduler.WorkerCapabilities{GPUCount: 0})
	worker2 := fixture.CreateWorker("e2e-worker-2", scheduler.WorkerCapabilities{GPUCount: 0})

	// Submit batch job
	jobID := "e2e-job-001"
	fixture.SubmitJob(scheduler.JobSpec{
		ID:      jobID,
		Type:    scheduler.JobTypeBatch,
		Command: []string{"python", "train.py"},
	})

	// Both workers signal ready
	worker1.SignalReady(scheduler.SlotStatus{BatchTotal: 4, BatchInUse: 0}, "polling")
	worker2.SignalReady(scheduler.SlotStatus{BatchTotal: 4, BatchInUse: 0}, "polling")

	// One worker should receive the job
	var assignedWorker *fixtures.MockWorker
	select {
	case msg := <-worker1.RecvCh:
		if msg.Type == scheduler.MsgJobAssign {
			assignedWorker = worker1
		}
	case msg := <-worker2.RecvCh:
		if msg.Type == scheduler.MsgJobAssign {
			assignedWorker = worker2
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for job assignment")
	}

	require.NotNil(t, assignedWorker, "one worker should receive the job")

	// Accept and complete the job
	assignedWorker.AcceptJob(jobID)
	assignedWorker.CompleteJob(jobID, 0, "Job completed successfully")

	// Verify task state
	task := fixture.GetTask(jobID)
	require.NotNil(t, task)

	t.Log("End-to-end job lifecycle test passed")
}

// TestSchedulerPersistenceWithCrash simulates crash and recovery
func TestSchedulerPersistenceWithCrash(t *testing.T) {
	tempDir := t.TempDir()

	// Create initial scheduler
	cfg := fixtures.DefaultHubConfig()
	cfg.StateDir = tempDir

	hub, err := scheduler.NewHub(cfg, nil)
	require.NoError(t, err)

	err = hub.Start()
	require.NoError(t, err)

	// Submit job
	jobID := "crash-test-job"
	err = hub.SubmitJob(scheduler.JobSpec{
		ID:      jobID,
		Type:    scheduler.JobTypeBatch,
		Command: []string{"python", "train.py"},
	})
	require.NoError(t, err)

	// Simulate crash by not calling Stop()
	// In real crash, state file should still have the job
	time.Sleep(100 * time.Millisecond)

	// Simulate crash - force close without cleanup
	// Note: In real implementation, we'd use syscall to kill process
	// For testing, we'll just create a new hub with same state
	hub = nil

	// Create new scheduler with same state dir
	hub2, err := scheduler.NewHub(cfg, nil)
	require.NoError(t, err)
	defer hub2.Stop()

	// Start should recover state
	err = hub2.Start()
	require.NoError(t, err)

	// Verify job survived
	task := hub2.GetTask(jobID)
	require.NotNil(t, task, "job should survive simulated crash")

	t.Log("Scheduler persistence with crash test passed")
}
