package extensionflow

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"github.com/tidbcloud/AgentTeam/internal/store/extension"
)

const (
	defaultWorkspaceCatalogReconcileBatchSize = 100
	workspaceCatalogReconcileLeaseDuration    = 2 * time.Minute
	workspaceCatalogReconcileWorkerOwner      = "workspace-catalog-reconcile-worker"
)

type workspaceCatalogReconcileBatchResult struct {
	Scanned           int
	Changed           int
	NextCursorAgentID string
	HasMore           bool
	Warnings          []error
}

// EnqueueWorkspaceCatalogReconcile records one coalesced Workspace dirty job in
// the caller's transaction. Source-state mutations use this instead of
// enumerating every Agent synchronously.
func (service *Service) EnqueueWorkspaceCatalogReconcile(ctx context.Context, tx coredb.Tx, workspaceID string, now time.Time) error {
	if err := service.validate(); err != nil {
		return err
	}
	_, err := service.store.InTx(tx).EnqueueWorkspaceCatalogReconcile(ctx, workspaceID, now)
	return err
}

// StartWorkspaceCatalogReconcileWorker drains coalesced Workspace reconcile jobs
// in bounded Agent pages. Idle ticks inspect only the small dirty-job table; they
// do not scan Agents.
func (service *Service) StartWorkspaceCatalogReconcileWorker(ctx context.Context, interval time.Duration, onError func(error)) {
	if service == nil || interval <= 0 {
		return
	}
	logging.Go(nil, func() {
		run := func() {
			service.drainPendingWorkspaceCatalogReconcileJobs(ctx, onError)
		}
		run()
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				run()
			}
		}
	})
}

func (service *Service) drainPendingWorkspaceCatalogReconcileJobs(ctx context.Context, onError func(error)) {
	_, err := service.processOneWorkspaceCatalogReconcileJob(ctx)
	if err != nil && onError != nil {
		onError(err)
	}
}

func (service *Service) processOneWorkspaceCatalogReconcileJob(ctx context.Context) (bool, error) {
	if err := service.validate(); err != nil {
		return false, err
	}
	if err := service.requireRegistry(); err != nil {
		return false, err
	}
	now := service.now().UTC()
	job, found, err := service.claimWorkspaceCatalogReconcileJob(ctx, now)
	if err != nil || !found {
		return found, err
	}
	result, reconcileErr := service.reconcileWorkspaceCatalogJobBatch(
		ctx,
		job.WorkspaceID,
		job.CursorAgentID,
		defaultWorkspaceCatalogReconcileBatchSize,
		now,
	)
	releaseNow := service.now().UTC()
	if reconcileErr != nil {
		if code, message, blocked := catalogReconcileBlocked(reconcileErr); blocked {
			blockErr := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
				return service.store.InTx(tx).BlockWorkspaceCatalogReconcileJob(
					ctx,
					job.WorkspaceID,
					job.LeaseOwner,
					job.ProcessingGeneration,
					code,
					message,
					releaseNow,
				)
			})
			return true, errors.Join(reconcileErr, blockErr)
		}
		delay := workspaceCatalogReconcileRetryDelay(job.WorkspaceID, job.Attempts)
		requeueErr := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
			return service.store.InTx(tx).RequeueWorkspaceCatalogReconcileJob(
				ctx,
				job.WorkspaceID,
				job.LeaseOwner,
				job.ProcessingGeneration,
				releaseNow.Add(delay),
				"transient",
				reconcileErr.Error(),
				releaseNow,
			)
		})
		return true, errors.Join(reconcileErr, requeueErr)
	}
	releaseErr := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
		return service.store.InTx(tx).ReleaseWorkspaceCatalogReconcileBatch(
			ctx,
			job.WorkspaceID,
			job.LeaseOwner,
			job.ProcessingGeneration,
			result.NextCursorAgentID,
			result.HasMore,
			releaseNow,
		)
	})
	if releaseErr != nil {
		return true, releaseErr
	}
	return true, errors.Join(result.Warnings...)
}

func (service *Service) claimWorkspaceCatalogReconcileJob(ctx context.Context, now time.Time) (extension.WorkspaceCatalogReconcileJob, bool, error) {
	var job extension.WorkspaceCatalogReconcileJob
	var found bool
	err := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
		claimed, ok, err := service.store.InTx(tx).ClaimNextWorkspaceCatalogReconcileJob(
			ctx,
			workspaceCatalogReconcileWorkerOwner,
			workspaceCatalogReconcileLeaseDuration,
			now,
		)
		job = claimed
		found = ok
		return err
	})
	return job, found, err
}

func (service *Service) reconcileWorkspaceCatalogJobBatch(
	ctx context.Context,
	workspaceID string,
	afterAgentID string,
	batchSize int,
	now time.Time,
) (workspaceCatalogReconcileBatchResult, error) {
	if err := service.validate(); err != nil {
		return workspaceCatalogReconcileBatchResult{}, err
	}
	if err := service.requireRegistry(); err != nil {
		return workspaceCatalogReconcileBatchResult{}, err
	}
	if batchSize <= 0 {
		return workspaceCatalogReconcileBatchResult{}, fmt.Errorf("batch size must be positive")
	}
	var result workspaceCatalogReconcileBatchResult
	var warnings []error
	err := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
		var txErr error
		result, warnings, txErr = service.reconcileWorkspaceCatalogJobBatchInTx(
			ctx,
			tx,
			workspaceID,
			afterAgentID,
			batchSize,
			now,
		)
		return txErr
	})
	result.Warnings = warnings
	return result, err
}

func (service *Service) reconcileWorkspaceCatalogJobBatchInTx(
	ctx context.Context,
	tx coredb.Tx,
	workspaceID string,
	afterAgentID string,
	batchSize int,
	now time.Time,
) (workspaceCatalogReconcileBatchResult, []error, error) {
	extTx := service.store.InTx(tx)
	agents, err := extTx.ListWorkspaceCatalogAgentBatch(
		ctx,
		workspaceID,
		afterAgentID,
		batchSize,
	)
	if err != nil {
		return workspaceCatalogReconcileBatchResult{}, nil, err
	}
	result := workspaceCatalogReconcileBatchResult{
		Scanned: len(agents),
	}
	if len(agents) == 0 {
		return result, nil, nil
	}
	result.NextCursorAgentID = agents[len(agents)-1].AgentID
	result.HasMore = len(agents) == batchSize
	agentIDs := make([]string, 0, len(agents))
	for _, agent := range agents {
		agentIDs = append(agentIDs, agent.AgentID)
	}
	authorizationsByAgent, err :=
		extTx.ListAgentAppAuthorizationDecisionsForAgents(
			ctx,
			workspaceID,
			agentIDs,
		)
	if err != nil {
		return result, nil, err
	}
	catalogsByAgent, rowErrors, err :=
		extTx.LoadAgentExtensionCatalogsForUpdate(
			ctx,
			workspaceID,
			agentIDs,
		)
	if err != nil {
		return result, nil, err
	}
	selections, err := extTx.ListActiveAppHostedExtensionSelections(
		ctx,
		workspaceID,
	)
	if err != nil {
		return result, nil, err
	}
	rowErrorsByAgent := make(
		map[string]extension.CatalogRowError,
		len(rowErrors),
	)
	for _, rowErr := range rowErrors {
		rowErrorsByAgent[rowErr.AgentID] = rowErr
	}
	var warnings []error
	for _, agent := range agents {
		agentID := agent.AgentID
		if rowErr, ok := rowErrorsByAgent[agentID]; ok {
			warnings = append(warnings, newCatalogReconcileBlockedError(
				CatalogReconcileErrorCorruptCatalog,
				"%s",
				rowErr.Error(),
			))
			continue
		}
		decisionsByInstallation, err := validateAppAuthorizationDecisions(
			workspaceID,
			agentID,
			authorizationsByAgent[agentID],
		)
		if err != nil {
			warnings = append(warnings, newCatalogReconcileBlockedError(
				CatalogReconcileErrorCorruptCatalog,
				"%s",
				err.Error(),
			))
			continue
		}
		current, found := catalogsByAgent[agentID]
		changed, err := service.reconcileAgentCatalogInTx(
			ctx,
			tx,
			extTx,
			workspaceID,
			agentID,
			selections,
			decisionsByInstallation,
			current,
			found,
			now,
		)
		if err != nil {
			return result, warnings, err
		}
		if !changed {
			continue
		}
		result.Changed++
	}
	return result, warnings, nil
}

func (service *Service) reconcileAgentCatalogInTx(
	ctx context.Context,
	tx coredb.Tx,
	extTx *extension.TxStore,
	workspaceID string,
	agentID string,
	selections []extension.AppHostedExtensionSelection,
	decisionsByInstallation map[string]appAuthorizationDecision,
	current extension.AgentExtensionCatalog,
	found bool,
	now time.Time,
) (bool, error) {
	entries, enabled, err :=
		service.projectAppHostedCatalogEntriesFromSelections(
			ctx,
			tx,
			workspaceID,
			selections,
			decisionsByInstallation,
		)
	if err != nil {
		return false, err
	}
	currentVersion := int64(0)
	storedState := ""
	if found {
		currentVersion = current.Version
		storedState = current.State
	}
	state, newEnabled, changed, err := projectedCatalogStateChange(
		entries,
		enabled,
		currentVersion,
		storedState,
	)
	if err != nil || !changed {
		return changed, err
	}
	if _, err := extTx.ApplyPrecomputedExtensionCatalogChange(
		ctx,
		workspaceID,
		agentID,
		current,
		found,
		state,
		newEnabled,
		now,
	); err != nil {
		return false, err
	}
	return true, nil
}

func workspaceCatalogReconcileRetryDelay(workspaceID string, attempts int) time.Duration {
	if attempts < 1 {
		attempts = 1
	}
	shift := attempts - 1
	if shift > 5 {
		shift = 5
	}
	delay := 5 * time.Second * time.Duration(1<<shift)
	if delay > 5*time.Minute {
		delay = 5 * time.Minute
	}
	jitterMax := delay / 4
	if jitterMax <= 0 {
		return delay
	}
	h := fnv.New64a()
	_, _ = h.Write([]byte(workspaceID))
	_, _ = h.Write([]byte(fmt.Sprintf("/%d", attempts)))
	jitter := time.Duration(h.Sum64() % uint64(jitterMax))
	return delay + jitter
}