package providers

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

	"pentagi/pkg/cast"
	"pentagi/pkg/csum"
	"pentagi/pkg/database"
	"pentagi/pkg/docker"
	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"
)

const (
	RepeatingToolCallThreshold   = 3
	maxQASectionsAfterRestore    = 3
	keepQASectionsAfterRestore   = 1
	lastSecBytesAfterRestore     = 16 * 1024 // 16 KB
	maxBPBytesAfterRestore       = 8 * 1024  // 8 KB
	maxQABytesAfterRestore       = 20 * 1024 // 20 KB
	msgLogResultSummarySizeLimit = 70 * 1024 // 70 KB
	msgLogResultEntrySizeLimit   = 1024      // 1 KB
	extractLastMessagesCount     = 30
	extractToolCallsCount        = 10
	toolCallsHistorySeparator    = "---------------TOOL_CALLS_HISTORY---------------"
	ChainContinuationMessage     = "Continue from where we left off"
)

type dummyMessage struct {
	Message string `json:"message"`
}

func toolCallIDTemplateOrDefault(template string) string {
	if template != "" {
		return template
	}

	return cast.ToolCallIDTemplate
}

func wrapToolCallIDTemplateError(err error) error {
	if err == nil {
		return nil
	}
	if strings.Contains(err.Error(), "does not support tools") {
		return fmt.Errorf(
			"failed to determine tool call ID template: the selected model "+
				"does not support tool/function calling, which PentAGI requires "+
				"for tool execution in flows and assistant sessions; select an "+
				"Ollama model whose metadata advertises tool/function calling "+
				"support and update the provider configuration: %w",
			err)
	}
	return fmt.Errorf("failed to determine tool call ID template: %w", err)
}

type reflectorRetryContextKey struct{}

// isReflectorRetry checks if we are already in a reflector retry cycle
func isReflectorRetry(ctx context.Context) bool {
	if isRetry, ok := ctx.Value(reflectorRetryContextKey{}).(bool); ok {
		return isRetry
	}
	return false
}

// markReflectorRetry marks context as being in a reflector retry cycle
func markReflectorRetry(ctx context.Context) context.Context {
	return context.WithValue(ctx, reflectorRetryContextKey{}, true)
}

type repeatingDetector struct {
	funcCalls []llms.FunctionCall
	// lastResponse is what the tool answered on the last call that reached it,
	// kept so a repeat can hand that answer back instead of replacing it with a
	// notice. The answer is usually the only thing that explains the repetition
	// -- a tool that does not exist says so, and the model never saw it.
	lastResponse string
}

func (rd *repeatingDetector) detect(toolCall llms.ToolCall) bool {
	if toolCall.FunctionCall == nil {
		return false
	}

	funcCall := rd.clearCallArguments(toolCall.FunctionCall)

	if len(rd.funcCalls) == 0 {
		rd.funcCalls = append(rd.funcCalls, funcCall)
		return false
	}

	lastToolCall := rd.funcCalls[len(rd.funcCalls)-1]
	if lastToolCall.Name != funcCall.Name || lastToolCall.Arguments != funcCall.Arguments {
		rd.funcCalls = []llms.FunctionCall{funcCall}
		rd.lastResponse = ""
		return false
	}

	rd.funcCalls = append(rd.funcCalls, funcCall)

	return len(rd.funcCalls) >= RepeatingToolCallThreshold
}

// repeatingToolResponse is what a repeated tool call is answered with.
//
// It hands back the answer the tool actually gave, when there is one. Replacing
// that answer with a notice costs the model the diagnosis: in the case that
// produced this guard's only observed aborts, the suppressed answer was the
// executor's own "not found in available tools list", so the model was told
// seven times that it was repeating and never once what was wrong.
func repeatingToolResponse(funcName string, repeats int, previous string) string {
	if previous == "" {
		return fmt.Sprintf("tool call '%s' is repeating, please try another tool", funcName)
	}

	return fmt.Sprintf(
		"you have called '%s' with identical arguments %d times and it answered the same way each "+
			"time. Here is that answer again:\n\n%s\n\nIt will not change. Alter the arguments, "+
			"use a different tool, or answer from what you already have.",
		funcName, repeats, previous[:min(len(previous), repeatEchoLimit)])
}

func (rd *repeatingDetector) clearCallArguments(toolCall *llms.FunctionCall) llms.FunctionCall {
	var v map[string]any
	if err := json.Unmarshal([]byte(toolCall.Arguments), &v); err != nil {
		return *toolCall
	}

	delete(v, "message")
	var keys []string
	for k := range v {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var buffer strings.Builder
	for _, k := range keys {
		fmt.Fprintf(&buffer, "%s: %v\n", k, v[k])
	}

	return llms.FunctionCall{
		Name:      toolCall.Name,
		Arguments: buffer.String(),
	}
}

type executionMonitorBuilder func() *executionMonitor

// executionMonitor detects when to invoke mentor (adviser agent) for execution monitoring
type executionMonitor struct {
	sameToolCount  int
	totalCallCount int
	lastToolName   string
	sameThreshold  int
	totalThreshold int
	enabled        bool
}

// shouldInvokeMentor checks if mentor (adviser agent) should be invoked based on tool call patterns
func (emd *executionMonitor) shouldInvokeMentor(toolCall llms.ToolCall) bool {
	if !emd.enabled || toolCall.FunctionCall == nil {
		return false
	}

	emd.totalCallCount++

	if toolCall.FunctionCall.Name == emd.lastToolName {
		emd.sameToolCount++
	} else {
		emd.sameToolCount = 1
		emd.lastToolName = toolCall.FunctionCall.Name
	}

	return emd.sameToolCount >= emd.sameThreshold || emd.totalCallCount >= emd.totalThreshold
}

// reset resets the execution monitor state after mentor (adviser agent) invocation
func (emd *executionMonitor) reset() {
	emd.sameToolCount = 0
	emd.totalCallCount = 0
	emd.lastToolName = ""
}

func (fp *flowProvider) getTasksInfo(ctx context.Context, taskID int64) (*tasksInfo, error) {
	var (
		err  error
		info tasksInfo
	)

	ctx, observation := obs.Observer.NewObservation(ctx)
	evaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("get tasks info"),
		langfuse.WithEvaluatorInput(map[string]any{
			"task_id": taskID,
		}),
	)
	ctx, _ = evaluator.Observation(ctx)

	info.Tasks, err = fp.db.GetFlowTasks(ctx, fp.flowID)
	if err != nil {
		return nil, wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to get flow tasks", err)
	}

	for idx, t := range info.Tasks {
		if t.ID == taskID {
			info.Task = t
			info.Tasks = append(info.Tasks[:idx], info.Tasks[idx+1:]...)
			break
		}
	}

	info.Subtasks, err = fp.db.GetFlowSubtasks(ctx, fp.flowID)
	if err != nil {
		return nil, wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to get flow subtasks", err)
	}

	evaluator.End(
		langfuse.WithEvaluatorOutput(map[string]any{
			"task":           info.Task,
			"subtasks":       info.Subtasks,
			"tasks_count":    len(info.Tasks),
			"subtasks_count": len(info.Subtasks),
		}),
		langfuse.WithEvaluatorStatus("success"),
		langfuse.WithEvaluatorLevel(langfuse.ObservationLevelDebug),
	)

	return &info, nil
}

func (fp *flowProvider) getSubtasksInfo(taskID int64, subtasks []database.Subtask) *subtasksInfo {
	var info subtasksInfo
	for _, subtask := range subtasks {
		if subtask.TaskID != taskID && taskID != 0 {
			continue
		}

		switch subtask.Status {
		case database.SubtaskStatusCreated:
			info.Planned = append(info.Planned, subtask)
		case database.SubtaskStatusFinished, database.SubtaskStatusFailed:
			info.Completed = append(info.Completed, subtask)
		default:
			info.Subtask = &subtask
		}
	}

	return &info
}

func (fp *flowProvider) updateMsgChainResult(chain []llms.MessageContent, name, result string) ([]llms.MessageContent, error) {
	if len(chain) == 0 {
		return []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, result)}, nil
	}

	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return nil, fmt.Errorf("failed to create chain ast: %w", err)
	}

	lastSection := ast.Sections[len(ast.Sections)-1]
	if len(lastSection.Body) == 0 {
		ast.AppendHumanMessage(result)
		return ast.Messages(), nil
	}

	lastBody := lastSection.Body[len(lastSection.Body)-1]
	switch lastBody.Type {
	case cast.Completion, cast.Summarization:
		ast.AppendHumanMessage(result)
		return ast.Messages(), nil
	case cast.RequestResponse:
		for _, msg := range lastBody.ToolMessages {
			for pdx, part := range msg.Parts {
				toolCallResp, ok := part.(llms.ToolCallResponse)
				if !ok {
					continue
				}

				if toolCallResp.Name == name {
					toolCallResp.Content = result
					msg.Parts[pdx] = toolCallResp
					return ast.Messages(), nil
				}
			}
		}

		ast.AppendHumanMessage(result)
		return ast.Messages(), nil
	default:
		return nil, fmt.Errorf("unknown message type: %d", lastBody.Type)
	}
}

// Makes chain consistent by adding default responses for any pending tool calls
func (fp *flowProvider) ensureChainConsistency(chain []llms.MessageContent) ([]llms.MessageContent, error) {
	if len(chain) == 0 {
		return chain, nil
	}

	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return nil, fmt.Errorf("failed to create chain ast: %w", err)
	}

	return ast.Messages(), nil
}

func (fp *flowProvider) getTaskPrimaryAgentChainSummary(
	ctx context.Context,
	taskID int64,
	summarizerHandler tools.SummarizeHandler,
) (string, error) {
	ctx, observation := obs.Observer.NewObservation(ctx)
	evaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("get task primary agent chain summary"),
		langfuse.WithEvaluatorInput(map[string]any{
			"task_id": taskID,
		}),
	)
	ctx, _ = evaluator.Observation(ctx)

	msgChain, err := fp.db.GetFlowTaskTypeLastMsgChain(ctx, database.GetFlowTaskTypeLastMsgChainParams{
		FlowID: fp.flowID,
		TaskID: database.Int64ToNullInt64(&taskID),
		Type:   database.MsgchainTypePrimaryAgent,
	})
	if err != nil || isEmptyChain(msgChain.Chain) {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to get task primary agent chain", err)
	}

	chain := []llms.MessageContent{}
	if err := json.Unmarshal(msgChain.Chain, &chain); err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to unmarshal task primary agent chain", err)
	}

	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to create refiner chain ast", err)
	}

	var humanMessages, aiMessages []llms.MessageContent
	for _, section := range ast.Sections {
		if section.Header.HumanMessage != nil {
			humanMessages = append(humanMessages, *section.Header.HumanMessage)
		}
		for _, pair := range section.Body {
			aiMessages = append(aiMessages, pair.Messages()...)
		}
	}

	humanSummary, err := csum.GenerateSummary(ctx, summarizerHandler, humanMessages, nil)
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to generate human summary", err)
	}

	aiSummary, err := csum.GenerateSummary(ctx, summarizerHandler, humanMessages, aiMessages)
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to generate ai summary", err)
	}

	summary := fmt.Sprintf(`## Task Summary

### User Requirements
*Summarized input from user:*

%s

### Execution Results
*Summarized actions and outcomes:*

%s`, humanSummary, aiSummary)

	evaluator.End(
		langfuse.WithEvaluatorOutput(summary),
		langfuse.WithEvaluatorStatus("success"),
		langfuse.WithEvaluatorLevel(langfuse.ObservationLevelDebug),
	)

	return summary, nil
}

func (fp *flowProvider) getTaskMsgLogsSummary(
	ctx context.Context,
	taskID int64,
	summarizerHandler tools.SummarizeHandler,
) (string, error) {
	ctx, observation := obs.Observer.NewObservation(ctx)
	evaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("get task msg logs summary"),
		langfuse.WithEvaluatorInput(map[string]any{
			"task_id": taskID,
			"flow_id": fp.flowID,
		}),
	)
	ctx, _ = evaluator.Observation(ctx)

	msgLogs, err := fp.db.GetTaskMsgLogs(ctx, database.Int64ToNullInt64(&taskID))
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to get task msg logs", err)
	}

	if len(msgLogs) == 0 {
		evaluator.End(
			langfuse.WithEvaluatorOutput("no msg logs"),
			langfuse.WithEvaluatorStatus("success"),
			langfuse.WithEvaluatorLevel(langfuse.ObservationLevelDebug),
		)
		return "no msg logs", nil
	}

	// truncate msg logs result to cut down the size the message to summarize
	for _, msgLog := range msgLogs {
		if len(msgLog.Result) > msgLogResultEntrySizeLimit {
			msgLog.Result = msgLog.Result[:msgLogResultEntrySizeLimit] + textTruncateMessage
		}
	}

	message, err := fp.prompter.RenderTemplate(templates.PromptTypeExecutionLogs, map[string]any{
		"MsgLogs": msgLogs,
	})
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to render task msg logs template", err)
	}

	for l := len(msgLogs) / 2; l > 2; l /= 2 {
		if len(message) < msgLogResultSummarySizeLimit {
			break
		}

		msgLogs = msgLogs[len(msgLogs)-l:]
		message, err = fp.prompter.RenderTemplate(templates.PromptTypeExecutionLogs, map[string]any{
			"MsgLogs": msgLogs,
		})
		if err != nil {
			return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to render task msg logs template", err)
		}
	}

	summary, err := summarizerHandler(ctx, message)
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to summarize task msg logs", err)
	}

	evaluator.End(
		langfuse.WithEvaluatorOutput(summary),
		langfuse.WithEvaluatorStatus("success"),
		langfuse.WithEvaluatorLevel(langfuse.ObservationLevelDebug),
	)

	return summary, nil
}

func (fp *flowProvider) restoreChain(
	ctx context.Context,
	taskID, subtaskID *int64,
	optAgentType pconfig.ProviderOptionsType,
	msgChainType database.MsgchainType,
	systemPrompt, humanPrompt string,
) (int64, []llms.MessageContent, error) {
	ctx, observation := obs.Observer.NewObservation(ctx)

	// Get raw chain from DB for observation input
	msgChain, err := fp.db.GetFlowTaskTypeLastMsgChain(ctx, database.GetFlowTaskTypeLastMsgChainParams{
		FlowID: fp.flowID,
		TaskID: database.Int64ToNullInt64(taskID),
		Type:   msgChainType,
	})

	var rawChain []llms.MessageContent
	if err == nil && !isEmptyChain(msgChain.Chain) {
		_ = json.Unmarshal(msgChain.Chain, &rawChain)
	}

	metadata := langfuse.Metadata{
		"msg_chain_type": string(msgChainType),
		"msg_chain_id":   msgChain.ID,
		"agent_type":     string(optAgentType),
	}
	if taskID != nil {
		metadata["task_id"] = *taskID
	}
	if subtaskID != nil {
		metadata["subtask_id"] = *subtaskID
	}

	chainObs := observation.Chain(
		langfuse.WithChainName("restore message chain"),
		langfuse.WithChainInput(rawChain),
		langfuse.WithChainMetadata(metadata),
	)
	ctx, observation = chainObs.Observation(ctx)
	wrapErrorWithEvent := func(msg string, err error) error {
		observation.Event(
			langfuse.WithEventName("error on restoring message chain"),
			langfuse.WithEventInput(rawChain),
			langfuse.WithEventMetadata(metadata),
			langfuse.WithEventStatus(err.Error()),
			langfuse.WithEventLevel(langfuse.ObservationLevelWarning),
		)

		if err != nil {
			logrus.WithContext(ctx).WithError(err).Warn(msg)
			return fmt.Errorf("%s: %w", msg, err)
		}

		logrus.WithContext(ctx).Warn(msg)
		return errors.New(msg)
	}

	var chain []llms.MessageContent
	fallback := func() {
		chain = []llms.MessageContent{
			llms.TextParts(llms.ChatMessageTypeSystem, systemPrompt),
		}
		if humanPrompt != "" {
			chain = append(chain, llms.TextParts(llms.ChatMessageTypeHuman, humanPrompt))
		}
	}

	if err != nil || isEmptyChain(msgChain.Chain) {
		fallback()
	} else {
		err = func() error {
			err = json.Unmarshal(msgChain.Chain, &chain)
			if err != nil {
				return wrapErrorWithEvent("failed to unmarshal msg chain", err)
			}

			ast, err := cast.NewChainAST(chain, true)
			if err != nil {
				return wrapErrorWithEvent("failed to create refiner chain ast", err)
			}

			if len(ast.Sections) == 0 {
				return wrapErrorWithEvent("failed to get sections from refiner chain ast", nil)
			}

			systemMessage := llms.TextParts(llms.ChatMessageTypeSystem, systemPrompt)
			ast.Sections[0].Header.SystemMessage = &systemMessage
			if humanPrompt != "" {
				lastSection := ast.Sections[len(ast.Sections)-1]
				if len(lastSection.Body) == 0 {
					// do not add a new human message if the previous human message is not yet completed
					lastSection.Header.HumanMessage = nil
				} else {
					lastBody := lastSection.Body[len(lastSection.Body)-1]
					if lastBody.Type == cast.RequestResponse && len(lastBody.ToolMessages) == 0 {
						// prevent using incomplete chain without tool call response
						lastSection.Body = lastSection.Body[:len(lastSection.Body)-1]
					}
				}
				ast.AppendHumanMessage(humanPrompt)
			}

			if err := ast.NormalizeToolCallIDs(fp.tcIDTemplate); err != nil {
				return wrapErrorWithEvent("failed to normalize tool call IDs", err)
			}

			ast.SanitizeToolCallArguments()

			if err := ast.ClearReasoning(); err != nil {
				return wrapErrorWithEvent("failed to clear reasoning", err)
			}

			summarizeHandler := fp.GetSummarizeResultHandler(taskID, subtaskID)
			summarizer := csum.NewSummarizer(csum.SummarizerConfig{
				PreserveLast:   true,
				UseQA:          true,
				SummHumanInQA:  true,
				LastSecBytes:   lastSecBytesAfterRestore,
				MaxBPBytes:     maxBPBytesAfterRestore,
				MaxQASections:  maxQASectionsAfterRestore,
				MaxQABytes:     maxQABytesAfterRestore,
				KeepQASections: keepQASectionsAfterRestore,
			})

			chain, err = summarizer.SummarizeChain(ctx, summarizeHandler, ast.Messages(), fp.tcIDTemplate)
			if err != nil {
				_ = wrapErrorWithEvent("failed to summarize chain", err) // non critical error, just log it
				chain = ast.Messages()
			}

			return nil
		}()
		if err != nil {
			fallback()
		}
	}

	chainObs.End(
		langfuse.WithChainOutput(chain),
		langfuse.WithChainStatus("success"),
	)

	chainBlob, err := json.Marshal(chain)
	if err != nil {
		return 0, nil, fmt.Errorf("failed to marshal msg chain: %w", err)
	}

	msgChain, err = fp.db.CreateMsgChain(ctx, database.CreateMsgChainParams{
		Type:          msgChainType,
		Model:         fp.Model(optAgentType),
		ModelProvider: string(fp.Type()),
		Chain:         chainBlob,
		FlowID:        fp.flowID,
		TaskID:        database.Int64ToNullInt64(taskID),
		SubtaskID:     database.Int64ToNullInt64(subtaskID),
	})
	if err != nil {
		return 0, nil, fmt.Errorf("failed to create msg chain: %w", err)
	}

	return msgChain.ID, chain, nil
}

// Eliminates code duplication by abstracting database operations on message chains
// optAgentType is used to detect model changes within same provider
func (fp *flowProvider) processChain(
	ctx context.Context,
	optAgentType pconfig.ProviderOptionsType,
	msgChainID int64,
	logger *logrus.Entry,
	transform func([]llms.MessageContent) ([]llms.MessageContent, error),
) error {
	msgChain, err := fp.db.GetMsgChain(ctx, msgChainID)
	if err != nil {
		obs.LogErrorOrCancel(logger, err, "failed to get message chain")
		return fmt.Errorf("failed to get message chain %d: %w", msgChainID, err)
	}

	var chain []llms.MessageContent
	if err := json.Unmarshal(msgChain.Chain, &chain); err != nil {
		obs.LogErrorOrCancel(logger, err, "failed to unmarshal message chain")
		return fmt.Errorf("failed to unmarshal message chain %d: %w", msgChainID, err)
	}

	updatedChain, err := transform(chain)
	if err != nil {
		obs.LogErrorOrCancel(logger, err, "failed to transform chain")
		return fmt.Errorf("failed to transform chain: %w", err)
	}

	// Check if provider or model has changed since chain was created
	currentProvider := string(fp.Type())
	currentModel := fp.Model(optAgentType)
	providerChanged := msgChain.ModelProvider != currentProvider
	modelChanged := msgChain.Model != currentModel

	if providerChanged || modelChanged {
		logger.WithFields(logrus.Fields{
			"old_provider":     msgChain.ModelProvider,
			"new_provider":     currentProvider,
			"old_model":        msgChain.Model,
			"new_model":        currentModel,
			"provider_changed": providerChanged,
			"model_changed":    modelChanged,
		}).Info("provider or model changed, normalizing chain")

		// Normalize chain for new provider/model
		ast, err := cast.NewChainAST(updatedChain, true)
		if err != nil {
			logger.WithError(err).Warn("failed to create chain AST for normalization")
		} else {
			// Normalize tool call IDs to new format
			if err := ast.NormalizeToolCallIDs(fp.tcIDTemplate); err != nil {
				logger.WithError(err).Warn("failed to normalize tool call IDs")
			}

			ast.SanitizeToolCallArguments()

			// Clear provider-specific reasoning signatures
			if err := ast.ClearReasoning(); err != nil {
				logger.WithError(err).Warn("failed to clear reasoning")
			}

			// Check if last message is Human - if not, add continuation message
			messages := ast.Messages()
			if len(messages) > 0 && messages[len(messages)-1].Role != llms.ChatMessageTypeHuman {
				ast.AppendHumanMessage(ChainContinuationMessage)
				logger.Debug("added continuation human message after provider switch")
			}

			updatedChain = ast.Messages()
		}
	}

	chainBlob, err := json.Marshal(updatedChain)
	if err != nil {
		obs.LogErrorOrCancel(logger, err, "failed to marshal updated chain")
		return fmt.Errorf("failed to marshal updated chain %d: %w", msgChainID, err)
	}

	_, err = fp.db.UpdateMsgChain(ctx, database.UpdateMsgChainParams{
		Chain:           chainBlob,
		DurationSeconds: 0.0,
		Model:           currentModel,
		ModelProvider:   currentProvider,
		ID:              msgChainID,
	})
	if err != nil {
		obs.LogErrorOrCancel(logger, err, "failed to update message chain")
		return fmt.Errorf("failed to update message chain %d: %w", msgChainID, err)
	}

	return nil
}

func (fp *flowProvider) prepareExecutionContext(ctx context.Context, taskID, subtaskID int64) (string, error) {
	ctx, observation := obs.Observer.NewObservation(ctx)
	evaluator := observation.Evaluator(
		langfuse.WithEvaluatorName("prepare execution context"),
		langfuse.WithEvaluatorInput(map[string]any{
			"task_id":    taskID,
			"subtask_id": subtaskID,
			"flow_id":    fp.flowID,
		}),
	)
	ctx, _ = evaluator.Observation(ctx)

	tasksInfo, err := fp.getTasksInfo(ctx, taskID)
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to get tasks info", err)
	}

	subtasksInfo := fp.getSubtasksInfo(taskID, tasksInfo.Subtasks)
	if subtasksInfo.Subtask == nil {
		subtasks := make([]database.Subtask, 0, len(subtasksInfo.Planned)+len(subtasksInfo.Completed))
		subtasks = append(subtasks, subtasksInfo.Planned...)
		subtasks = append(subtasks, subtasksInfo.Completed...)
		slices.SortFunc(subtasks, func(a, b database.Subtask) int {
			return int(a.ID - b.ID)
		})

		for i, subtask := range subtasks {
			if subtask.ID == subtaskID {
				subtasksInfo.Subtask = &subtask
				subtasksInfo.Planned = subtasks[i+1:]
				subtasksInfo.Completed = subtasks[:i]
				break
			}
		}
	}

	executionContextRaw, err := fp.prompter.RenderTemplate(templates.PromptTypeFullExecutionContext, map[string]any{
		"Task":              tasksInfo.Task,
		"Tasks":             tasksInfo.Tasks,
		"CompletedSubtasks": subtasksInfo.Completed,
		"Subtask":           subtasksInfo.Subtask,
		"PlannedSubtasks":   subtasksInfo.Planned,
	})
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to render execution context", err)
	}

	summarizeHandler := fp.GetSummarizeResultHandler(&taskID, &subtaskID)
	executionContext, err := summarizeHandler(ctx, executionContextRaw)
	if err != nil {
		return "", wrapErrorEndEvaluatorSpan(ctx, evaluator, "failed to summarize execution context", err)
	}

	evaluator.End(
		langfuse.WithEvaluatorOutput(executionContext),
		langfuse.WithEvaluatorStatus("success"),
		langfuse.WithEvaluatorLevel(langfuse.ObservationLevelDebug),
	)

	return executionContext, nil
}

func (fp *flowProvider) getExecutionContext(ctx context.Context, taskID, subtaskID *int64) (string, error) {
	if taskID != nil && subtaskID != nil {
		return fp.getExecutionContextBySubtask(ctx, *taskID, *subtaskID)
	}

	if taskID != nil {
		return fp.getExecutionContextByTask(ctx, *taskID)
	}

	return fp.getExecutionContextByFlow(ctx)
}

func (fp *flowProvider) getExecutionContextBySubtask(ctx context.Context, taskID, subtaskID int64) (string, error) {
	subtask, err := fp.db.GetSubtask(ctx, subtaskID)
	if err == nil && subtask.TaskID == taskID && subtask.Context != "" {
		return subtask.Context, nil
	}

	return fp.getExecutionContextByTask(ctx, taskID)
}

func (fp *flowProvider) getExecutionContextByTask(ctx context.Context, taskID int64) (string, error) {
	tasksInfo, err := fp.getTasksInfo(ctx, taskID)
	if err != nil {
		return fp.getExecutionContextByFlow(ctx)
	}

	subtasksInfo := fp.getSubtasksInfo(taskID, tasksInfo.Subtasks)
	executionContext, err := fp.prompter.RenderTemplate(templates.PromptTypeShortExecutionContext, map[string]any{
		"Task":              tasksInfo.Task,
		"Tasks":             tasksInfo.Tasks,
		"CompletedSubtasks": subtasksInfo.Completed,
		"Subtask":           subtasksInfo.Subtask,
		"PlannedSubtasks":   subtasksInfo.Planned,
	})
	if err != nil {
		return fp.getExecutionContextByFlow(ctx)
	}

	return executionContext, nil
}

func (fp *flowProvider) getExecutionContextByFlow(ctx context.Context) (string, error) {
	tasks, err := fp.db.GetFlowTasks(ctx, fp.flowID)
	if err != nil {
		return "", fmt.Errorf("failed to get flow tasks: %w", err)
	}

	if len(tasks) == 0 {
		return "flow has no tasks, it's using in assistant mode", nil
	}

	subtasks, err := fp.db.GetFlowSubtasks(ctx, fp.flowID)
	if err != nil {
		return "", fmt.Errorf("failed to get flow subtasks: %w", err)
	}

	for tid := len(tasks) - 1; tid >= 0; tid-- {
		taskID := tasks[tid].ID

		subtasksInfo := fp.getSubtasksInfo(taskID, subtasks)
		executionContext, err := fp.prompter.RenderTemplate(templates.PromptTypeShortExecutionContext, map[string]any{
			"Task":              tasks[tid],
			"Tasks":             tasks,
			"CompletedSubtasks": subtasksInfo.Completed,
			"Subtask":           subtasksInfo.Subtask,
			"PlannedSubtasks":   subtasksInfo.Planned,
		})
		if err != nil {
			continue
		}

		return executionContext, nil
	}

	subtasksInfo := fp.getSubtasksInfo(0, subtasks)
	executionContext, err := fp.prompter.RenderTemplate(templates.PromptTypeShortExecutionContext, map[string]any{
		"Tasks":             tasks,
		"CompletedSubtasks": subtasksInfo.Completed,
		"Subtask":           subtasksInfo.Subtask,
		"PlannedSubtasks":   subtasksInfo.Planned,
	})
	if err != nil {
		return "", fmt.Errorf("failed to render execution context: %w", err)
	}

	return executionContext, nil
}

func (fp *flowProvider) subtasksToMarkdown(subtasks []tools.SubtaskInfo) string {
	var buffer strings.Builder
	for sid, subtask := range subtasks {
		fmt.Fprintf(&buffer, "# Subtask %d\n\n", sid+1)
		fmt.Fprintf(&buffer, "## %s\n\n%s\n\n", subtask.Title, subtask.Description)
	}

	return buffer.String()
}

func (fp *flowProvider) getContainerPortsDescription() string {
	ports := docker.GetPrimaryContainerPorts(fp.cfg.WorkerPortsBase(), fp.flowID)
	var buffer strings.Builder

	buffer.WriteString("**OOB Attack Infrastructure:**\n\n")

	// Host network mode: container has direct access to host network interfaces
	if fp.cfg.WorkerNetwork() == "host" {
		buffer.WriteString("This container uses **host network mode** with direct access to host network interfaces.\n\n")
		buffer.WriteString("**MANDATORY PORTS - YOU MUST USE ONLY THESE:**\n\n")

		for _, port := range ports {
			fmt.Fprintf(&buffer, "- Port %d/tcp (REQUIRED)\n", port)
		}

		buffer.WriteString("\n**Network Access:**\n")
		buffer.WriteString("- Direct access to all host network interfaces\n")
		buffer.WriteString("- No port forwarding - services bind directly to host ports\n")
		buffer.WriteString("- **CRITICAL**: DO NOT use any other ports - this may interfere with other running applications on the host system\n")
		buffer.WriteString("- Using non-allocated ports could cause conflicts with system services and impact host stability\n")
		buffer.WriteString("- All host network interfaces are accessible for binding\n\n")

		buffer.WriteString("**Usage for OOB Attacks:**\n")
		if fp.cfg.WorkerPublicIP() == "0.0.0.0" {
			buffer.WriteString("To determine the public IP for callbacks:\n")
			buffer.WriteString("1. Discover your public IP: `curl -s https://api.ipify.org` or `curl -s ipinfo.io/ip`\n")
			buffer.WriteString("2. Use discovered IP in exploit payloads for callbacks\n")
			buffer.WriteString("3. Listen on allocated ports (shown above) to receive connections\n\n")
			buffer.WriteString("**Important:** Check Task.Input - user may have specified the public IP to use.\n")
		} else {
			fmt.Fprintf(&buffer, "Your external IP is: **%s**\n\n", fp.cfg.WorkerPublicIP())
			buffer.WriteString("Use this IP in exploit payloads requiring callbacks (DNS exfiltration, reverse shells, XXE OOB, SSRF verification, etc.)\n")
			buffer.WriteString("Listen on the allocated ports above to receive incoming connections.\n")
		}
	} else {
		// Bridge network mode: traditional port forwarding
		buffer.WriteString("**MANDATORY FORWARDED PORTS - YOU MUST USE ONLY THESE:**\n\n")

		for _, port := range ports {
			fmt.Fprintf(&buffer, "- Port %d/tcp (container) → %s:%d (external)\n", port, fp.cfg.WorkerPublicIP(), port)
		}

		buffer.WriteString("\n**Port Usage Rules:**\n")
		buffer.WriteString("- **YOU MUST use ONLY the ports listed above** for all listeners and reverse connections\n")
		buffer.WriteString("- **CRITICAL**: Any reverse connections (shells, callbacks) will FAIL on other ports - only allocated ports are forwarded\n")
		buffer.WriteString("- Standard ports like 4444, 8080, 9001 are NOT forwarded and will NOT work\n")
		fmt.Fprintf(&buffer, "- Example: For Metasploit reverse shell, use LPORT=%d (not 4444)\n\n", ports[0])

		buffer.WriteString("**Usage for OOB Attacks:**\n")

		if fp.cfg.WorkerPublicIP() == "0.0.0.0" {
			buffer.WriteString("The bind IP is 0.0.0.0 (all interfaces). To receive external callbacks:\n")
			buffer.WriteString("1. Discover your public IP: `curl -s https://api.ipify.org` or `curl -s ipinfo.io/ip`\n")
			buffer.WriteString("2. Use discovered IP in exploit payloads for callbacks\n")
			buffer.WriteString("3. Listen on container ports (shown above) to receive connections\n\n")
			buffer.WriteString("**Important:** Check Task.Input - user may have specified the public IP to use.\n")
		} else {
			fmt.Fprintf(&buffer, "Your external IP is: %s\n", fp.cfg.WorkerPublicIP())
			buffer.WriteString("Use this IP in exploit payloads requiring callbacks (DNS exfiltration, reverse shells, XXE OOB, SSRF verification, etc.)\n")
			buffer.WriteString("Listen on the container ports above to receive incoming connections.\n")
		}
	}

	return buffer.String()
}

func getCurrentTime() string {
	return time.Now().Format("2006-01-02 15:04:05")
}

func isEmptyChain(msgChain json.RawMessage) bool {
	var msgList []llms.MessageContent

	if err := json.Unmarshal(msgChain, &msgList); err != nil {
		return true
	}

	return len(msgList) == 0
}

func getToolCallMessage(toolCall *llms.FunctionCall) map[string]string {
	var msg dummyMessage

	if toolCall == nil {
		return nil
	}

	if err := json.Unmarshal(json.RawMessage(toolCall.Arguments), &msg); err != nil {
		return nil
	}

	return map[string]string{
		"name": toolCall.Name,
		"msg":  msg.Message,
	}
}

// getRecentMessages returns the last section messages from the chain
func getRecentMessages(chain []llms.MessageContent) []map[string]string {
	var messages []map[string]string

	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return messages
	}

	lastSection := ast.Sections[len(ast.Sections)-1]
	if len(lastSection.Body) == 0 {
		return messages
	}

	for idx := len(lastSection.Body) - 1; idx >= 0 && len(messages) < extractLastMessagesCount; idx-- {
		pair := lastSection.Body[idx]
		if pair.Type != cast.RequestResponse {
			continue
		}

		for _, tc := range pair.GetToolCallsInfo().CompletedToolCalls {
			message := getToolCallMessage(tc.ToolCall.FunctionCall)
			if message != nil {
				messages = append(messages, message)
			}
		}
	}

	slices.Reverse(messages)

	return messages
}

func cutString(s string, maxLength int) string {
	if len(s) <= maxLength {
		return s
	}

	return s[:maxLength] + "...[truncated full size is " + strconv.Itoa(len(s)) + " bytes]"
}

func formatToolCallArguments(args string) string {
	var v map[string]any
	if err := json.Unmarshal(json.RawMessage(args), &v); err != nil {
		return ""
	}

	delete(v, "message")
	var keys []string
	for k := range v {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var buffer strings.Builder
	for _, k := range keys {
		value, err := json.Marshal(v[k])
		if err != nil {
			continue
		}
		fmt.Fprintf(&buffer, "<field name=\"%s\">%s</field>\n", k, cutString(string(value), 256))
	}

	return buffer.String()
}

func getToolCallInfo(toolCall *cast.ToolCallPair) map[string]string {
	if toolCall == nil {
		return nil
	}

	if toolCall.ToolCall.FunctionCall == nil {
		return nil
	}

	return map[string]string{
		"name":   toolCall.ToolCall.FunctionCall.Name,
		"args":   formatToolCallArguments(toolCall.ToolCall.FunctionCall.Arguments),
		"result": cutString(toolCall.Response.Content, 1024),
	}
}

// extractToolCallsFromChain extracts all tool calls from the message chain
func extractToolCallsFromChain(chain []llms.MessageContent) []map[string]string {
	var toolCalls []map[string]string

	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return toolCalls
	}

	lastSection := ast.Sections[len(ast.Sections)-1]
	if len(lastSection.Body) == 0 {
		return toolCalls
	}

	for idx := len(lastSection.Body) - 1; idx >= 0 && len(toolCalls) < extractToolCallsCount; idx-- {
		pair := lastSection.Body[idx]
		if pair.Type != cast.RequestResponse {
			continue
		}

		toolCallsInfo := pair.GetToolCallsInfo()
		for _, tc := range toolCallsInfo.CompletedToolCalls {
			info := getToolCallInfo(tc)
			if info != nil {
				toolCalls = append(toolCalls, info)
			}
		}
	}

	slices.Reverse(toolCalls)

	return toolCalls
}

// extractAgentPromptFromChain extracts the agent prompt from the message chain
func extractAgentPromptFromChain(chain []llms.MessageContent) string {
	ast, err := cast.NewChainAST(chain, true)
	if err != nil {
		return ""
	}

	lastSection := ast.Sections[len(ast.Sections)-1]
	if len(lastSection.Body) == 0 {
		return ""
	}

	if lastSection.Header == nil {
		return ""
	}

	humanMessage := lastSection.Header.HumanMessage
	if humanMessage == nil {
		return ""
	}

	var parts []string
	for _, part := range humanMessage.Parts {
		if text, ok := part.(llms.TextContent); ok && text.Text != "" {
			parts = append(parts, text.Text)
		}
	}

	return strings.Join(parts, "\n")
}

// formatEnhancedToolResponse formats tool response with optional mentor analysis
func formatEnhancedToolResponse(originalResult, mentorAnalysis string) string {
	if mentorAnalysis == "" {
		return originalResult
	}

	return fmt.Sprintf(`<enhanced_response>
<original_result>
%s
</original_result>

<mentor_analysis>
%s
</mentor_analysis>
</enhanced_response>`, originalResult, mentorAnalysis)
}

func extractHistoryFromHumanMessage(msg *llms.MessageContent) string {
	if msg == nil {
		return ""
	}

	if msg.Role != llms.ChatMessageTypeHuman {
		return ""
	}

	var parts []string
	for _, part := range msg.Parts {
		if text, ok := part.(llms.TextContent); ok && text.Text != "" {
			parts = append(parts, text.Text)
		}
	}

	msgText := strings.Join(parts, "\n")
	msgParts := strings.Split(msgText, toolCallsHistorySeparator)
	if len(msgParts) < 2 {
		return ""
	}

	return strings.Trim(msgParts[len(msgParts)-1], "\n\t\r ")
}

func appendNewToolCallsToHistory(history string, toolCalls []map[string]string) string {
	var buffer strings.Builder

	buffer.WriteString(history)
	if history != "" {
		buffer.WriteString("\n")
	}

	for _, toolCall := range toolCalls {
		fmt.Fprintf(&buffer, "<tool_call>\n<name>%s</name>\n<arguments>\n%s\n</arguments>\n<result>%s</result>\n</tool_call>\n",
			toolCall["name"], toolCall["args"], toolCall["result"])
	}

	return buffer.String()
}

func combineHistoryToolCallsToHumanMessage(history, msg string) string {
	return fmt.Sprintf("%s\n\n%s\n\n%s", msg, toolCallsHistorySeparator, history)
}

func enrichLogrusFields(flowID int64, taskID, subtaskID *int64, fields logrus.Fields) logrus.Fields {
	if fields == nil {
		fields = logrus.Fields{}
	}

	fields["flow_id"] = flowID
	if taskID != nil {
		fields["task_id"] = *taskID
	}
	if subtaskID != nil {
		fields["subtask_id"] = *subtaskID
	}

	return fields
}