package server

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/database"
)

type workerStartStage struct {
	name     string
	start    func() error
	rollback func(context.Context) error
}

func startWorkerStages(stages []workerStartStage, beforeStart func(string) error) error {
	started := make([]workerStartStage, 0, len(stages))
	for _, stage := range stages {
		if err := beforeStart(stage.name); err != nil {
			return rollbackWorkerStages(started, fmt.Errorf("start %s: %w", stage.name, err))
		}
		if err := stage.start(); err != nil {
			return rollbackWorkerStages(started, fmt.Errorf("start %s: %w", stage.name, err))
		}
		started = append(started, stage)
	}
	return nil
}

func rollbackWorkerStages(started []workerStartStage, startErr error) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	joined := startErr
	for i := len(started) - 1; i >= 0; i-- {
		if started[i].rollback != nil {
			if err := started[i].rollback(ctx); err != nil {
				joined = errors.Join(joined, fmt.Errorf("rollback %s: %w", started[i].name, err))
			}
		}
	}
	return joined
}

func cancelAndWait(ctx context.Context, cancel context.CancelFunc, done <-chan struct{}) error {
	if cancel == nil {
		return nil
	}
	cancel()
	if done == nil {
		return nil
	}
	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Shutdown drains workers before closing handlers and database dependencies.
func (s *Server) Shutdown(ctx context.Context) error {
	if s == nil {
		return nil
	}
	s.shutdownMu.Lock()
	defer s.shutdownMu.Unlock()
	if s.shutdownComplete {
		return nil
	}
	if err := s.QuiesceBackgroundWorkers(ctx); err != nil {
		return fmt.Errorf("shutdown before dependency close: %w", err)
	}
	s.handlerMu.Lock()
	s.handler = nil
	s.handlerMu.Unlock()
	if s.resourceActivityBridge != nil {
		s.resourceActivityBridge.Close()
	}
	if s.automationBranchBridge != nil {
		s.automationBranchBridge.Close()
	}
	if s.unregisterGit != nil {
		s.unregisterGit()
		s.unregisterGit = nil
	}
	var errs []string
	if s.ags != nil {
		if err := s.ags.Shutdown(ctx); err != nil {
			errs = append(errs, err.Error())
		} else {
			s.ags = nil
		}
	}
	if s.contentPayload != nil {
		if err := s.contentPayload.Close(); err != nil {
			errs = append(errs, err.Error())
		} else {
			s.contentPayload = nil
		}
	}
	if s.meshxDB != nil {
		if err := database.Close(s.meshxDB); err != nil {
			errs = append(errs, err.Error())
		} else {
			s.meshxDB = nil
		}
	}
	if s.agentteamCoreDB != nil {
		if err := database.Close(s.agentteamCoreDB); err != nil {
			errs = append(errs, err.Error())
		} else {
			s.agentteamCoreDB = nil
		}
	}
	if len(errs) > 0 {
		return fmt.Errorf("shutdown: %s", strings.Join(errs, "; "))
	}
	s.shutdownComplete = true
	return nil
}

type WorkerOwnership struct {
	WorkspaceScope    WorkspaceScope
	GlobalMaintenance bool
}

// StartBackgroundWorkers starts the configured worker set exactly once.
func (s *Server) StartBackgroundWorkers(ownership ...WorkerOwnership) error {
	if s == nil {
		return fmt.Errorf("start background workers: server is nil")
	}
	s.workersMu.Lock()
	defer s.workersMu.Unlock()
	if s.workersStopped || s.retiring {
		return fmt.Errorf("start background workers: workers were already stopped")
	}
	if s.workersStarted {
		return nil
	}
	if len(ownership) > 0 {
		s.workspaceScope = ownership[0].WorkspaceScope
		s.globalMaintenance = ownership[0].GlobalMaintenance
	} else {
		s.globalMaintenance = true
	}
	if s.startBackgroundWorkers != nil {
		if err := s.startBackgroundWorkers(); err != nil {
			s.workersStopped = true
			return fmt.Errorf("start background workers: %w", err)
		}
	}
	s.workersStarted = true
	return nil
}

// QuiesceBackgroundWorkers stops new work and waits for active loops to exit.
func (s *Server) QuiesceBackgroundWorkers(ctx context.Context) error {
	if s == nil {
		return nil
	}
	s.workersMu.Lock()
	s.retiring = true
	s.ready.Store(false)
	if s.workersDrained {
		s.workersMu.Unlock()
		return nil
	}
	if !s.workersStopped {
		s.workersStopped = true
	}
	s.workersMu.Unlock()
	err := s.stopWorkers(ctx)
	if err == nil {
		s.workersMu.Lock()
		s.workersDrained = true
		s.workersMu.Unlock()
	}
	return err
}

func (s *Server) stopWorkers(ctx context.Context) error {
	var errs []string
	for name, bridge := range map[string]interface{ Quiesce(context.Context) error }{
		"resource activity bridge": s.resourceActivityBridge,
		"automation bridge":        s.automationBranchBridge,
	} {
		if bridge != nil {
			if err := bridge.Quiesce(ctx); err != nil {
				errs = append(errs, name+": "+err.Error())
			}
		}
	}
	stopAndWait := func(cancel context.CancelFunc, done <-chan struct{}) {
		if cancel == nil {
			return
		}
		cancel()
		if done != nil {
			select {
			case <-done:
			case <-ctx.Done():
				errs = append(errs, ctx.Err().Error())
			}
		}
	}
	if s.serverRunner != nil {
		if err := s.serverRunner.Quiesce(ctx); err != nil {
			errs = append(errs, err.Error())
		}
	}
	stopAndWait(s.runtimeCancel, s.runtimeDone)
	if s.serverRunner != nil {
		if err := s.serverRunner.Close(); err != nil {
			errs = append(errs, err.Error())
		}
	}
	stopAndWait(s.kbWorkerCancel, s.kbWorkerDone)
	stopAndWait(s.contentOutboxCancel, s.contentOutboxDone)
	stopAndWait(s.knowledgeQACleanupCancel, s.knowledgeQACleanupDone)
	if len(errs) > 0 {
		return fmt.Errorf("quiesce background workers: %s", strings.Join(errs, "; "))
	}
	return nil
}