package server

import (
	"context"
	"log/slog"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/config"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"github.com/tidbcloud/AgentTeam/internal/service/kb"
)

type kbPeriodicWorker struct {
	interval       time.Duration
	runImmediately bool
	component      string
	run            func(context.Context) error
	errMsg         string
}

func startKBPeriodicWorkers(ctx context.Context, logger *slog.Logger, workers ...kbPeriodicWorker) <-chan struct{} {
	if logger == nil {
		logger = slog.Default().With(slog.String("domain", "kb"))
	}
	done := make(chan struct{})
	var wg sync.WaitGroup
	for _, worker := range workers {
		if worker.run == nil {
			continue
		}
		worker := worker
		workerLogger := logger
		if component := strings.TrimSpace(worker.component); component != "" {
			workerLogger = logger.With(slog.String("component", component))
		}
		wg.Add(1)
		logging.Go(workerLogger, func() {
			defer wg.Done()
			if worker.runImmediately {
				if err := worker.run(ctx); err != nil {
					workerLogger.ErrorContext(ctx, worker.errMsg, slog.String("error_class", "worker_failed"))
				}
			}
			ticker := time.NewTicker(worker.interval)
			defer ticker.Stop()
			for {
				select {
				case <-ctx.Done():
					return
				case <-ticker.C:
					if err := worker.run(ctx); err != nil {
						workerLogger.ErrorContext(ctx, worker.errMsg, slog.String("error_class", "worker_failed"))
					}
				}
			}
		})
	}
	logging.Go(logger, func() {
		defer close(done)
		wg.Wait()
	})
	return done
}

func startKBCodeWikiWorker(cfg config.KBConfig, kbService *kb.KBService, sourceMirrorSyncer kb.AGSHTTPSourceMirrorSyncer, publicRepoCache publicRepoCacheWorker, workspaceScope WorkspaceScope, logger *slog.Logger, globalMaintenance ...bool) (context.CancelFunc, <-chan struct{}) {
	if kbService == nil {
		return nil, nil
	}
	ctx, cancel := context.WithCancel(context.Background())
	if logger == nil {
		logger = slog.Default().With(slog.String("domain", "kb"))
	}
	sourceMirrorLogger := logger.With(slog.String("component", "source_mirror"))
	wikiLogger := logger.With(slog.String("component", "wiki_generation"))
	startGlobalMaintenance := len(globalMaintenance) == 0 || globalMaintenance[0]
	mirrorWorkerInterval := cfg.SourceMirrorWorkerInterval
	if mirrorWorkerInterval <= 0 {
		mirrorWorkerInterval = 30 * time.Second
	}
	mirrorWorkerBatch := cfg.SourceMirrorWorkerBatch
	if mirrorWorkerBatch <= 0 {
		mirrorWorkerBatch = 3
	}
	mirrorSchedulerInterval := cfg.SourceMirrorSchedulerInterval
	if mirrorSchedulerInterval <= 0 {
		mirrorSchedulerInterval = 60 * time.Second
	}
	mirrorSchedulerBatch := cfg.SourceMirrorSchedulerBatch
	if mirrorSchedulerBatch <= 0 {
		mirrorSchedulerBatch = 25
	}
	reclaimStaleKBJobs(ctx, kbService, workspaceScope, sourceMirrorLogger, wikiLogger)
	workers := make([]kbPeriodicWorker, 0, 4)
	if startGlobalMaintenance {
		workers = append(workers, publicRepoCachePeriodicWorkers(cfg, publicRepoCache)...)
		workers = append(workers, kbPeriodicWorker{
			interval:       kb.SourceMirrorWorkspaceCleanupInterval,
			runImmediately: true,
			component:      "source_mirror_cleanup",
			errMsg:         "kb source mirror workspace cleanup failed",
			run: func(ctx context.Context) error {
				result, err := sourceMirrorSyncer.CleanupStaleWorkspaces(ctx, time.Now())
				if result.Removed > 0 {
					sourceMirrorLogger.InfoContext(ctx, "kb source mirror stale workspaces removed", "count", result.Removed)
				}
				return err
			},
		})
	}
	workers = append(workers,
		kbPeriodicWorker{
			interval:  mirrorWorkerInterval,
			component: "source_mirror_worker",
			errMsg:    "kb source mirror worker tick failed",
			run: func(ctx context.Context) error {
				workspaceIDs, err := loadWorkspaceScope(ctx, workspaceScope)
				if err != nil {
					return err
				}
				_, err = kbService.RunPendingSourceMirrorJobs(ctx, kb.SourceMirrorRunOptions{Limit: mirrorWorkerBatch, LeaseOwner: "meshx-kb-source-mirror-worker", LeaseDuration: kb.SourceMirrorLeaseDuration(cfg.SourceMirrorTimeout), WorkspaceIDs: workspaceIDs})
				return err
			},
		},
		kbPeriodicWorker{
			interval:  mirrorSchedulerInterval,
			component: "source_mirror_scheduler",
			errMsg:    "kb source mirror scheduler tick failed",
			run: func(ctx context.Context) error {
				workspaceIDs, err := loadWorkspaceScope(ctx, workspaceScope)
				if err != nil {
					return err
				}
				_, err = kbService.EnqueueDueSourceMirrorSyncs(ctx, kb.SourceMirrorScheduleOptions{Limit: mirrorSchedulerBatch, WorkspaceIDs: workspaceIDs})
				return err
			},
		},
	)
	workers = append(workers, kbPeriodicWorker{
		interval:       mirrorWorkerInterval,
		runImmediately: true,
		component:      "wiki_terminal_projection",
		errMsg:         "kb Wiki generation terminal notification worker tick failed",
		run: func(ctx context.Context) error {
			workspaceIDs, err := loadWorkspaceScope(ctx, workspaceScope)
			if err != nil {
				return err
			}
			_, err = kbService.ProjectPendingWikiGenerationTerminalNotifications(ctx, mirrorWorkerBatch, workspaceIDs)
			return err
		},
	})
	if kbService.WikiGenerationEnabled() {
		workers = append(workers, kbPeriodicWorker{
			interval:  mirrorWorkerInterval,
			component: "wiki_generation_worker",
			errMsg:    "kb Wiki generation worker tick failed",
			run: func(ctx context.Context) error {
				workspaceIDs, err := loadWorkspaceScope(ctx, workspaceScope)
				if err != nil {
					return err
				}
				_, err = kbService.RunPendingWikiGenerationJobs(ctx, kb.WikiGenerationRunOptions{
					Limit: mirrorWorkerBatch, LeaseOwner: "meshx-kb-wiki-generation-worker",
					LeaseDuration: cfg.WikiGenerationTimeout + 5*time.Minute, WorkspaceIDs: workspaceIDs,
				})
				return err
			},
		})
	}
	done := startKBPeriodicWorkers(ctx, logger, workers...)
	return cancel, done
}

// reclaimStaleKBJobs makes jobs left running by a previous process eligible
// before the periodic workers begin polling.
func reclaimStaleKBJobs(ctx context.Context, kbService *kb.KBService, workspaceScope WorkspaceScope, sourceMirrorLogger, wikiLogger *slog.Logger) {
	workspaceIDs, scopeErr := loadWorkspaceScope(ctx, workspaceScope)
	if scopeErr != nil {
		sourceMirrorLogger.ErrorContext(ctx, "kb source mirror startup workspace scope failed", slog.String("error_class", "workspace_scope_failed"))
		return
	}
	if reclaimed, err := kbService.ReclaimStaleSourceMirrorJobs(ctx, workspaceIDs); err != nil {
		sourceMirrorLogger.ErrorContext(ctx, "kb source mirror startup reclaim failed", slog.String("error_class", "startup_reclaim_failed"))
	} else if reclaimed > 0 {
		sourceMirrorLogger.InfoContext(ctx, "kb source mirror startup reclaimed stale jobs", "count", reclaimed)
	}
	if !kbService.WikiGenerationEnabled() {
		return
	}
	if reclaimed, err := kbService.ReclaimStaleWikiGenerationJobs(ctx, workspaceIDs); err != nil {
		wikiLogger.ErrorContext(ctx, "kb Wiki generation startup reclaim failed", slog.String("error_class", "startup_reclaim_failed"))
	} else if reclaimed > 0 {
		wikiLogger.InfoContext(ctx, "kb Wiki generation startup reclaimed stale jobs", "count", reclaimed)
	}
}

func loadWorkspaceScope(ctx context.Context, scope WorkspaceScope) ([]string, error) {
	if scope == nil {
		return nil, nil
	}
	workspaceIDs, err := scope(ctx)
	if err != nil {
		return nil, err
	}
	if workspaceIDs == nil {
		return []string{}, nil
	}
	return workspaceIDs, nil
}