package agent

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/promptcompiler"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

type transientRuntimeRetryState struct {
	MessageID   string
	SessionID   string
	Strategy    string
	RecoveryRef string
	Input       drivercontract.RuntimeInput
	Attempts    int
}

func (ap *Process) shouldRetryTransientRuntimeError(ev drivercontract.ParsedEvent) (transientRuntimeErrorClassification, int, bool) {
	classification := classifyTransientRuntimeError(ev)
	if !classification.Retryable {
		return transientRuntimeErrorClassification{}, 0, false
	}

	ap.mu.Lock()
	defer ap.mu.Unlock()
	if ap.stopping || ap.disableRecovery || ap.transientRuntimeRetryToolCallObserved || ap.pendingWake == nil {
		return classification, 0, false
	}
	messageID := strings.TrimSpace(ap.pendingWake.Wake.MessageID)
	if messageID == "" {
		return classification, 0, false
	}
	if ap.transientRuntimeRetry.MessageID != messageID {
		ap.transientRuntimeRetry = transientRuntimeRetryState{MessageID: messageID}
	}
	if ap.transientRuntimeRetryPending || ap.transientRuntimeRetry.Attempts >= transientRuntimeRetryMaxAttempts {
		return classification, ap.transientRuntimeRetry.Attempts, false
	}

	strategy := strings.TrimSpace(ap.transientRuntimeRetry.Strategy)
	recoveryRef := strings.TrimSpace(ap.transientRuntimeRetry.RecoveryRef)
	if strategy == "" || recoveryRef == "" {
		var supported bool
		var err error
		strategy, recoveryRef, supported, err = prepareDriverRecoveryPoint(ap.driver, messageID)
		if err != nil {
			slog.Warn("prepare transient runtime retry recovery point failed",
				"agent_id", ap.AgentID,
				"message_id", messageID,
				"error", err,
			)
			return classification, ap.transientRuntimeRetry.Attempts, false
		}
		if !supported || strings.TrimSpace(strategy) == "" || strings.TrimSpace(recoveryRef) == "" {
			return classification, ap.transientRuntimeRetry.Attempts, false
		}
		ap.transientRuntimeRetry.Strategy = strings.TrimSpace(strategy)
		ap.transientRuntimeRetry.RecoveryRef = strings.TrimSpace(recoveryRef)
		ap.transientRuntimeRetry.SessionID = strings.TrimSpace(firstNonEmptyString(ev.SessionID, ap.SessionID))
	}

	ap.transientRuntimeRetry.Attempts++
	ap.transientRuntimeRetryPending = true
	ap.transientRuntimeRetryAwaitingTurnEnd = true
	return classification, ap.transientRuntimeRetry.Attempts, true
}

func (ap *Process) retryTransientRuntimeError(ev drivercontract.ParsedEvent, classification transientRuntimeErrorClassification, attempt int) {
	delay := transientRuntimeRetryDelayFunc(attempt)
	detail := transientRuntimeRetryDetail(classification, attempt, delay)

	ap.RecordCheckpoint("transient_runtime_retry", detail)
	ap.emitRuntimeEvent(RuntimeEvent{
		MessageID:  ap.pendingWakeMessageID(),
		WakeToken:  ap.pendingWakeTokenSnapshot(),
		EventClass: consts.RuntimeEventClassRuntimeRetrying,
		Target:     ap.pendingWakeTarget(),
		Detail:     detail,
	})

	go ap.runTransientRuntimeRetry(ev, classification, attempt, delay)
}

func (ap *Process) runTransientRuntimeRetry(ev drivercontract.ParsedEvent, classification transientRuntimeErrorClassification, attempt int, delay time.Duration) {
	if !ap.waitForTransientRuntimeRetryTurnEnd() {
		return
	}

	if delay > 0 {
		timer := time.NewTimer(delay)
		select {
		case <-timer.C:
		case <-ap.processDone():
			timer.Stop()
			return
		}
	}

	state, ok := ap.transientRuntimeRetrySnapshot()
	if !ok {
		return
	}

	ctx, cancel := ap.newProtocolDeliveryContextFor(context.Background())
	defer cancel()
	if err := ap.restoreOrRestartForTransientRuntimeRetry(ctx, state); err != nil {
		ap.mu.Lock()
		if ap.transientRuntimeRetry.MessageID == state.MessageID {
			ap.transientRuntimeRetryPending = false
		}
		ap.mu.Unlock()
		slog.Warn("transient runtime retry restore failed",
			"agent_id", ap.AgentID,
			"message_id", state.MessageID,
			"attempt", attempt,
			"reason", classification.Reason,
			"error", err,
		)
		detail, _ := ap.formatRuntimeErrorDetail(ev)
		ap.failRuntimeError(drivercontract.ParsedEvent{
			Kind:    consts.ParsedEventKindError,
			Message: fmt.Sprintf("%s; transient retry restore failed after attempt %d: %v", detail, attempt, err),
		})
		return
	}
	ap.mu.Lock()
	if ap.transientRuntimeRetry.MessageID == state.MessageID {
		ap.transientRuntimeRetryPending = false
		ap.transientRuntimeRetryAwaitingTurnEnd = false
		ap.isIdle = false
	}
	ap.mu.Unlock()
}

func (ap *Process) waitForTransientRuntimeRetryTurnEnd() bool {
	ap.mu.Lock()
	if !ap.transientRuntimeRetryAwaitingTurnEnd {
		ap.mu.Unlock()
		return true
	}
	ap.mu.Unlock()

	ticker := time.NewTicker(50 * time.Millisecond)
	defer ticker.Stop()
	timer := time.NewTimer(transientRuntimeRetryTurnEndWait)
	defer timer.Stop()
	for {
		select {
		case <-ap.processDone():
			return false
		case <-timer.C:
			ap.mu.Lock()
			ap.transientRuntimeRetryAwaitingTurnEnd = false
			ap.mu.Unlock()
			return true
		case <-ticker.C:
			ap.mu.Lock()
			awaiting := ap.transientRuntimeRetryAwaitingTurnEnd
			ap.mu.Unlock()
			if !awaiting {
				return true
			}
		}
	}
}

func (ap *Process) restoreOrRestartForTransientRuntimeRetry(ctx context.Context, state transientRuntimeRetryState) error {
	if strings.TrimSpace(state.Strategy) == "" || strings.TrimSpace(state.RecoveryRef) == "" {
		return fmt.Errorf("transient retry recovery point unavailable")
	}
	if err := ap.acquireDeliveryGate(ctx); err != nil {
		return err
	}
	defer ap.releaseDeliveryGate()
	var err error
	if ap.SupportsLiveRecovery(state.Strategy) {
		err = ap.restoreLiveForTransientRuntimeRetry(ctx, state)
	} else {
		err = ap.redeliverForTransientRuntimeRetry(ctx, state)
	}
	if err != nil {
		return err
	}
	ap.markRuntimeHandoffSucceeded()
	return nil
}

func (ap *Process) redeliverForTransientRuntimeRetry(ctx context.Context, state transientRuntimeRetryState) error {
	ap.mu.Lock()
	if ap.stopping {
		ap.mu.Unlock()
		return fmt.Errorf("runtime is stopping")
	}
	stdinPipe := ap.stdinPipe
	sessionID := strings.TrimSpace(firstNonEmptyString(state.SessionID, ap.SessionID))
	ap.mu.Unlock()

	if strings.TrimSpace(state.MessageID) == "" {
		return fmt.Errorf("transient retry message id unavailable")
	}
	if sessionID == "" {
		return fmt.Errorf("session id unavailable")
	}
	input := state.Input
	if input.Empty() {
		return fmt.Errorf("transient retry input unavailable")
	}
	if strings.TrimSpace(input.ID) == "" {
		input.ID = strings.TrimSpace(state.MessageID)
	}
	var encoded []byte
	structured, hasStructured := ap.driver.(drivercontract.RuntimeInputDeliveryDriver)
	acked, hasAcked := ap.driver.(drivercontract.AckedBusyRuntimeInputDriver)
	if !hasStructured && !(hasAcked && acked.SupportsAckedBusyRuntimeInputDelivery()) {
		if stdinPipe == nil {
			return fmt.Errorf("stdin delivery unavailable")
		}
		encoded = ap.driver.EncodeStdinMessage(input.PromptText(), sessionID)
		if len(encoded) == 0 || strings.TrimSpace(string(encoded)) == "{}" {
			return fmt.Errorf("runtime cannot encode retry input")
		}
	}
	if err := ap.markResumeSafetyUnsafe(); err != nil {
		return err
	}

	ap.mu.Lock()
	if ap.stopping {
		ap.mu.Unlock()
		return fmt.Errorf("runtime is stopping")
	}
	ap.transientRuntimeRetryPending = false
	ap.runtimeErrorDetail = ""
	ap.runtimeErrorEventEmitted = false
	ap.runtimeErrorWakeFailed = false
	ap.isIdle = false
	ap.transientRuntimeRetryAwaitingTurnEnd = false
	ap.mu.Unlock()

	ap.rememberTransientRuntimeRetryInput(state.MessageID, input, sessionID)
	if hasStructured {
		if err := structured.DeliverRuntimeInput(ctx, input, sessionID); err != nil {
			return err
		}
	} else if hasAcked && acked.SupportsAckedBusyRuntimeInputDelivery() {
		if err := acked.DeliverAckedBusyRuntimeInput(ctx, input, sessionID); err != nil {
			return err
		}
	} else {
		if _, err := stdinPipe.Write(encoded); err != nil {
			return fmt.Errorf("stdin retry write: %w", err)
		}
	}
	ap.emitParsedEvent(newPromptInputParsedEvent(
		input.PromptText(),
		consts.PromptKindRuntimeRetry,
		ap.driver.Name(),
		consts.PromptKindRuntimeRetry,
		sessionID,
	))
	ap.markTransientRuntimeRetryRedelivered(state.MessageID)
	return nil
}

func (ap *Process) restoreLiveForTransientRuntimeRetry(ctx context.Context, state transientRuntimeRetryState) error {
	ap.mu.Lock()
	recoveryDriver, ok := ap.driver.(drivercontract.RecoveryDriver)
	if !ok {
		ap.mu.Unlock()
		return fmt.Errorf("runtime %q does not support live recovery", ap.driver.Name())
	}
	if ap.stopping {
		ap.mu.Unlock()
		return fmt.Errorf("runtime is stopping")
	}
	stdinPipe := ap.stdinPipe
	sessionID := strings.TrimSpace(firstNonEmptyString(state.SessionID, ap.SessionID))
	ap.mu.Unlock()

	if _, err := ValidateRecoveryCapability(strings.TrimSpace(ap.driver.Name()), state.Strategy); err != nil {
		return err
	}
	if sessionID == "" {
		return fmt.Errorf("session id unavailable")
	}
	input := state.Input
	if input.Empty() {
		return fmt.Errorf("transient retry input unavailable")
	}
	if strings.TrimSpace(input.ID) == "" {
		input.ID = strings.TrimSpace(state.MessageID)
	}
	structured, hasStructured := ap.driver.(drivercontract.RuntimeInputDeliveryDriver)
	var encoded []byte
	if !hasStructured {
		if stdinPipe == nil {
			return fmt.Errorf("stdin delivery unavailable")
		}
		encoded = ap.driver.EncodeStdinMessage(input.PromptText(), sessionID)
		if len(encoded) == 0 || strings.TrimSpace(string(encoded)) == "{}" {
			return fmt.Errorf("runtime cannot encode retry input")
		}
	}
	if err := ap.markResumeSafetyUnsafe(); err != nil {
		return err
	}
	ap.mu.Lock()
	if ap.stopping {
		ap.mu.Unlock()
		return fmt.Errorf("runtime is stopping")
	}
	ap.isIdle = false
	ap.mu.Unlock()

	if err := recoveryDriver.RestoreRecoveryPoint(ctx, state.Strategy, state.RecoveryRef); err != nil {
		ap.mu.Lock()
		ap.isIdle = true
		ap.mu.Unlock()
		return err
	}

	ap.mu.Lock()
	ap.runtimeErrorDetail = ""
	ap.runtimeErrorEventEmitted = false
	ap.runtimeErrorWakeFailed = false
	ap.transientRuntimeRetryPending = false
	ap.transientRuntimeRetryToolCallObserved = false
	ap.transientRuntimeRetryAwaitingTurnEnd = false
	messageID := strings.TrimSpace(state.MessageID)
	if messageID != "" {
		ap.ackedMessages = removeStringValue(ap.ackedMessages, messageID)
		delete(ap.ackedSet, messageID)
	}
	ap.isIdle = false
	ap.mu.Unlock()

	ap.rememberTransientRuntimeRetryInput(state.MessageID, input, sessionID)
	if hasStructured {
		if err := structured.DeliverRuntimeInput(ctx, input, sessionID); err != nil {
			return err
		}
	} else {
		if _, err := stdinPipe.Write(encoded); err != nil {
			return fmt.Errorf("stdin retry write: %w", err)
		}
	}
	ap.emitParsedEvent(newPromptInputParsedEvent(
		input.PromptText(),
		consts.PromptKindRuntimeRetry,
		ap.driver.Name(),
		consts.PromptKindRuntimeRetry,
		sessionID,
	))
	ap.markTransientRuntimeRetryRedelivered(state.MessageID)
	return nil
}

func (ap *Process) markTransientRuntimeRetryRedelivered(messageID string) {
	messageID = strings.TrimSpace(messageID)
	if messageID == "" {
		return
	}
	ap.mu.Lock()
	if strings.TrimSpace(ap.transientRuntimeRetry.MessageID) == messageID {
		ap.transientRuntimeRetryRedelivered = true
	}
	ap.mu.Unlock()
}

func (ap *Process) clearTransientRuntimeRetryAwaitingTurnEnd() {
	ap.mu.Lock()
	ap.transientRuntimeRetryAwaitingTurnEnd = false
	ap.mu.Unlock()
}

func (ap *Process) transientRuntimeRetrySnapshot() (transientRuntimeRetryState, bool) {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	state := ap.transientRuntimeRetry
	return state, strings.TrimSpace(state.MessageID) != ""
}

func transientRuntimeRetryDetail(classification transientRuntimeErrorClassification, attempt int, delay time.Duration) string {
	reason := strings.TrimSpace(classification.Reason)
	if reason == "" {
		reason = "transient_provider_error"
	}
	return promptcompiler.RuntimeRetryDetail(classification.HTTPStatus, reason, attempt, transientRuntimeRetryMaxAttempts, delay.Round(time.Millisecond).String())
}

func (ap *Process) pendingWakeMessageID() string {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	if ap.pendingWake == nil {
		return ""
	}
	return strings.TrimSpace(ap.pendingWake.Wake.MessageID)
}

func (ap *Process) pendingWakeTarget() string {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	if ap.pendingWake == nil {
		return ""
	}
	return strings.TrimSpace(ap.pendingWake.Wake.Target)
}

func (ap *Process) pendingWakeTokenSnapshot() string {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	return strings.TrimSpace(ap.pendingWakeToken)
}

func (ap *Process) processDone() <-chan struct{} {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	if ap.exited != nil {
		return ap.exited
	}
	done := make(chan struct{})
	return done
}

func removeStringValue(values []string, value string) []string {
	value = strings.TrimSpace(value)
	if value == "" || len(values) == 0 {
		return values
	}
	next := values[:0]
	for _, candidate := range values {
		if strings.TrimSpace(candidate) == value {
			continue
		}
		next = append(next, candidate)
	}
	return next
}

func (ap *Process) clearTransientRuntimeRetryStateForWakeLocked(messageID string) {
	if strings.TrimSpace(ap.transientRuntimeRetry.MessageID) == strings.TrimSpace(messageID) {
		ap.transientRuntimeRetry = transientRuntimeRetryState{}
		ap.transientRuntimeRetryPending = false
		ap.transientRuntimeRetryAwaitingTurnEnd = false
		ap.transientRuntimeRetryRedelivered = false
	}
}

func (ap *Process) recordSpawnTransientRuntimeRetryPoint(point *drivercontract.PreparedRecoveryPoint) {
	if point == nil {
		return
	}
	if strings.TrimSpace(point.MessageID) == "" || strings.TrimSpace(point.Strategy) == "" || strings.TrimSpace(point.RecoveryRef) == "" {
		return
	}
	ap.mu.Lock()
	messageID := strings.TrimSpace(point.MessageID)
	if ap.transientRuntimeRetry.MessageID != messageID {
		ap.transientRuntimeRetry = transientRuntimeRetryState{MessageID: messageID}
	}
	ap.transientRuntimeRetry.SessionID = strings.TrimSpace(point.ThreadID)
	ap.transientRuntimeRetry.Strategy = strings.TrimSpace(point.Strategy)
	ap.transientRuntimeRetry.RecoveryRef = strings.TrimSpace(point.RecoveryRef)
	ap.mu.Unlock()
}

func (ap *Process) rememberTransientRuntimeRetryInput(messageID string, input drivercontract.RuntimeInput, sessionID string) {
	messageID = strings.TrimSpace(messageID)
	if messageID == "" {
		messageID = strings.TrimSpace(input.ID)
	}
	if messageID == "" || input.Empty() {
		return
	}
	if strings.TrimSpace(input.ID) == "" {
		input.ID = messageID
	}
	ap.mu.Lock()
	if ap.transientRuntimeRetry.MessageID != messageID {
		ap.transientRuntimeRetry = transientRuntimeRetryState{MessageID: messageID}
	}
	ap.transientRuntimeRetry.Input = input
	if trimmedSessionID := strings.TrimSpace(sessionID); trimmedSessionID != "" {
		ap.transientRuntimeRetry.SessionID = trimmedSessionID
	}
	ap.mu.Unlock()
}

func (ap *Process) recordTransientRuntimeRetryPoint(messageID string, strategy string, recoveryRef string) {
	messageID = strings.TrimSpace(messageID)
	strategy = strings.TrimSpace(strategy)
	recoveryRef = strings.TrimSpace(recoveryRef)
	if messageID == "" || strategy == "" || recoveryRef == "" {
		return
	}
	ap.mu.Lock()
	if ap.transientRuntimeRetry.MessageID != messageID {
		ap.transientRuntimeRetry = transientRuntimeRetryState{MessageID: messageID}
	}
	ap.transientRuntimeRetry.Strategy = strategy
	ap.transientRuntimeRetry.RecoveryRef = recoveryRef
	if strings.TrimSpace(ap.transientRuntimeRetry.SessionID) == "" {
		ap.transientRuntimeRetry.SessionID = strings.TrimSpace(ap.SessionID)
	}
	ap.mu.Unlock()
}

func (ap *Process) transientRuntimeRetryRedeliveryMatches(messageID string) bool {
	messageID = strings.TrimSpace(messageID)
	if messageID == "" {
		return false
	}
	ap.mu.Lock()
	defer ap.mu.Unlock()
	return ap.transientRuntimeRetryRedelivered &&
		strings.TrimSpace(ap.transientRuntimeRetry.MessageID) == messageID
}

func initialRuntimeInputForTransientRetry(cfg drivercontract.LaunchConfig) drivercontract.RuntimeInput {
	return drivercontract.RuntimeInput{
		ID:             strings.TrimSpace(cfg.InitialDeliveryMessageID),
		Resume:         strings.TrimSpace(cfg.SessionID) != "",
		CompiledPrompt: cfg.CompiledPrompt,
		Attachments:    append([]types.RuntimeAttachment(nil), cfg.PromptAttachments...),
	}
}

func launchSessionIDForTransientRetry(cfg drivercontract.LaunchConfig, point *drivercontract.PreparedRecoveryPoint) string {
	if point != nil && strings.TrimSpace(point.ThreadID) != "" {
		return strings.TrimSpace(point.ThreadID)
	}
	return strings.TrimSpace(cfg.SessionID)
}