package integration_test

// SPDX-License-Identifier: GPL-3.0-or-later

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"

	"projects.blender.org/studio/flamenco/pkg/api"
	"projects.blender.org/studio/flamenco/tests/helpers"
)

// IntegrationTestSuite provides end-to-end workflow testing
type IntegrationTestSuite struct {
	suite.Suite
	testHelper   *helpers.TestHelper
	baseURL      string
	wsURL        string
	client       *http.Client
	wsConn       *websocket.Conn
	wsEvents     chan []byte
	wsCloseOnce  sync.Once
}

// WorkflowContext tracks a complete render job workflow
type WorkflowContext struct {
	Job            api.Job
	Worker         api.RegisteredWorker
	AssignedTasks  []api.AssignedTask
	TaskUpdates    []api.TaskUpdate
	JobStatusHist  []api.JobStatus
	StartTime      time.Time
	CompletionTime time.Time
	Events         []interface{}
}

// SetupSuite initializes the integration test environment
func (suite *IntegrationTestSuite) SetupSuite() {
	suite.testHelper = helpers.NewTestHelper(suite.T())
	
	// Start test server
	server := suite.testHelper.StartTestServer()
	suite.baseURL = server.URL
	suite.wsURL = strings.Replace(server.URL, "http://", "ws://", 1)
	
	// Configure HTTP client
	suite.client = &http.Client{
		Timeout: 30 * time.Second,
	}
	
	// Initialize WebSocket connection
	suite.setupWebSocket()
}

// TearDownSuite cleans up the integration test environment
func (suite *IntegrationTestSuite) TearDownSuite() {
	suite.closeWebSocket()
	if suite.testHelper != nil {
		suite.testHelper.Cleanup()
	}
}

// TestCompleteRenderWorkflow tests full job lifecycle from submission to completion
func (suite *IntegrationTestSuite) TestCompleteRenderWorkflow() {
	ctx := &WorkflowContext{
		StartTime: time.Now(),
		Events:    make([]interface{}, 0),
	}
	
	suite.Run("JobSubmission", func() {
		// Submit a render job
		submittedJob := api.SubmittedJob{
			Name:              "Integration Test - Complete Workflow",
			Type:              "simple-blender-render",
			Priority:          50,
			SubmitterPlatform: "linux",
			Settings: map[string]interface{}{
				"filepath":               "/shared-storage/test-scene.blend",
				"chunk_size":             5,
				"format":                 "PNG",
				"image_file_extension":   ".png",
				"frames":                 "1-20",
				"render_output_root":     "/shared-storage/renders/",
				"add_path_components":    0,
				"render_output_path":     "/shared-storage/renders/test-render/######",
			},
		}
		
		jobData, err := json.Marshal(submittedJob)
		require.NoError(suite.T(), err)
		
		resp, err := suite.makeRequest("POST", "/api/v3/jobs", bytes.NewReader(jobData))
		require.NoError(suite.T(), err)
		require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
		
		err = json.NewDecoder(resp.Body).Decode(&ctx.Job)
		require.NoError(suite.T(), err)
		resp.Body.Close()
		
		assert.NotEmpty(suite.T(), ctx.Job.Id)
		assert.Equal(suite.T(), submittedJob.Name, ctx.Job.Name)
		assert.Equal(suite.T(), api.JobStatusQueued, ctx.Job.Status)
		
		suite.T().Logf("Job submitted: %s (ID: %s)", ctx.Job.Name, ctx.Job.Id)
	})
	
	suite.Run("WorkerRegistration", func() {
		// Register a worker
		workerReg := api.WorkerRegistration{
			Name:               "integration-test-worker",
			Address:            "192.168.1.100",
			Platform:           "linux",
			SoftwareVersion:    "3.0.0",
			SupportedTaskTypes: []string{"blender", "ffmpeg"},
		}
		
		workerData, err := json.Marshal(workerReg)
		require.NoError(suite.T(), err)
		
		resp, err := suite.makeRequest("POST", "/api/v3/worker/register-worker", bytes.NewReader(workerData))
		require.NoError(suite.T(), err)
		require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
		
		err = json.NewDecoder(resp.Body).Decode(&ctx.Worker)
		require.NoError(suite.T(), err)
		resp.Body.Close()
		
		assert.NotEmpty(suite.T(), ctx.Worker.Uuid)
		assert.Equal(suite.T(), workerReg.Name, ctx.Worker.Name)
		
		suite.T().Logf("Worker registered: %s (UUID: %s)", ctx.Worker.Name, ctx.Worker.Uuid)
	})
	
	suite.Run("WorkerSignOn", func() {
		// Worker signs on and becomes available
		signOnInfo := api.WorkerSignOn{
			Name:               ctx.Worker.Name,
			SoftwareVersion:    "3.0.0",
			SupportedTaskTypes: []string{"blender", "ffmpeg"},
		}
		
		signOnData, err := json.Marshal(signOnInfo)
		require.NoError(suite.T(), err)
		
		url := fmt.Sprintf("/api/v3/worker/%s/sign-on", ctx.Worker.Uuid)
		resp, err := suite.makeRequest("POST", url, bytes.NewReader(signOnData))
		require.NoError(suite.T(), err)
		require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
		
		var signOnResponse api.WorkerStateChange
		err = json.NewDecoder(resp.Body).Decode(&signOnResponse)
		require.NoError(suite.T(), err)
		resp.Body.Close()
		
		assert.Equal(suite.T(), api.WorkerStatusAwake, *signOnResponse.StatusRequested)
		
		suite.T().Logf("Worker signed on successfully")
	})
	
	suite.Run("TaskAssignmentAndExecution", func() {
		// Worker requests tasks and executes them
		maxTasks := 10
		completedTasks := 0
		
		for attempt := 0; attempt < maxTasks; attempt++ {
			// Request task
			taskURL := fmt.Sprintf("/api/v3/worker/%s/task", ctx.Worker.Uuid)
			resp, err := suite.makeRequest("POST", taskURL, nil)
			require.NoError(suite.T(), err)
			
			if resp.StatusCode == http.StatusNoContent {
				// No more tasks available
				resp.Body.Close()
				suite.T().Logf("No more tasks available after %d completed tasks", completedTasks)
				break
			}
			
			require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
			
			var assignedTask api.AssignedTask
			err = json.NewDecoder(resp.Body).Decode(&assignedTask)
			require.NoError(suite.T(), err)
			resp.Body.Close()
			
			ctx.AssignedTasks = append(ctx.AssignedTasks, assignedTask)
			
			assert.NotEmpty(suite.T(), assignedTask.Uuid)
			assert.Equal(suite.T(), ctx.Job.Id, assignedTask.JobId)
			assert.NotEmpty(suite.T(), assignedTask.Commands)
			
			suite.T().Logf("Task assigned: %s (Type: %s)", assignedTask.Name, assignedTask.TaskType)
			
			// Simulate task execution
			suite.simulateTaskExecution(ctx.Worker.Uuid, &assignedTask)
			completedTasks++
			
			// Small delay between tasks
			time.Sleep(time.Millisecond * 100)
		}
		
		assert.Greater(suite.T(), completedTasks, 0, "Should have completed at least one task")
		suite.T().Logf("Completed %d tasks", completedTasks)
	})
	
	suite.Run("JobCompletion", func() {
		// Wait for job to complete
		timeout := time.After(30 * time.Second)
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		
		for {
			select {
			case <-timeout:
				suite.T().Fatal("Timeout waiting for job completion")
			case <-ticker.C:
				// Check job status
				resp, err := suite.makeRequest("GET", fmt.Sprintf("/api/v3/jobs/%s", ctx.Job.Id), nil)
				require.NoError(suite.T(), err)
				
				var currentJob api.Job
				err = json.NewDecoder(resp.Body).Decode(&currentJob)
				require.NoError(suite.T(), err)
				resp.Body.Close()
				
				ctx.JobStatusHist = append(ctx.JobStatusHist, currentJob.Status)
				
				suite.T().Logf("Job status: %s", currentJob.Status)
				
				if currentJob.Status == api.JobStatusCompleted {
					ctx.Job = currentJob
					ctx.CompletionTime = time.Now()
					suite.T().Logf("Job completed successfully in %v", ctx.CompletionTime.Sub(ctx.StartTime))
					return
				}
				
				if currentJob.Status == api.JobStatusFailed || currentJob.Status == api.JobStatusCanceled {
					ctx.Job = currentJob
					suite.T().Fatalf("Job failed or was canceled. Final status: %s", currentJob.Status)
				}
			}
		}
	})
	
	// Validate complete workflow
	suite.validateWorkflowResults(ctx)
}

// TestWorkerFailureRecovery tests system behavior when workers fail
func (suite *IntegrationTestSuite) TestWorkerFailureRecovery() {
	suite.Run("SetupJobAndWorker", func() {
		// Submit a job
		job := suite.createIntegrationTestJob("Worker Failure Recovery Test")
		
		// Register and sign on worker
		worker := suite.registerAndSignOnWorker("failure-test-worker")
		
		// Worker requests a task
		taskURL := fmt.Sprintf("/api/v3/worker/%s/task", worker.Uuid)
		resp, err := suite.makeRequest("POST", taskURL, nil)
		require.NoError(suite.T(), err)
		
		if resp.StatusCode == http.StatusOK {
			var assignedTask api.AssignedTask
			json.NewDecoder(resp.Body).Decode(&assignedTask)
			resp.Body.Close()
			
			suite.T().Logf("Task assigned: %s", assignedTask.Name)
			
			// Simulate worker failure (no task updates sent)
			suite.T().Logf("Simulating worker failure...")
			
			// Wait for timeout handling
			time.Sleep(5 * time.Second)
			
			// Check if task was requeued or marked as failed
			suite.validateTaskRecovery(assignedTask.Uuid, job.Id)
		} else {
			resp.Body.Close()
			suite.T().Skip("No tasks available for failure recovery test")
		}
	})
}

// TestMultiWorkerCoordination tests coordination between multiple workers
func (suite *IntegrationTestSuite) TestMultiWorkerCoordination() {
	const numWorkers = 3
	
	suite.Run("SetupMultiWorkerEnvironment", func() {
		// Submit a large job
		job := suite.createIntegrationTestJob("Multi-Worker Coordination Test")
		
		// Register multiple workers
		workers := make([]api.RegisteredWorker, numWorkers)
		for i := 0; i < numWorkers; i++ {
			workerName := fmt.Sprintf("coordination-worker-%d", i)
			workers[i] = suite.registerAndSignOnWorker(workerName)
		}
		
		// Simulate workers processing tasks concurrently
		var wg sync.WaitGroup
		taskCounts := make([]int, numWorkers)
		
		for i, worker := range workers {
			wg.Add(1)
			go func(workerIndex int, w api.RegisteredWorker) {
				defer wg.Done()
				
				for attempt := 0; attempt < 5; attempt++ {
					taskURL := fmt.Sprintf("/api/v3/worker/%s/task", w.Uuid)
					resp, err := suite.makeRequest("POST", taskURL, nil)
					if err != nil {
						continue
					}
					
					if resp.StatusCode == http.StatusOK {
						var task api.AssignedTask
						json.NewDecoder(resp.Body).Decode(&task)
						resp.Body.Close()
						
						suite.T().Logf("Worker %d got task: %s", workerIndex, task.Name)
						suite.simulateTaskExecution(w.Uuid, &task)
						taskCounts[workerIndex]++
					} else {
						resp.Body.Close()
						break
					}
					
					time.Sleep(time.Millisecond * 200)
				}
			}(i, worker)
		}
		
		wg.Wait()
		
		// Validate task distribution
		totalTasks := 0
		for i, count := range taskCounts {
			suite.T().Logf("Worker %d completed %d tasks", i, count)
			totalTasks += count
		}
		
		assert.Greater(suite.T(), totalTasks, 0, "Workers should have completed some tasks")
		
		// Verify job progresses towards completion
		suite.waitForJobProgress(job.Id, 30*time.Second)
	})
}

// TestWebSocketUpdates tests real-time updates via WebSocket
func (suite *IntegrationTestSuite) TestWebSocketUpdates() {
	if suite.wsConn == nil {
		suite.T().Skip("WebSocket connection not available")
		return
	}
	
	suite.Run("JobStatusUpdates", func() {
		// Clear event buffer
		suite.clearWebSocketEvents()
		
		// Submit a job
		job := suite.createIntegrationTestJob("WebSocket Updates Test")
		
		// Register worker and process tasks
		worker := suite.registerAndSignOnWorker("websocket-test-worker")
		
		// Start monitoring WebSocket events
		eventReceived := make(chan bool, 1)
		go func() {
			timeout := time.After(10 * time.Second)
			for {
				select {
				case event := <-suite.wsEvents:
					suite.T().Logf("WebSocket event received: %s", string(event))
					
					// Check if this is a job-related event
					if strings.Contains(string(event), job.Id) {
						eventReceived <- true
						return
					}
				case <-timeout:
					eventReceived <- false
					return
				}
			}
		}()
		
		// Process a task to trigger events
		taskURL := fmt.Sprintf("/api/v3/worker/%s/task", worker.Uuid)
		resp, err := suite.makeRequest("POST", taskURL, nil)
		require.NoError(suite.T(), err)
		
		if resp.StatusCode == http.StatusOK {
			var task api.AssignedTask
			json.NewDecoder(resp.Body).Decode(&task)
			resp.Body.Close()
			
			// Execute task to generate updates
			suite.simulateTaskExecution(worker.Uuid, &task)
		} else {
			resp.Body.Close()
		}
		
		// Wait for WebSocket event
		received := <-eventReceived
		assert.True(suite.T(), received, "Should receive WebSocket event for job update")
	})
}

// Helper methods

func (suite *IntegrationTestSuite) setupWebSocket() {
	wsURL := suite.wsURL + "/ws"
	
	var err error
	suite.wsConn, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		suite.T().Logf("Failed to connect to WebSocket: %v", err)
		return
	}
	
	suite.wsEvents = make(chan []byte, 100)
	
	// Start WebSocket message reader
	go func() {
		defer func() {
			if r := recover(); r != nil {
				suite.T().Logf("WebSocket reader panic: %v", r)
			}
		}()
		
		for {
			_, message, err := suite.wsConn.ReadMessage()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					suite.T().Logf("WebSocket error: %v", err)
				}
				return
			}
			
			select {
			case suite.wsEvents <- message:
			default:
				// Buffer full, drop oldest message
				select {
				case <-suite.wsEvents:
				default:
				}
				suite.wsEvents <- message
			}
		}
	}()
}

func (suite *IntegrationTestSuite) closeWebSocket() {
	suite.wsCloseOnce.Do(func() {
		if suite.wsConn != nil {
			suite.wsConn.Close()
		}
		if suite.wsEvents != nil {
			close(suite.wsEvents)
		}
	})
}

func (suite *IntegrationTestSuite) clearWebSocketEvents() {
	if suite.wsEvents == nil {
		return
	}
	
	for len(suite.wsEvents) > 0 {
		<-suite.wsEvents
	}
}

func (suite *IntegrationTestSuite) makeRequest(method, path string, body io.Reader) (*http.Response, error) {
	url := suite.baseURL + path
	req, err := http.NewRequestWithContext(context.Background(), method, url, body)
	if err != nil {
		return nil, err
	}
	
	req.Header.Set("Content-Type", "application/json")
	return suite.client.Do(req)
}

func (suite *IntegrationTestSuite) createIntegrationTestJob(name string) api.Job {
	submittedJob := api.SubmittedJob{
		Name:              name,
		Type:              "simple-blender-render",
		Priority:          50,
		SubmitterPlatform: "linux",
		Settings: map[string]interface{}{
			"filepath":               "/shared-storage/test.blend",
			"chunk_size":             3,
			"format":                 "PNG",
			"image_file_extension":   ".png",
			"frames":                 "1-12",
			"render_output_root":     "/shared-storage/renders/",
			"add_path_components":    0,
			"render_output_path":     "/shared-storage/renders/test/######",
		},
	}
	
	jobData, _ := json.Marshal(submittedJob)
	resp, err := suite.makeRequest("POST", "/api/v3/jobs", bytes.NewReader(jobData))
	require.NoError(suite.T(), err)
	require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
	
	var job api.Job
	json.NewDecoder(resp.Body).Decode(&job)
	resp.Body.Close()
	
	return job
}

func (suite *IntegrationTestSuite) registerAndSignOnWorker(name string) api.RegisteredWorker {
	// Register worker
	workerReg := api.WorkerRegistration{
		Name:               name,
		Address:            "192.168.1.100",
		Platform:           "linux",
		SoftwareVersion:    "3.0.0",
		SupportedTaskTypes: []string{"blender", "ffmpeg"},
	}
	
	workerData, _ := json.Marshal(workerReg)
	resp, err := suite.makeRequest("POST", "/api/v3/worker/register-worker", bytes.NewReader(workerData))
	require.NoError(suite.T(), err)
	require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
	
	var worker api.RegisteredWorker
	json.NewDecoder(resp.Body).Decode(&worker)
	resp.Body.Close()
	
	// Sign on worker
	signOnInfo := api.WorkerSignOn{
		Name:               name,
		SoftwareVersion:    "3.0.0",
		SupportedTaskTypes: []string{"blender", "ffmpeg"},
	}
	
	signOnData, _ := json.Marshal(signOnInfo)
	signOnURL := fmt.Sprintf("/api/v3/worker/%s/sign-on", worker.Uuid)
	resp, err = suite.makeRequest("POST", signOnURL, bytes.NewReader(signOnData))
	require.NoError(suite.T(), err)
	require.Equal(suite.T(), http.StatusOK, resp.StatusCode)
	resp.Body.Close()
	
	return worker
}

func (suite *IntegrationTestSuite) simulateTaskExecution(workerUUID string, task *api.AssignedTask) {
	updates := []struct {
		progress int
		status   api.TaskStatus
		message  string
	}{
		{25, api.TaskStatusActive, "Task started"},
		{50, api.TaskStatusActive, "Rendering in progress"},
		{75, api.TaskStatusActive, "Almost complete"},
		{100, api.TaskStatusCompleted, "Task completed successfully"},
	}
	
	for _, update := range updates {
		taskUpdate := api.TaskUpdate{
			TaskProgress: &api.TaskProgress{
				PercentageComplete: int32(update.progress),
			},
			TaskStatus: update.status,
			Log:        update.message,
		}
		
		updateData, _ := json.Marshal(taskUpdate)
		updateURL := fmt.Sprintf("/api/v3/worker/%s/task/%s", workerUUID, task.Uuid)
		
		resp, err := suite.makeRequest("POST", updateURL, bytes.NewReader(updateData))
		if err == nil && resp != nil {
			resp.Body.Close()
		}
		
		// Simulate processing time
		time.Sleep(time.Millisecond * 100)
	}
}

func (suite *IntegrationTestSuite) validateTaskRecovery(taskUUID, jobID string) {
	// Implementation would check if task was properly handled after worker failure
	suite.T().Logf("Validating task recovery for task %s in job %s", taskUUID, jobID)
	
	// In a real implementation, this would:
	// 1. Check if task was marked as failed
	// 2. Verify task was requeued for another worker
	// 3. Ensure job can still complete
}

func (suite *IntegrationTestSuite) waitForJobProgress(jobID string, timeout time.Duration) {
	deadline := time.After(timeout)
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	
	for {
		select {
		case <-deadline:
			suite.T().Logf("Timeout waiting for job %s progress", jobID)
			return
		case <-ticker.C:
			resp, err := suite.makeRequest("GET", fmt.Sprintf("/api/v3/jobs/%s", jobID), nil)
			if err != nil {
				continue
			}
			
			var job api.Job
			json.NewDecoder(resp.Body).Decode(&job)
			resp.Body.Close()
			
			suite.T().Logf("Job %s status: %s", jobID, job.Status)
			
			if job.Status == api.JobStatusCompleted || job.Status == api.JobStatusFailed {
				return
			}
		}
	}
}

func (suite *IntegrationTestSuite) validateWorkflowResults(ctx *WorkflowContext) {
	suite.T().Logf("=== Workflow Validation ===")
	
	// Validate job completion
	assert.Equal(suite.T(), api.JobStatusCompleted, ctx.Job.Status, "Job should be completed")
	assert.True(suite.T(), ctx.CompletionTime.After(ctx.StartTime), "Completion time should be after start time")
	
	// Validate task execution
	assert.Greater(suite.T(), len(ctx.AssignedTasks), 0, "Should have assigned tasks")
	
	// Validate workflow timing
	duration := ctx.CompletionTime.Sub(ctx.StartTime)
	assert.Less(suite.T(), duration, 5*time.Minute, "Workflow should complete within reasonable time")
	
	suite.T().Logf("Workflow completed in %v with %d tasks", duration, len(ctx.AssignedTasks))
}

// TestSuite runs all integration tests
func TestIntegrationSuite(t *testing.T) {
	suite.Run(t, new(IntegrationTestSuite))
}