package knowledgeqa

import (
	"context"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/forwarder"
	knowledgeqastore "github.com/tidbcloud/AgentTeam/internal/store/knowledgeqa"
)

const (
	accountingRecoveryPollInterval = 5 * time.Second
	accountingRecoveryReadTimeout  = 5 * time.Second
	accountingRecoveryCrashGrace   = time.Minute
	maxAccountingRecoveryBatch     = 16
)

// postDispatchError marks failures after the durable generating boundary. The
// caller must terminalize the turn first and schedule accounting reconciliation
// because MeshX may already hold settled usage that was not returned or saved.
type postDispatchError struct{ cause error }

func (err *postDispatchError) Error() string {
	if err == nil || err.cause == nil {
		return "Knowledge Q&A post-dispatch execution failed"
	}
	return err.cause.Error()
}

func (err *postDispatchError) Unwrap() error {
	if err == nil {
		return nil
	}
	return err.cause
}

func (worker *Worker) runAccountingReconciler(ctx context.Context) {
	ticker := time.NewTicker(accountingRecoveryPollInterval)
	defer ticker.Stop()
	for {
		worker.reconcilePendingAccounting(ctx)
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (worker *Worker) reconcilePendingAccounting(ctx context.Context) {
	limit := min(worker.capacity, maxAccountingRecoveryBatch)
	items, err := worker.store.ListPendingAccountingRecoveries(ctx, worker.now(), limit)
	if err != nil {
		if ctx.Err() == nil {
			worker.logger.Warn("Knowledge Q&A accounting recovery scan failed", "error", err)
		}
		return
	}
	for _, item := range items {
		if ctx.Err() != nil {
			return
		}
		worker.reconcileAttemptAccounting(ctx, item)
	}
}

func (worker *Worker) reconcileAttemptAccounting(ctx context.Context, item knowledgeqastore.AccountingRecovery) {
	now := worker.now()
	recoveryCtx, cancel := context.WithTimeout(ctx, accountingRecoveryReadTimeout)
	result, err := worker.runtime.RecoverKnowledgeQATurnAccounting(
		recoveryCtx,
		forwarder.KnowledgeQAAccountingRecoveryRequest{
			Principal: forwarder.KnowledgeQARuntimePrincipal{
				OrganizationID: item.OrganizationID,
				WorkspaceID:    item.WorkspaceID,
				RequesterID:    item.RequesterID,
				TurnID:         item.TurnID,
			},
			AttemptID: item.AttemptID,
		},
	)
	cancel()
	recoveryErr := err
	if err == nil && result.Found && (result.Terminal || accountingCheckpointIsQuiescent(item, now)) {
		recoveryErr = worker.store.ResolveAccountingRecovery(ctx, knowledgeqastore.AccountingRecoveryResolution{
			WorkspaceID: item.WorkspaceID, TurnID: item.TurnID, AttemptID: item.AttemptID,
			Usage: result.Usage, SettlementRef: result.SettlementRef,
			SettlementRefs: result.SettlementRefs, Now: now,
		})
		if recoveryErr == nil {
			return
		}
	}
	expired := !item.DeadlineAt.After(now)
	if retryErr := worker.store.RecordAccountingRecoveryRetry(
		ctx,
		item.AttemptID,
		now.Add(accountingRecoveryBackoff(item.Attempts)),
		expired,
	); retryErr != nil && ctx.Err() == nil {
		worker.logger.Warn("Knowledge Q&A accounting recovery retry update failed", "attempt_id", item.AttemptID, "error", retryErr)
		return
	}
	if recoveryErr != nil && !expired && ctx.Err() == nil {
		worker.logger.Warn("Knowledge Q&A accounting recovery attempt failed", "attempt_id", item.AttemptID, "error", recoveryErr)
	}
}

func accountingCheckpointIsQuiescent(item knowledgeqastore.AccountingRecovery, now time.Time) bool {
	if item.StartedAt == nil {
		return !item.DeadlineAt.After(now)
	}
	return !item.StartedAt.Add(turnExecutionTimeout + accountingRecoveryCrashGrace).After(now)
}

func accountingRecoveryBackoff(attempts int64) time.Duration {
	switch {
	case attempts <= 0:
		return 5 * time.Second
	case attempts == 1:
		return 15 * time.Second
	case attempts == 2:
		return 30 * time.Second
	case attempts == 3:
		return time.Minute
	case attempts == 4:
		return 2 * time.Minute
	case attempts == 5:
		return 5 * time.Minute
	case attempts == 6:
		return 10 * time.Minute
	default:
		return 15 * time.Minute
	}
}