package knowledgeqa

import (
	"context"
	"time"

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

const (
	knowledgeQAProgressPollInterval = 10 * time.Second
	knowledgeQAProgressReadTimeout  = 2 * time.Second
)

var knowledgeQAProgressMilestones = [...]int{1, 2, 4, 8, 16, 32, 64, 128, 256, 512}

func (worker *Worker) executeKnowledgeQATurnWithProgress(
	ctx context.Context,
	session knowledgeqastore.Session,
	turn knowledgeqastore.Turn,
	request forwarder.KnowledgeQARuntimeRequest,
) (forwarder.KnowledgeQARuntimeResult, error) {
	progressCtx, cancelProgress := context.WithCancel(ctx)
	progressDone := make(chan struct{})
	logging.Go(nil, func() {
		defer close(progressDone)
		worker.monitorKnowledgeQATurnProgress(progressCtx, session, turn, request)
	})

	result, err := worker.runtime.ExecuteKnowledgeQATurn(ctx, request)
	cancelProgress()
	<-progressDone
	return result, err
}

func (worker *Worker) monitorKnowledgeQATurnProgress(
	ctx context.Context,
	session knowledgeqastore.Session,
	turn knowledgeqastore.Turn,
	request forwarder.KnowledgeQARuntimeRequest,
) {
	interval := worker.progressPollInterval
	if interval <= 0 {
		interval = knowledgeQAProgressPollInterval
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	lastMilestone := 0
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}

		readCtx, cancelRead := context.WithTimeout(ctx, knowledgeQAProgressReadTimeout)
		checkpoint, err := worker.runtime.RecoverKnowledgeQATurnAccounting(
			readCtx,
			forwarder.KnowledgeQAAccountingRecoveryRequest{
				Principal: request.Principal,
				AttemptID: request.IdempotencyKey,
			},
		)
		cancelRead()
		if ctx.Err() != nil {
			return
		}
		if err != nil {
			worker.logger.Debug(
				"Knowledge Q&A progress checkpoint read failed",
				"turn_id", turn.ID,
				"error", err,
			)
			continue
		}
		if checkpoint.Terminal {
			return
		}
		if !checkpoint.Found || checkpoint.Generation <= 0 {
			continue
		}
		milestone := knowledgeQAProgressMilestone(checkpoint.Generation)
		if milestone <= lastMilestone {
			continue
		}
		if worker.appendActivity(
			ctx,
			session,
			turn,
			protocol.RuntimeActivity{
				Kind: "tool", Status: "completed", ToolName: "model",
				ModelCalls: checkpoint.Generation,
			},
			map[string]any{
				"stage": "exploration_progress", "modelCalls": checkpoint.Generation,
			},
			worker.now(),
		) {
			lastMilestone = milestone
		}
	}
}

func knowledgeQAProgressMilestone(generation int) int {
	result := 0
	for _, milestone := range knowledgeQAProgressMilestones {
		if generation < milestone {
			break
		}
		result = milestone
	}
	return result
}