package task

import (
	"errors"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/tidbcloud/tidb-management-service/tidb-x-migration/internal/config"
	"github.com/tidbcloud/tidb-management-service/tidb-x-migration/internal/export"
	"github.com/tidbcloud/tidb-management-service/tidb-x-migration/internal/importer"
	"github.com/tidbcloud/tidb-management-service/tidb-x-migration/internal/storage"

	"github.com/pingcap/log"
	"go.uber.org/zap"
)

// Task represents a migration task
type Task struct {
	ID               string             `json:"id"`
	Status           storage.TaskStatus `json:"status"`
	ConnectionOrigin string             `json:"connectionOrigin"`
	ConnectionTarget string             `json:"connectionTarget"`
	CreatedAt        string             `json:"createdAt"`
	UpdatedAt        string             `json:"updatedAt"`
	Error            string             `json:"error,omitempty"`
	RetryCount       int                `json:"retryCount"`
	Duration         int64              `json:"duration,omitempty"` // Total duration (seconds)
}

// Manager manages migration tasks
type Manager struct {
	storage      storage.Storage
	exporter     *export.Exporter
	importer     *importer.Importer
	maxRetries   int
	workerCount  int
	pollInterval time.Duration
	workerID     string
	staleTaskTTL time.Duration
}

// NewManager creates a new task manager instance
func NewManager(store storage.Storage, maxRetries int, workerCount int, pollInterval time.Duration, staleTaskTTL time.Duration) *Manager {
	exporter := export.NewExporter()
	importer := importer.NewImporter()

	if workerCount < 1 {
		workerCount = 1
	}
	if pollInterval <= 0 {
		pollInterval = 500 * time.Millisecond
	}
	if staleTaskTTL <= 0 {
		staleTaskTTL = 15 * time.Minute
	}

	manager := &Manager{
		storage:      store,
		exporter:     exporter,
		importer:     importer,
		maxRetries:   maxRetries,
		workerCount:  workerCount,
		pollInterval: pollInterval,
		workerID:     fmt.Sprintf("%s-%d", hostname(), os.Getpid()),
		staleTaskTTL: staleTaskTTL,
	}
	manager.startWorkers()
	return manager
}

// CreateTask creates a new migration task
func (m *Manager) CreateTask(id, origin, target string) (*Task, error) {
	now := time.Now().UTC().Format(time.RFC3339)

	task := &Task{
		ID:               id,
		Status:           storage.StatusPending,
		ConnectionOrigin: origin,
		ConnectionTarget: target,
		CreatedAt:        now,
		UpdatedAt:        now,
		RetryCount:       0,
	}

	if err := m.storage.CreateTask(&storage.Task{
		ID:               task.ID,
		Status:           task.Status,
		ConnectionOrigin: task.ConnectionOrigin,
		ConnectionTarget: task.ConnectionTarget,
		CreatedAt:        task.CreatedAt,
		UpdatedAt:        task.UpdatedAt,
		Error:            task.Error,
		RetryCount:       task.RetryCount,
	}); err != nil {
		if errors.Is(err, storage.ErrTaskAlreadyExists) {
			existing, getErr := m.GetTask(id)
			if getErr == nil {
				return existing, nil
			}
		}
		return nil, fmt.Errorf("failed to create task: %v", err)
	}

	return task, nil
}

func (m *Manager) startWorkers() {
	staleBefore := time.Now().UTC().Add(-m.staleTaskTTL)
	if count, err := m.storage.RequeueStaleTasks(staleBefore); err != nil {
		log.Error("Failed to requeue stale tasks on startup", zap.Error(err))
	} else if count > 0 {
		log.Info("Requeued stale tasks on startup", zap.Int64("count", count))
	}

	for idx := 0; idx < m.workerCount; idx++ {
		go m.workerLoop(idx)
	}
}

func (m *Manager) workerLoop(workerIdx int) {
	logger := log.With(
		zap.String("worker", fmt.Sprintf("%s-%d", m.workerID, workerIdx)),
	)

	for {
		storageTask, err := m.storage.ClaimNextPendingTask(fmt.Sprintf("%s-%d", m.workerID, workerIdx))
		if err != nil {
			if err != storage.ErrNoPendingTask {
				logger.Error("Failed to claim next pending task", zap.Error(err))
			}
			time.Sleep(m.pollInterval)
			continue
		}
		m.safeExecuteTask(logger, fromStorageTask(storageTask))
	}
}

func (m *Manager) safeExecuteTask(logger *zap.Logger, task *Task) {
	defer func() {
		if r := recover(); r != nil {
			task.Duration = 0
			m.updateTaskStatus(logger, task, storage.StatusFailed, fmt.Sprintf("panic recovered: %v", r))
		}
	}()
	m.executeTask(task)
}

// GetTask retrieves a task by ID
func (m *Manager) GetTask(id string) (*Task, error) {
	storageTask, err := m.storage.GetTask(id)
	if err != nil {
		return nil, err
	}

	// Convert to task instance
	task := &Task{
		ID:               storageTask.ID,
		Status:           storageTask.Status,
		ConnectionOrigin: storageTask.ConnectionOrigin,
		ConnectionTarget: storageTask.ConnectionTarget,
		CreatedAt:        storageTask.CreatedAt,
		UpdatedAt:        storageTask.UpdatedAt,
		Error:            storageTask.Error,
		RetryCount:       storageTask.RetryCount,
		Duration:         storageTask.Duration,
	}
	return task, nil
}

// executeTask executes a migration task
func (m *Manager) executeTask(task *Task) {
	logger := log.With(zap.String("taskID", task.ID))
	startTime := time.Now()

	for attempt := 0; attempt <= m.maxRetries; attempt++ {
		if attempt == 0 {
			m.updateTaskStatus(logger, task, storage.StatusRunning, "")
		} else {
			task.RetryCount = attempt
			m.updateTaskStatus(logger, task, storage.StatusRetrying, fmt.Sprintf("Retry %d/%d", attempt, m.maxRetries))
			interval := withJitter(time.Duration(1<<uint(attempt-1))*time.Second, 300*time.Millisecond)
			time.Sleep(interval)
			m.updateTaskStatus(logger, task, storage.StatusRunning, "")
		}

		err := m.runExportImport(logger, task)
		if err == nil {
			task.Duration = int64(time.Since(startTime).Seconds())
			m.updateTaskStatus(logger, task, storage.StatusSuccess, "")
			logger.Info("Task completed successfully", zap.Duration("time.elapsed", time.Since(startTime)))
			return
		}

		// Import errors are already retried in executeImportWithRetry.
		if !strings.Contains(err.Error(), "export failed") || attempt == m.maxRetries {
			task.Duration = int64(time.Since(startTime).Seconds())
			m.updateTaskStatus(logger, task, storage.StatusFailed, err.Error())
			return
		}
	}
}

// runExportImport executes the export and import process
func (m *Manager) runExportImport(logger *zap.Logger, task *Task) error {
	// Use dataDir as base directory for all runtime files
	tempDir := filepath.Join(config.DataDir, "temp_export")
	if err := os.MkdirAll(tempDir, 0755); err != nil {
		return fmt.Errorf("Failed to create temporary directory: %v", err)
	}

	// Execute export
	logger.Info("Starting data export", zap.String("origin", task.ConnectionOrigin))
	exportStart := time.Now()
	// Use unique filename with task ID to avoid conflicts
	exportFile := filepath.Join(tempDir, fmt.Sprintf("%s_data", task.ID))
	err := m.exporter.Export(logger, task.ConnectionOrigin, exportFile)
	exportElapsed := time.Since(exportStart)
	if err != nil {
		logger.Error("Export failed",
			zap.Duration("time.elapsed", exportElapsed), zap.Error(err))
		// Clean up failed export file
		os.RemoveAll(exportFile)
		return fmt.Errorf("export failed: %v", err)
	}
	logger.Info("Export successful",
		zap.Duration("time.elapsed", exportElapsed), zap.String("outputFile", exportFile))

	// Execute import with retry
	importErr := m.executeImportWithRetry(logger, task, exportFile)
	if importErr != nil {
		// Clean up failed export file
		os.RemoveAll(exportFile)
		return fmt.Errorf("import failed: %v", importErr)
	}

	// Clean up successful export file
	os.RemoveAll(exportFile)

	return nil
}

// executeImportWithRetry executes import with retry mechanism
func (m *Manager) executeImportWithRetry(logger *zap.Logger, task *Task, dataFile string) error {
	maxRetries := m.maxRetries
	retryCount := 0

	for retryCount <= maxRetries {
		// Execute import
		logger.Info("Starting data import",
			zap.String("target", task.ConnectionTarget), zap.Int("attempt", retryCount+1), zap.Int("maxAttempts", maxRetries+1))
		importStart := time.Now()
		err := m.importer.Import(logger, task.ConnectionTarget, dataFile, task.ID)
		importElapsed := time.Since(importStart)

		if err == nil {
			logger.Info("Import successful", zap.Duration("time.elapsed", importElapsed))
			return nil
		}

		logger.Error("Import failed",
			zap.Duration("time.elapsed", importElapsed), zap.Error(err))

		// Check if retry is needed
		if retryCount < maxRetries {
			retryCount++
			task.RetryCount++
			// Update task status to retrying
			m.updateTaskStatus(logger, task, storage.StatusRetrying, fmt.Sprintf("Retry import %d/%d: %v", retryCount, maxRetries, err))
			// Calculate retry interval (exponential backoff)
			interval := withJitter(time.Duration(1<<uint(retryCount-1))*time.Second, 300*time.Millisecond)
			logger.Info("Retrying import", zap.Duration("interval", interval))
			time.Sleep(interval)
		} else {
			// Max retries reached
			return err
		}
	}

	return fmt.Errorf("max retries reached")
}

// updateTaskStatus updates the status of a task
func (m *Manager) updateTaskStatus(logger *zap.Logger, task *Task, status storage.TaskStatus, errMsg string) {
	// Update task information
	task.Status = status
	task.Error = errMsg
	task.UpdatedAt = time.Now().UTC().Format(time.RFC3339)

	// Save to storage
	if err := m.storage.UpdateTask(&storage.Task{
		ID:               task.ID,
		Status:           task.Status,
		ConnectionOrigin: task.ConnectionOrigin,
		ConnectionTarget: task.ConnectionTarget,
		CreatedAt:        task.CreatedAt,
		UpdatedAt:        task.UpdatedAt,
		Error:            task.Error,
		RetryCount:       task.RetryCount,
		Duration:         task.Duration,
	}); err != nil {
		// Log error but don't interrupt process
		logger.Error("Failed to update task status", zap.Error(err))
	}
}

func hostname() string {
	h, err := os.Hostname()
	if err != nil || h == "" {
		return "unknown-host"
	}
	return h
}

func withJitter(base time.Duration, jitter time.Duration) time.Duration {
	if jitter <= 0 {
		return base
	}
	delta := time.Duration(rand.Int63n(int64(jitter)))
	return base + delta
}

func fromStorageTask(s *storage.Task) *Task {
	return &Task{
		ID:               s.ID,
		Status:           s.Status,
		ConnectionOrigin: s.ConnectionOrigin,
		ConnectionTarget: s.ConnectionTarget,
		CreatedAt:        s.CreatedAt,
		UpdatedAt:        s.UpdatedAt,
		Error:            s.Error,
		RetryCount:       s.RetryCount,
		Duration:         s.Duration,
	}
}