package providers

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strconv"
	"strings"
	"time"

	"pentagi/pkg/cast"
	"pentagi/pkg/csum"
	"pentagi/pkg/database"
	"pentagi/pkg/graphiti"
	obs "pentagi/pkg/observability"
	"pentagi/pkg/observability/langfuse"
	"pentagi/pkg/providers/pconfig"
	"pentagi/pkg/templates"
	"pentagi/pkg/tools"

	"github.com/sirupsen/logrus"
	"github.com/vxcontrol/langchaingo/llms"
	"github.com/vxcontrol/langchaingo/llms/reasoning"
	"github.com/vxcontrol/langchaingo/llms/streaming"
)

const (
	maxRetriesToCallSimpleChain = 3
	maxRetriesToCallAgentChain  = 3
	maxRetriesToCallFunction    = 3
	// repeatEchoLimit caps the previous tool answer quoted back on a repeat, so
	// echoing it cannot re-inflate a chain that is already being summarized.
	repeatEchoLimit                = 4096
	maxReflectorCallsPerChain      = 3
	maxGeneralAgentChainIterations = 100
	maxLimitedAgentChainIterations = 20
	maxAgentShutdownIterations     = 3
	maxSoftDetectionsBeforeAbort   = 4
	delayBetweenRetries            = 5 * time.Second
)

// ErrAgentModelCall marks failures while obtaining a usable model response.
// Persistence and tool execution errors deliberately do not carry this marker:
// callers may fall back from an optional model step only when the saved plan is
// still trustworthy.
var ErrAgentModelCall = errors.New("agent model call failed")

type callResult struct {
	streamID  int64
	funcCalls []llms.ToolCall
	info      map[string]any
	thinking  *reasoning.ContentReasoning
	content   string
}

func (fp *flowProvider) performAgentChain(
	ctx context.Context,
	optAgentType pconfig.ProviderOptionsType,
	chainID int64,
	taskID, subtaskID *int64,
	chain []llms.MessageContent,
	executor tools.ContextToolsExecutor,
	summarizer csum.Summarizer,
) error {
	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "providers.flowProvider.performAgentChain")
	defer span.End()

	var (
		wantToStop        bool
		monitor           = fp.buildMonitor()
		detector          = &repeatingDetector{}
		summarizerHandler = fp.GetSummarizeResultHandler(taskID, subtaskID)
	)

	logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(fp.flowID, taskID, subtaskID, logrus.Fields{
		"provider":     fp.Type(),
		"agent":        optAgentType,
		"msg_chain_id": chainID,
	}))

	lastUpdateTime := time.Now()
	rollLastUpdateTime := func() float64 {
		durationDelta := time.Since(lastUpdateTime).Seconds()
		lastUpdateTime = time.Now()
		return durationDelta
	}

	executionContext, err := fp.getExecutionContext(ctx, taskID, subtaskID)
	if err != nil {
		logger.WithError(err).Error("failed to get execution context")
		return fmt.Errorf("failed to get execution context: %w", err)
	}

	groupID := fp.cfg.GroupID(fp.flowID)
	toolTypeMapping := tools.GetToolTypeMapping()

	var maxCallsLimit int
	switch optAgentType {
	case pconfig.OptionsTypeAssistant, pconfig.OptionsTypePrimaryAgent,
		pconfig.OptionsTypePentester, pconfig.OptionsTypeCoder, pconfig.OptionsTypeInstaller:
		if fp.maxGACallsLimit <= 0 {
			maxCallsLimit = maxGeneralAgentChainIterations
		} else {
			maxCallsLimit = max(fp.maxGACallsLimit, maxAgentShutdownIterations*2)
		}
	default:
		if fp.maxLACallsLimit <= 0 {
			maxCallsLimit = maxLimitedAgentChainIterations
		} else {
			maxCallsLimit = max(fp.maxLACallsLimit, maxAgentShutdownIterations*2)
		}
	}

	for iteration := 0; ; iteration++ {
		if iteration >= maxCallsLimit {
			msg := fmt.Sprintf("agent chain exceeded maximum iterations (%d)", maxCallsLimit)
			logger.WithField("iteration", iteration).Error(msg)
			return fmt.Errorf("%w: %s", ErrAgentModelCall, msg)
		}

		var result *callResult
		if iteration >= maxCallsLimit-maxAgentShutdownIterations {
			logger.WithFields(logrus.Fields{
				"iteration": iteration,
				"limit":     maxCallsLimit,
			}).Warn("max tool calls limit will be reached soon, invoking reflector for graceful termination")

			// Format reflector message for graceful termination
			result = &callResult{
				content: fmt.Sprintf(
					"I can’t continue this multi-turn chain because I’m too close to the AI agent iteration limit (%d).",
					maxCallsLimit,
				),
			}
		} else {
			if summarizer != nil {
				if compacted, done := compactChainForWindow(
					ctx, fp, optAgentType, chain, executor.Tools(), summarizerHandler, fp.tcIDTemplate,
				); done {
					chain = compacted
					if err := fp.updateMsgChain(ctx, optAgentType, chainID, chain, rollLastUpdateTime()); err != nil {
						obs.LogErrorOrCancel(logger, err, "failed to update msg chain")
						return err
					}
				}
			}

			result, err = fp.callWithRetries(ctx, optAgentType, chainID, taskID, subtaskID, chain, executor, executionContext)
			if err != nil {
				obs.LogErrorOrCancel(logger, err, "failed to call agent chain")
				return err
			}

			if err := fp.updateMsgChainUsage(ctx, chainID, optAgentType, result.info, rollLastUpdateTime()); err != nil {
				obs.LogErrorOrCancel(logger, err, "failed to update msg chain usage")
				return err
			}
		}

		if len(result.funcCalls) == 0 {
			if optAgentType == pconfig.OptionsTypeAssistant {
				fp.storeAgentResponseToGraphiti(ctx, groupID, optAgentType, result, taskID, subtaskID, chainID)
				return fp.processAssistantResult(ctx, logger, chainID, chain, result, summarizer, summarizerHandler, rollLastUpdateTime())
			} else {
				// Build AI message with reasoning for reflector (universal pattern)
				reflectorMsg := llms.MessageContent{Role: llms.ChatMessageTypeAI}
				if part, ok := aiResponsePart(result.content, result.thinking); ok {
					reflectorMsg.Parts = append(reflectorMsg.Parts, part)
				}
				result, err = fp.performReflector(
					ctx, optAgentType, chainID, taskID, subtaskID,
					append(chain, reflectorMsg), executor,
					fp.getLastHumanMessage(chain), result.content, executionContext, 1)
				if err != nil {
					fields := logrus.Fields{}
					if result != nil {
						fields["content"] = result.content[:min(1000, len(result.content))]
						if !result.thinking.IsEmpty() {
							fields["thinking"] = result.thinking.Content[:min(1000, len(result.thinking.Content))]
						}
						fields["execution"] = executionContext[:min(1000, len(executionContext))]
					}
					obs.LogErrorOrCancel(logger.WithFields(fields), err, "failed to perform reflector")
					return err
				}
			}
		}

		fp.storeAgentResponseToGraphiti(ctx, groupID, optAgentType, result, taskID, subtaskID, chainID)

		msg := llms.MessageContent{Role: llms.ChatMessageTypeAI}
		// Universal pattern: preserve content with or without reasoning (works for all providers thanks to deduplication)
		if part, ok := aiResponsePart(result.content, result.thinking); ok {
			msg.Parts = append(msg.Parts, part)
		}
		for _, toolCall := range result.funcCalls {
			msg.Parts = append(msg.Parts, toolCall)
		}
		chain = append(chain, msg)

		if err := fp.updateMsgChain(ctx, optAgentType, chainID, chain, rollLastUpdateTime()); err != nil {
			obs.LogErrorOrCancel(logger, err, "failed to update msg chain")
			return err
		}

		for idx, toolCall := range result.funcCalls {
			if toolCall.FunctionCall == nil {
				continue
			}

			funcName := toolCall.FunctionCall.Name
			response, err := fp.execToolCall(
				ctx, optAgentType, chainID, idx, result, monitor, detector, executor, taskID, subtaskID, chain,
			)

			if toolTypeMapping[funcName] != tools.AgentToolType {
				fp.storeToolExecutionToGraphiti(
					ctx, groupID, optAgentType, toolCall, response, err, executor, taskID, subtaskID, chainID,
				)
			}

			if err != nil {
				obs.LogErrorOrCancel(logger.WithFields(logrus.Fields{
					"func_name": funcName,
					"func_args": toolCall.FunctionCall.Arguments,
				}), err, "failed to exec tool call")
				return err
			}

			chain = append(chain, llms.MessageContent{
				Role: llms.ChatMessageTypeTool,
				Parts: []llms.ContentPart{
					llms.ToolCallResponse{
						ToolCallID: toolCall.ID,
						Name:       funcName,
						Content:    response,
					},
				},
			})
			if err := fp.updateMsgChain(ctx, optAgentType, chainID, chain, rollLastUpdateTime()); err != nil {
				obs.LogErrorOrCancel(logger, err, "failed to update msg chain")
				return err
			}

			if executor.IsBarrierFunction(funcName) {
				wantToStop = true
			}
		}

		if wantToStop {
			return nil
		}

		if summarizer != nil {
			// it returns the same chain state if error occurs
			chain, err = summarizer.SummarizeChain(ctx, summarizerHandler, chain, fp.tcIDTemplate)
			if err != nil {
				// log swallowed error
				_, observation := obs.Observer.NewObservation(ctx)
				observation.Event(
					langfuse.WithEventName("chain summarization error swallowed"),
					langfuse.WithEventInput(chain),
					langfuse.WithEventStatus(err.Error()),
					langfuse.WithEventLevel(langfuse.ObservationLevelWarning),
					langfuse.WithEventMetadata(langfuse.Metadata{
						"tc_id_template": fp.tcIDTemplate,
						"msg_chain_id":   chainID,
						"error":          err.Error(),
					}),
				)
				logger.WithError(err).Warn("failed to summarize chain")
			} else if err := fp.updateMsgChain(ctx, optAgentType, chainID, chain, rollLastUpdateTime()); err != nil {
				obs.LogErrorOrCancel(logger, err, "failed to update msg chain")
				return err
			}
		}
	}
}

func (fp *flowProvider) execToolCall(
	ctx context.Context,
	optAgentType pconfig.ProviderOptionsType,
	chainID int64,
	toolCallIDx int,
	result *callResult,
	monitor *executionMonitor,
	detector *repeatingDetector,
	executor tools.ContextToolsExecutor,
	taskID, subtaskID *int64,
	chain []llms.MessageContent,
) (string, error) {
	var (
		streamID int64
		thinking string
	)

	// use streamID and thinking only for first tool call to minimize content
	if toolCallIDx == 0 {
		streamID = result.streamID
		if !result.thinking.IsEmpty() {
			thinking = result.thinking.Content
		}
	}

	toolCall := result.funcCalls[toolCallIDx]
	if toolCall.FunctionCall == nil {
		return "", fmt.Errorf("tool call function call is nil")
	}

	funcName := toolCall.FunctionCall.Name
	funcArgs := json.RawMessage(toolCall.FunctionCall.Arguments)

	logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(fp.flowID, taskID, subtaskID, logrus.Fields{
		"agent":        fp.Type(),
		"func_name":    funcName,
		"func_args":    string(funcArgs)[:min(1000, len(funcArgs))],
		"tool_call_id": toolCall.ID,
		"msg_chain_id": chainID,
	}))

	if detector.detect(toolCall) {
		if len(detector.funcCalls) >= RepeatingToolCallThreshold+maxSoftDetectionsBeforeAbort {
			errMsg := fmt.Sprintf("tool '%s' repeated %d times consecutively, aborting chain", funcName, len(detector.funcCalls))
			logger.WithField("repeat_count", len(detector.funcCalls)).Error(errMsg)
			return "", errors.New(errMsg)
		}

		response := repeatingToolResponse(funcName, len(detector.funcCalls), detector.lastResponse)

		_, observation := obs.Observer.NewObservation(ctx)
		observation.Event(
			langfuse.WithEventName("repeating tool call detected"),
			langfuse.WithEventInput(funcArgs),
			langfuse.WithEventMetadata(map[string]any{
				"tool_call_id": toolCall.ID,
				"tool_name":    funcName,
				"msg_chain_id": chainID,
			}),
			langfuse.WithEventStatus("failed"),
			langfuse.WithEventLevel(langfuse.ObservationLevelError),
			langfuse.WithEventOutput(response),
		)
		logger.Warn("failed to exec function: tool call is repeating")

		return response, nil
	}

	var (
		err      error
		response string
	)

	for idx := 0; idx <= maxRetriesToCallFunction; idx++ {
		if idx == maxRetriesToCallFunction {
			err = fmt.Errorf("reached max retries to call function: %w", err)
			obs.LogErrorOrCancel(logger, err, "failed to exec function")
			return "", fmt.Errorf("failed to exec function '%s': %w", funcName, err)
		}

		response, err = executor.Execute(ctx, streamID, toolCall.ID, funcName, funcName, thinking, funcArgs)
		if err != nil {
			if errors.Is(err, context.Canceled) {
				return "", err
			}

			if errors.Is(err, tools.ErrFlowStateGuard) {
				// Not an argument-format problem: the flow/task/subtask is in a state
				// that makes this call impossible right now (task running, no plan
				// yet, stale subtask, ...). Retrying fixToolCallArgs can't change
				// that, so surface it to the agent immediately instead of burning
				// retries on "corrected" arguments that won't help.
				return "", err
			}

			logger.WithError(err).Warn("failed to exec function")

			funcExecErr := err
			funcSchema, err := executor.GetToolSchema(funcName)
			if err != nil {
				logger.WithError(err).Error("failed to get tool schema")
				return "", fmt.Errorf("failed to get tool schema: %w", err)
			}

			funcArgs, err = fp.fixToolCallArgs(ctx, funcName, funcArgs, funcSchema, funcExecErr)
			if err != nil {
				obs.LogErrorOrCancel(logger, err, "failed to fix tool call args")
				return "", correctorFailure(ctx, err, funcName, funcExecErr)
			}
		} else {
			detector.lastResponse = response
			break
		}
	}

	if monitor.shouldInvokeMentor(toolCall) && executor.IsFunctionExists(tools.AdviceToolName) {
		logger.WithFields(logrus.Fields{
			"same_tool_count":  monitor.sameToolCount,
			"total_call_count": monitor.totalCallCount,
		}).Debug("execution monitor threshold reached, invoking mentor for progress review")

		mentorResponse, err := fp.performMentor(
			ctx, optAgentType, chainID, taskID, subtaskID, chain, executor, toolCall, response,
		)
		if err != nil {
			logger.WithError(err).Warn("failed to invoke execution mentor, continuing with normal execution")
		} else {
			monitor.reset()
			response = formatEnhancedToolResponse(response, mentorResponse)
		}
	}

	return response, nil
}

func (fp *flowProvider) callWithRetries(
	ctx context.Context,
	optAgentType pconfig.ProviderOptionsType,
	chainID int64,
	taskID, subtaskID *int64,
	chain []llms.MessageContent,
	executor tools.ContextToolsExecutor,
	executionContext string,
) (*callResult, error) {
	var (
		err     error
		errs    []error
		msgType = database.MsglogTypeAnswer
		resp    *llms.ContentResponse
		result  callResult
	)

	logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(fp.flowID, taskID, subtaskID, logrus.Fields{
		"agent":        fp.Type(),
		"msg_chain_id": chainID,
		"agent_type":   optAgentType,
	}))

	ticker := time.NewTicker(delayBetweenRetries)
	defer ticker.Stop()

	fillResult := func(resp *llms.ContentResponse) error {
		var stopReason string
		var parts []string

		if resp == nil || len(resp.Choices) == 0 {
			return fmt.Errorf("no choices in response")
		}

		for _, choice := range resp.Choices {
			if stopReason == "" {
				stopReason = choice.StopReason
			}

			if choice.GenerationInfo != nil {
				result.info = choice.GenerationInfo
			}

			// Extract reasoning for logging/analytics (provider-aware)
			if result.thinking.IsEmpty() {
				if !choice.Reasoning.IsEmpty() {
					result.thinking = choice.Reasoning
				} else if len(choice.ToolCalls) > 0 && !choice.ToolCalls[0].Reasoning.IsEmpty() {
					// Gemini puts reasoning in first tool call when tools are used
					result.thinking = choice.ToolCalls[0].Reasoning
				}
			}

			if strings.TrimSpace(choice.Content) != "" {
				parts = append(parts, choice.Content)
			}

			for _, toolCall := range choice.ToolCalls {
				if toolCall.FunctionCall == nil {
					continue
				}
				sanitizedArgs := cast.SanitizeJSONControlChars(toolCall.FunctionCall.Arguments)
				if !json.Valid([]byte(sanitizedArgs)) {
					logger.WithFields(logrus.Fields{
						"tool_call_id": toolCall.ID,
						"tool_name":    toolCall.FunctionCall.Name,
						"raw_args":     toolCall.FunctionCall.Arguments[:min(200, len(toolCall.FunctionCall.Arguments))],
					}).Warn("tool call has invalid JSON arguments, replacing with empty object to allow tool-call fixer to regenerate them")
					sanitizedArgs = "{}"
				}
				toolCall.FunctionCall.Arguments = sanitizedArgs
				result.funcCalls = append(result.funcCalls, toolCall)
			}
		}

		result.content = strings.Join(parts, "\n")
		if strings.Trim(result.content, "' \"\n\r\t") == "" && len(result.funcCalls) == 0 {
			return fmt.Errorf("no content and tool calls in response: stop reason '%s'", stopReason)
		}

		return nil
	}

	for idx := 0; idx <= maxRetriesToCallAgentChain; idx++ {
		if idx == maxRetriesToCallAgentChain {
			reflectorResult, err := fp.performCallerReflector(
				ctx, optAgentType, chainID, taskID, subtaskID, chain, executor, executionContext, errs,
			)
			if err != nil {
				msg := fmt.Sprintf("failed to call agent chain: max retries reached, %d", idx)
				// The final reflector failure determines whether fallback is safe.
				// Earlier model errors are context, not causes: joining them here
				// would make a later database failure look like a model refusal.
				return nil, fmt.Errorf("%s; prior attempts: %v: %w", msg, errors.Join(errs...), err)
			}

			return reflectorResult, nil
		}

		var streamCb streaming.Callback
		if fp.streamCb != nil {
			result.streamID = fp.callCounter.Add(1)
			streamCb = func(ctx context.Context, chunk streaming.Chunk) error {
				switch chunk.Type {
				case streaming.ChunkTypeReasoning:
					if chunk.Reasoning.IsEmpty() {
						return nil
					}
					return fp.streamCb(ctx, &StreamMessageChunk{
						Type:     StreamMessageChunkTypeThinking,
						MsgType:  msgType,
						Thinking: chunk.Reasoning,
						StreamID: result.streamID,
					})
				case streaming.ChunkTypeText:
					return fp.streamCb(ctx, &StreamMessageChunk{
						Type:     StreamMessageChunkTypeContent,
						MsgType:  msgType,
						Content:  chunk.Content,
						StreamID: result.streamID,
					})
				case streaming.ChunkTypeToolCall:
					// skip tool call chunks (we don't need them for now)
				case streaming.ChunkTypeDone:
					return fp.streamCb(ctx, &StreamMessageChunk{
						Type:     StreamMessageChunkTypeFlush,
						MsgType:  msgType,
						StreamID: result.streamID,
					})
				}
				return nil
			}
		}

		resp, err = fp.CallWithTools(ctx, optAgentType, chain, executor.Tools(), streamCb)
		if err != nil {
			err = fmt.Errorf("%w: %w", ErrAgentModelCall, err)
		}
		if isDeterministicError(err) {
			return nil, fmt.Errorf("failed to call agent chain: %w", errors.Join(append(errs, err)...))
		}
		if err == nil {
			err = fillResult(resp)
			if err != nil {
				err = fmt.Errorf("%w: %w", ErrAgentModelCall, err)
			}
		}
		if err == nil {
			break
		} else {
			errs = append(errs, err)

			// Saying "will retry" for a call the loop is about to abandon is
			// what makes this line unreadable in a log. The wrapped ctx.Err() is
			// the same error the select below would return one statement later,
			// which matters: the provider replaces a cancellation with a bare
			// message of its own, and the controllers that unwind a stopped flow
			// test for context.Canceled with errors.Is.
			if ctxErr := ctx.Err(); ctxErr != nil {
				return nil, fmt.Errorf("context canceled while waiting for retry: %w", ctxErr)
			}

			logger.WithFields(logrus.Fields{
				"retry_iteration": idx,
				"error":           err.Error()[:min(200, len(err.Error()))],
			}).Warn("agent chain call failed, will retry")
		}

		ticker.Reset(delayBetweenRetries)
		select {
		case <-ticker.C:
		case <-ctx.Done():
			return nil, fmt.Errorf("context canceled while waiting for retry: %w", ctx.Err())
		}
	}

	if fp.streamCb != nil && result.streamID != 0 {
		_ = fp.streamCb(ctx, &StreamMessageChunk{
			Type:     StreamMessageChunkTypeUpdate,
			MsgType:  msgType,
			Content:  result.content,
			Thinking: result.thinking,
			StreamID: result.streamID,
		})
		// don't update stream by ID if we got content separately from tool calls
		// because we stored thinking and content into standalone messages
		if len(result.funcCalls) > 0 && result.content != "" {
			result.streamID = 0
		}
	}

	return &result, nil
}

func (fp *flowProvider) performReflector(
	ctx context.Context,
	optOriginType pconfig.ProviderOptionsType,
	chainID int64,
	taskID, subtaskID *int64,
	chain []llms.MessageContent,
	executor tools.ContextToolsExecutor,
	humanMessage, content, executionContext string,
	iteration int,
) (*callResult, error) {
	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "providers.flowProvider.performReflector")
	defer span.End()

	var (
		msgChainType = database.MsgchainTypeReflector
		optAgentType = agentByChain[msgChainType]
	)

	logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(fp.flowID, taskID, subtaskID, logrus.Fields{
		"provider":     fp.Type(),
		"agent":        optAgentType,
		"origin":       optOriginType,
		"msg_chain_id": chainID,
		"iteration":    iteration,
	}))

	if iteration > maxReflectorCallsPerChain {
		msg := "reflector called too many times"
		_, observation := obs.Observer.NewObservation(ctx)
		observation.Event(
			langfuse.WithEventName("reflector limit calls reached"),
			langfuse.WithEventInput(content),
			langfuse.WithEventStatus("failed"),
			langfuse.WithEventLevel(langfuse.ObservationLevelError),
			langfuse.WithEventOutput(msg),
			langfuse.WithEventMetadata(map[string]any{
				"iteration": iteration,
			}),
		)
		logger.WithField("content", content[:min(1000, len(content))]).Warn(msg)
		return nil, fmt.Errorf("%w: %s", ErrAgentModelCall, msg)
	}

	logger.WithField("content", content[:min(1000, len(content))]).Warn("got message instead of tool call")

	reflectorContext := map[string]map[string]any{
		"user": {
			"Message":          content,
			"BarrierToolNames": executor.GetBarrierToolNames(),
		},
		"system": {
			"BarrierTools":     executor.GetBarrierTools(),
			"CurrentTime":      getCurrentTime(),
			"ExecutionContext": executionContext,
		},
	}

	if humanMessage != "" {
		reflectorContext["system"]["Request"] = humanMessage
	}

	ctx, observation := obs.Observer.NewObservation(ctx)
	reflectorAgent := observation.Agent(
		langfuse.WithAgentName("reflector"),
		langfuse.WithAgentInput(content),
		langfuse.WithAgentMetadata(langfuse.Metadata{
			"user_context":   reflectorContext["user"],
			"system_context": reflectorContext["system"],
		}),
	)
	ctx, observation = reflectorAgent.Observation(ctx)

	reflectorEvaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("render reflector agent prompts"),
		langfuse.WithEvaluatorInput(reflectorContext),
		langfuse.WithEvaluatorMetadata(langfuse.Metadata{
			"user_context":   reflectorContext["user"],
			"system_context": reflectorContext["system"],
			"lang":           fp.language,
		}),
	)

	userReflectorTmpl, err := fp.prompter.RenderTemplate(templates.PromptTypeQuestionReflector, reflectorContext["user"])
	if err != nil {
		msg := "failed to get user reflector template"
		return nil, wrapErrorEndEvaluatorSpan(ctx, reflectorEvaluator, msg, err)
	}

	systemReflectorTmpl, err := fp.prompter.RenderTemplate(templates.PromptTypeReflector, reflectorContext["system"])
	if err != nil {
		msg := "failed to get system reflector template"
		return nil, wrapErrorEndEvaluatorSpan(ctx, reflectorEvaluator, msg, err)
	}

	reflectorEvaluator.End(
		langfuse.WithEvaluatorOutput(map[string]any{
			"user_template":   userReflectorTmpl,
			"system_template": systemReflectorTmpl,
		}),
		langfuse.WithEvaluatorStatus("success"),
		langfuse.WithEvaluatorLevel(langfuse.ObservationLevelDebug),
	)

	advice, err := fp.performSimpleChain(ctx, taskID, subtaskID, optAgentType,
		msgChainType, systemReflectorTmpl, userReflectorTmpl)
	if err != nil {
		advice = ToolPlaceholder
	}

	opts := []langfuse.AgentOption{
		langfuse.WithAgentStatus("failed"),
		langfuse.WithAgentOutput(advice),
		langfuse.WithAgentLevel(langfuse.ObservationLevelWarning),
	}
	defer func() {
		reflectorAgent.End(opts...)
	}()

	chain = append(chain, llms.TextParts(llms.ChatMessageTypeHuman, advice))
	result, err := fp.callWithRetries(ctx, optOriginType, chainID, taskID, subtaskID, chain, executor, executionContext)
	if err != nil {
		obs.LogErrorOrCancel(logger, err, "failed to call agent chain by reflector")
		level := langfuse.ObservationLevelError
		if errors.Is(err, context.Canceled) {
			level = langfuse.ObservationLevelWarning
		}
		opts = append(opts,
			langfuse.WithAgentStatus(err.Error()),
			langfuse.WithAgentLevel(level),
		)
		return nil, err
	}

	// don't update duration delta for reflector because it's already included in the performAgentChain
	if err := fp.updateMsgChainUsage(ctx, chainID, optOriginType, result.info, 0); err != nil {
		logger.WithError(err).Error("failed to update msg chain usage")
		opts = append(opts,
			langfuse.WithAgentStatus(err.Error()),
			langfuse.WithAgentLevel(langfuse.ObservationLevelError),
		)
		return nil, err
	}

	// preserve reasoning in reflector response using universal pattern
	reflectorMsg := llms.MessageContent{Role: llms.ChatMessageTypeAI}
	if part, ok := aiResponsePart(result.content, result.thinking); ok {
		reflectorMsg.Parts = append(reflectorMsg.Parts, part)
	}
	chain = append(chain, reflectorMsg)
	if len(result.funcCalls) == 0 {
		// Check if we are already in a reflector retry cycle to prevent infinite recursion.
		// This blocks recursive performReflector calls after caller reflector was invoked.
		if isReflectorRetry(ctx) {
			logger.Error("reflector recursion detected: cannot recursively call reflector after caller reflector")
			return nil, fmt.Errorf("%w: reflector recursion detected: LLM returned no tool calls after reflector advice", ErrAgentModelCall)
		}

		return fp.performReflector(ctx, optOriginType, chainID, taskID, subtaskID, chain, executor,
			humanMessage, result.content, executionContext, iteration+1)
	}

	opts = append(opts, langfuse.WithAgentStatus("success"))
	return result, nil
}

func (fp *flowProvider) performCallerReflector(
	ctx context.Context,
	optAgentType pconfig.ProviderOptionsType,
	chainID int64,
	taskID, subtaskID *int64,
	chain []llms.MessageContent,
	executor tools.ContextToolsExecutor,
	executionContext string,
	errs []error,
) (*callResult, error) {
	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "providers.flowProvider.performCallerReflector")
	defer span.End()

	logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(fp.flowID, taskID, subtaskID, logrus.Fields{
		"provider":     fp.Type(),
		"agent":        optAgentType,
		"msg_chain_id": chainID,
		"errors_count": len(errs),
	})).WithError(errors.Join(errs...))

	// Check if we are already in a reflector retry cycle to prevent infinite recursion.
	// This blocks repeated calls to performCallerReflector after reflector advice failed.
	if isReflectorRetry(ctx) {
		logger.Error("reflector recursion detected: caller reflector already invoked in this chain")
		return nil, fmt.Errorf("%w: reflector recursion detected: cannot invoke caller reflector again after reflector advice failed", ErrAgentModelCall)
	}

	// Mark context to prevent any further reflector recursion.
	// This flag will be checked in:
	// 1. performCallerReflector (here) - if reflector advice fails again
	// 2. performReflector - before recursive call when no tool calls returned
	ctx = markReflectorRetry(ctx)
	logger = logger.WithContext(ctx)

	logger.Warn("max retries reached, invoking caller reflector for guidance")

	reflectorContent := fmt.Sprintf(
		"I'm having trouble generating a proper tool call response. "+
			"I've attempted %d times but each attempt failed with errors:\n\n%s\n\n"+
			"I'm not sure how to proceed correctly. Should I try a different approach, "+
			"or should I use one of the barrier tools to report this issue?",
		len(errs), errors.Join(errs...).Error(),
	)

	reflectorResult, err := fp.performReflector(
		ctx, optAgentType, chainID, taskID, subtaskID, chain, executor,
		fp.getLastHumanMessage(chain), reflectorContent, executionContext, 1,
	)
	if err == nil {
		return reflectorResult, nil
	}

	return nil, fmt.Errorf("failed to perform caller reflector: %w", err)
}

func (fp *flowProvider) getLastHumanMessage(chain []llms.MessageContent) string {
	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return ""
	}

	slices.Reverse(ast.Sections)
	for _, section := range ast.Sections {
		if section.Header.HumanMessage != nil {
			var hparts []string
			for _, part := range section.Header.HumanMessage.Parts {
				if text, ok := part.(llms.TextContent); ok {
					hparts = append(hparts, text.Text)
				}
			}
			return strings.Join(hparts, "\n")
		}
	}

	return ""
}

func (fp *flowProvider) processAssistantResult(
	ctx context.Context,
	logger *logrus.Entry,
	chainID int64,
	chain []llms.MessageContent,
	result *callResult,
	summarizer csum.Summarizer,
	summarizerHandler tools.SummarizeHandler,
	durationDelta float64,
) error {
	var err error

	processAssistantResultStartTime := time.Now()

	if fp.streamCb != nil {
		if result.streamID == 0 {
			result.streamID = fp.callCounter.Add(1)
		}
		err := fp.streamCb(ctx, &StreamMessageChunk{
			Type:     StreamMessageChunkTypeUpdate,
			MsgType:  database.MsglogTypeAnswer,
			Content:  result.content,
			Thinking: result.thinking,
			StreamID: result.streamID,
		})
		if err != nil {
			return fmt.Errorf("failed to stream assistant result: %w", err)
		}
	}

	if summarizer != nil {
		// it returns the same chain state if error occurs
		chain, err = summarizer.SummarizeChain(ctx, summarizerHandler, chain, fp.tcIDTemplate)
		if err != nil {
			// log swallowed error
			_, observation := obs.Observer.NewObservation(ctx)
			observation.Event(
				langfuse.WithEventName("chain summarization error swallowed"),
				langfuse.WithEventInput(chain),
				langfuse.WithEventStatus(err.Error()),
				langfuse.WithEventLevel(langfuse.ObservationLevelWarning),
				langfuse.WithEventMetadata(langfuse.Metadata{
					"tc_id_template": fp.tcIDTemplate,
					"msg_chain_id":   chainID,
					"error":          err.Error(),
				}),
			)
			logger.WithError(err).Warn("failed to summarize chain")
		}
	}

	// Preserve reasoning for assistant responses using universal pattern
	msg := llms.MessageContent{Role: llms.ChatMessageTypeAI}
	if part, ok := aiResponsePart(result.content, result.thinking); ok {
		msg.Parts = append(msg.Parts, part)
	}
	chain = append(chain, msg)
	durationDelta += time.Since(processAssistantResultStartTime).Seconds()
	if err := fp.updateMsgChain(ctx, pconfig.OptionsTypeAssistant, chainID, chain, durationDelta); err != nil {
		return fmt.Errorf("failed to update msg chain: %w", err)
	}

	return nil
}

func (fp *flowProvider) updateMsgChain(
	ctx context.Context,
	optAgentType pconfig.ProviderOptionsType,
	chainID int64,
	chain []llms.MessageContent,
	durationDelta float64,
) error {
	chainBlob, err := json.Marshal(chain)
	if err != nil {
		return fmt.Errorf("failed to marshal msg chain: %w", err)
	}

	_, err = fp.db.UpdateMsgChain(ctx, database.UpdateMsgChainParams{
		Chain:           chainBlob,
		DurationSeconds: durationDelta,
		Model:           fp.Model(optAgentType),
		ModelProvider:   string(fp.Type()),
		ID:              chainID,
	})
	if err != nil {
		return fmt.Errorf("failed to update msg chain in DB: %w", err)
	}

	return nil
}

func (fp *flowProvider) updateMsgChainUsage(
	ctx context.Context,
	chainID int64,
	optAgentType pconfig.ProviderOptionsType,
	info map[string]any,
	durationDelta float64,
) error {
	usage := fp.GetUsage(info)
	if usage.IsZero() {
		return nil
	}

	price := fp.GetPriceInfo(optAgentType)
	if price != nil {
		usage.UpdateCost(price)
	}

	_, err := fp.db.UpdateMsgChainUsage(ctx, database.UpdateMsgChainUsageParams{
		UsageIn:         usage.Input,
		UsageOut:        usage.Output,
		UsageCacheIn:    usage.CacheRead,
		UsageCacheOut:   usage.CacheWrite,
		UsageCostIn:     usage.CostInput,
		UsageCostOut:    usage.CostOutput,
		DurationSeconds: durationDelta,
		ID:              chainID,
	})
	if err != nil {
		return fmt.Errorf("failed to update msg chain usage in DB: %w", err)
	}

	return nil
}

// storeToGraphiti stores messages to Graphiti with timeout
func (fp *flowProvider) storeToGraphiti(
	ctx context.Context,
	observation langfuse.Observation,
	groupID string,
	messages []graphiti.Message,
) error {
	if fp.graphitiClient == nil || !fp.graphitiClient.IsEnabled() {
		return nil
	}

	storeCtx, cancel := context.WithTimeout(ctx, fp.graphitiClient.GetTimeout())
	defer cancel()

	err := fp.graphitiClient.AddMessages(storeCtx, graphiti.AddMessagesRequest{
		GroupID:  groupID,
		Messages: messages,
		Observation: &graphiti.Observation{
			ID:      observation.ID(),
			TraceID: observation.TraceID(),
			Time:    time.Now().UTC(),
		},
	})
	if err != nil {
		logrus.WithError(err).
			WithField("group_id", groupID).
			Warn("failed to store messages to graphiti")
	}

	return err
}

// storeAgentResponseToGraphiti stores agent response to Graphiti
func (fp *flowProvider) storeAgentResponseToGraphiti(
	ctx context.Context,
	groupID string,
	agentType pconfig.ProviderOptionsType,
	result *callResult,
	taskID, subtaskID *int64,
	chainID int64,
) {
	if fp.graphitiClient == nil || !fp.graphitiClient.IsEnabled() {
		return
	}

	if result.content == "" {
		return
	}

	tmpl, err := templates.ReadGraphitiTemplate("agent_response.tmpl")
	if err != nil {
		logrus.WithError(err).Warn("failed to read agent response template for graphiti")
		return
	}

	content, err := templates.RenderPrompt("agent_response", tmpl, map[string]any{
		"AgentType": string(agentType),
		"Response":  result.content,
		"TaskID":    templateScopeID(taskID),
		"SubtaskID": templateScopeID(subtaskID),
	})
	if err != nil {
		logrus.WithError(err).Warn("failed to render agent response template for graphiti")
		return
	}

	parts := []string{fmt.Sprintf("PentAGI %s agent execution in flow %d", agentType, fp.flowID)}
	if taskID != nil {
		parts = append(parts, fmt.Sprintf("task %d", *taskID))
	}
	if subtaskID != nil {
		parts = append(parts, fmt.Sprintf("subtask %d", *subtaskID))
	}
	sourceDescription := strings.Join(parts, ", ")

	messages := []graphiti.Message{
		{
			Content:           content,
			Author:            fmt.Sprintf("%s Agent", string(agentType)),
			Timestamp:         time.Now(),
			Name:              "agent_response",
			SourceDescription: sourceDescription,
		},
	}
	logrus.WithField("messages", messages).Debug("storing agent response to graphiti")

	ctx, observation := obs.Observer.NewObservation(ctx)
	storeEvaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("store messages to graphiti"),
		langfuse.WithEvaluatorInput(messages),
		langfuse.WithEvaluatorMetadata(langfuse.Metadata{
			"group_id":     groupID,
			"agent_type":   agentType,
			"task_id":      taskID,
			"subtask_id":   subtaskID,
			"msg_chain_id": chainID,
		}),
	)

	ctx, observation = storeEvaluator.Observation(ctx)
	if err := fp.storeToGraphiti(ctx, observation, groupID, messages); err != nil {
		storeEvaluator.End(
			langfuse.WithEvaluatorStatus(err.Error()),
			langfuse.WithEvaluatorLevel(langfuse.ObservationLevelError),
		)
		return
	}

	storeEvaluator.End(
		langfuse.WithEvaluatorStatus("success"),
	)
}

// storeToolExecutionToGraphiti stores tool execution to Graphiti
func (fp *flowProvider) storeToolExecutionToGraphiti(
	ctx context.Context,
	groupID string,
	agentType pconfig.ProviderOptionsType,
	toolCall llms.ToolCall,
	response string,
	execErr error,
	executor tools.ContextToolsExecutor,
	taskID, subtaskID *int64,
	chainID int64,
) {
	if fp.graphitiClient == nil || !fp.graphitiClient.IsEnabled() {
		return
	}

	if toolCall.FunctionCall == nil {
		return
	}

	funcName := toolCall.FunctionCall.Name
	funcArgs := toolCall.FunctionCall.Arguments

	registryDefs := tools.GetRegistryDefinitions()
	toolDef, ok := registryDefs[funcName]
	description := ""
	if ok {
		description = toolDef.Description
	}

	isBarrier := executor.IsBarrierFunction(funcName)

	status := "success"
	if execErr != nil {
		status = "failure"
		response = fmt.Sprintf("Error: %s", execErr.Error())
	}

	toolExecTmpl, err := templates.ReadGraphitiTemplate("tool_execution.tmpl")
	if err != nil {
		logrus.WithError(err).Warn("failed to read tool execution template for graphiti")
		return
	}

	toolExecContent, err := templates.RenderPrompt("tool_execution", toolExecTmpl, map[string]any{
		"ToolName":    funcName,
		"Description": description,
		"IsBarrier":   isBarrier,
		"Arguments":   funcArgs,
		"AgentType":   string(agentType),
		"Status":      status,
		"Result":      response,
		"TaskID":      templateScopeID(taskID),
		"SubtaskID":   templateScopeID(subtaskID),
	})
	if err != nil {
		logrus.WithError(err).Warn("failed to render tool execution template for graphiti")
		return
	}

	parts := []string{fmt.Sprintf("PentAGI tool execution in flow %d", fp.flowID)}
	if taskID != nil {
		parts = append(parts, fmt.Sprintf("task %d", *taskID))
	}
	if subtaskID != nil {
		parts = append(parts, fmt.Sprintf("subtask %d", *subtaskID))
	}
	sourceDescription := strings.Join(parts, ", ")

	messages := []graphiti.Message{
		{
			Content:           toolExecContent,
			Author:            fmt.Sprintf("%s Agent", string(agentType)),
			Timestamp:         time.Now(),
			Name:              fmt.Sprintf("tool_execution_%s", funcName),
			SourceDescription: sourceDescription,
		},
	}

	ctx, observation := obs.Observer.NewObservation(ctx)
	storeEvaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("store tool execution to graphiti"),
		langfuse.WithEvaluatorInput(messages),
		langfuse.WithEvaluatorMetadata(langfuse.Metadata{
			"group_id":     groupID,
			"agent_type":   agentType,
			"tool_name":    funcName,
			"tool_args":    funcArgs,
			"task_id":      taskID,
			"subtask_id":   subtaskID,
			"msg_chain_id": chainID,
		}),
	)

	ctx, observation = storeEvaluator.Observation(ctx)
	if err := fp.storeToGraphiti(ctx, observation, groupID, messages); err != nil {
		storeEvaluator.End(
			langfuse.WithEvaluatorStatus(err.Error()),
			langfuse.WithEvaluatorLevel(langfuse.ObservationLevelError),
		)
		return
	}

	storeEvaluator.End(
		langfuse.WithEvaluatorStatus("success"),
	)
}

// templateScopeID keeps a missing id out of the graphiti templates: text/template prints a nil
// pointer as "<nil>", and an assistant turn has no task, a generator or refiner run no subtask.
func templateScopeID(id *int64) string {
	if id == nil {
		return "-"
	}
	return strconv.FormatInt(*id, 10)
}

func aiResponsePart(content string, thinking *reasoning.ContentReasoning) (llms.ContentPart, bool) {
	if !thinking.IsEmpty() {
		return llms.TextPartWithReasoning(content, thinking), true
	}

	if content != "" {
		return llms.TextPart(content), true
	}

	return nil, false
}

func correctorFailure(ctx context.Context, correctorErr error, funcName string, funcExecErr error) error {
	if errors.Is(correctorErr, context.Canceled) || ctx.Err() != nil {
		return correctorErr
	}
	return fmt.Errorf("failed to exec function '%s' and could not correct its arguments: %w", funcName, funcExecErr)
}