package dispatcher

import (
	"context"
	"strings"

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

type compiledSessionPrompts struct {
	Carrier        *promptcompiler.CompiledPrompt
	PromptInputs   []types.TrajectoryEntry
	CompileMetrics []promptCompileMeasurement
}

type promptCompileMeasurement struct {
	compiled  promptcompiler.RuntimeInputEnvelope
	rendered  promptcompiler.RenderedPromptPayload
	stage     string
	freshness string
}

func compiledPromptsUseWorkspaceRootMemory(prompts compiledSessionPrompts) bool {
	for _, measurement := range prompts.CompileMetrics {
		for _, block := range measurement.compiled.Bundle.Blocks {
			if block.ID == "platform.workspace_memory" {
				for _, tag := range block.Tags {
					if tag == "memory:remote" {
						return false
					}
				}
				return true
			}
		}
	}
	return true
}

func appendPromptCompileMeasurement(prompts *compiledSessionPrompts, compiled promptcompiler.RuntimeInputEnvelope, rendered promptcompiler.RenderedPromptPayload, stage, freshness string) {
	prompts.CompileMetrics = append(prompts.CompileMetrics, promptCompileMeasurement{compiled: compiled, rendered: rendered, stage: stage, freshness: freshness})
}

func compileSessionInstructionsFingerprint(
	workDir string,
	ag agent.Config,
	rt runtimeconfig.RuntimeType,
	binding conversationBinding,
	msg types.InboxMessage,
	instructions resolvedInstructionContext,
	mcps []types.NamedMCPConfig,
	resuming bool,
) (string, error) {
	compiler := promptcompiler.New(promptcompiler.DefaultPolicy())
	input := buildPromptCompileInput(workDir, ag, rt, binding, msg, instructions)
	eventType := promptcompiler.EventFreshStart
	if resuming {
		eventType = promptcompiler.EventResume
	}
	carrier, err := compiler.CompileLaunchPrompt(context.Background(), promptcompiler.LaunchCompileDescriptor{
		EventID:               promptcompiler.NewEventID(eventType),
		TurnEventType:         eventType,
		SessionID:             firstNonEmpty(binding.serverSessionID, ""),
		Runtime:               string(rt),
		ContractMode:          promptcompiler.LaunchContractModeForRuntime(string(rt)),
		ContractCarrierFile:   promptcompiler.LaunchContractFileForRuntime(string(rt)),
		FreshFallbackOnResume: promptcompiler.LaunchNeedsFreshFallbackOnResume(string(rt)),
	}, promptcompiler.LaunchCompileInput{Input: input})
	if err != nil {
		return "", err
	}
	fingerprint, err := promptcompiler.ThreadContractInstructionsFingerprint(carrier.LaunchComponents()[0].Envelope)
	if err != nil {
		return "", err
	}
	if rt == runtimeconfig.TypeCodex {
		fingerprint = codexAppServerSessionFingerprint(fingerprint, mcps)
	}
	return fingerprint, nil
}

func resolveCompiledSessionID(
	workDir string,
	ag agent.Config,
	rt runtimeconfig.RuntimeType,
	binding conversationBinding,
	msg types.InboxMessage,
	instructions resolvedInstructionContext,
	mcps []types.NamedMCPConfig,
) (string, error) {
	sessionID := strings.TrimSpace(binding.serverSessionID)
	if rt != runtimeconfig.TypeCodex || sessionID == "" {
		return sessionID, nil
	}
	compiler := promptcompiler.New(promptcompiler.DefaultPolicy())
	input := buildPromptCompileInput(workDir, ag, rt, binding, msg, instructions)
	carrier, err := compiler.CompileLaunchPrompt(context.Background(), promptcompiler.LaunchCompileDescriptor{
		EventID:               promptcompiler.NewEventID(promptcompiler.EventResume),
		TurnEventType:         promptcompiler.EventResume,
		SessionID:             sessionID,
		Runtime:               string(rt),
		ContractMode:          promptcompiler.LaunchContractModeForRuntime(string(rt)),
		ContractCarrierFile:   promptcompiler.LaunchContractFileForRuntime(string(rt)),
		FreshFallbackOnResume: promptcompiler.LaunchNeedsFreshFallbackOnResume(string(rt)),
	}, promptcompiler.LaunchCompileInput{Input: input})
	if err != nil {
		return "", err
	}
	startCfg := drivercontract.LaunchConfig{SessionID: sessionID, CompiledPrompt: &carrier}
	currentFingerprint := codexAppServerSessionFingerprint(
		instructionsFingerprint(carrier.LaunchSurface("baseInstructions"), carrier.LaunchSurface("developerInstructions")),
		mcps,
	)
	applyCodexAppServerResumePolicy(workDir, &startCfg, currentFingerprint)
	return strings.TrimSpace(startCfg.SessionID), nil
}

func compileSessionPrompts(
	workDir string,
	ag agent.Config,
	rt runtimeconfig.RuntimeType,
	binding conversationBinding,
	msg types.InboxMessage,
	instructions resolvedInstructionContext,
) (compiledSessionPrompts, error) {
	compiler := promptcompiler.New(promptcompiler.DefaultPolicy())
	input := buildPromptCompileInput(workDir, ag, rt, binding, msg, instructions)
	sessionID := strings.TrimSpace(binding.serverSessionID)
	turnType := sessionTurnContextEventType(sessionID != "", msg)
	carrier, err := compiler.CompileLaunchPrompt(context.Background(), promptcompiler.LaunchCompileDescriptor{
		EventID:               promptcompiler.NewEventID(turnType),
		TurnEventType:         turnType,
		SessionID:             sessionID,
		Runtime:               string(rt),
		ContractMode:          promptcompiler.LaunchContractModeForRuntime(string(rt)),
		ContractCarrierFile:   promptcompiler.LaunchContractFileForRuntime(string(rt)),
		FreshFallbackOnResume: promptcompiler.LaunchNeedsFreshFallbackOnResume(string(rt)),
	}, promptcompiler.LaunchCompileInput{Input: input})
	if err != nil {
		return compiledSessionPrompts{}, err
	}
	components := carrier.LaunchComponents()
	prompts := compiledSessionPrompts{Carrier: &carrier, PromptInputs: carrier.TrajectoryInputs(string(rt), turnType)}
	for _, component := range components {
		appendPromptCompileMeasurement(&prompts, component.Envelope, component.Rendered, component.Stage, component.Freshness)
	}
	return prompts, nil
}

func compileDeferredContinuationPrompts(
	workDir string,
	ag agent.Config,
	rt runtimeconfig.RuntimeType,
	binding conversationBinding,
	msg types.InboxMessage,
	instructions resolvedInstructionContext,
	pendingBehind int,
) (compiledSessionPrompts, error) {
	compiler := promptcompiler.New(promptcompiler.DefaultPolicy())
	input := buildPromptCompileInput(workDir, ag, rt, binding, msg, instructions)
	input.Delivery.PendingBehind = pendingBehind
	sessionID := strings.TrimSpace(binding.serverSessionID)
	carrier, err := compiler.CompileLaunchPrompt(context.Background(), promptcompiler.LaunchCompileDescriptor{
		EventID:               promptcompiler.NewEventID(promptcompiler.EventDeferredContinuation),
		TurnEventType:         promptcompiler.EventDeferredContinuation,
		SessionID:             sessionID,
		Runtime:               string(rt),
		ContractMode:          promptcompiler.LaunchContractModeForRuntime(string(rt)),
		ContractCarrierFile:   promptcompiler.LaunchContractFileForRuntime(string(rt)),
		FreshFallbackOnResume: promptcompiler.LaunchNeedsFreshFallbackOnResume(string(rt)),
	}, promptcompiler.LaunchCompileInput{Input: input})
	if err != nil {
		return compiledSessionPrompts{}, err
	}
	prompts := compiledSessionPrompts{Carrier: &carrier, PromptInputs: carrier.TrajectoryInputs(string(rt), promptcompiler.EventDeferredContinuation)}
	for _, component := range carrier.LaunchComponents() {
		appendPromptCompileMeasurement(&prompts, component.Envelope, component.Rendered, component.Stage, component.Freshness)
	}
	return prompts, nil
}

func buildPromptCompileInput(
	workDir string,
	ag agent.Config,
	rt runtimeconfig.RuntimeType,
	binding conversationBinding,
	msg types.InboxMessage,
	instructions resolvedInstructionContext,
) promptcompiler.CompileInput {
	agentName := strings.TrimSpace(ag.Name)
	if agentName == "" {
		agentName = strings.TrimSpace(ag.ID)
	}
	activeTaskTarget := binding.activeTaskTarget.Normalize()
	if !activeTaskTarget.Valid() {
		activeTaskTarget = msg.ResolvedActiveTaskTarget()
	}
	return promptcompiler.CompileInput{
		AgentID:          strings.TrimSpace(ag.ID),
		AgentName:        agentName,
		AgentDescription: strings.TrimSpace(ag.Description),
		Runtime:          string(rt),
		WorkDir:          workDir,
		Binding: promptcompiler.ConversationBinding{
			ChannelID:        binding.channelID,
			ChannelType:      binding.targetScopeType,
			TargetID:         binding.targetID,
			ActiveTaskTarget: activeTaskTarget,
			WorkflowContext:  msg.ResolvedWorkflowContext(),
		},
		Message:                       msg,
		UnreadSummary:                 binding.unreadSummary,
		ExternalSkillCorePayload:      strings.TrimSpace(instructions.skillDocParts.Core),
		ExternalSkillReferencePayload: strings.TrimSpace(instructions.skillDocParts.Reference),
		SkillPlan:                     instructions.skillPlan,
		Exec:                          instructions.exec,
		Extension:                     instructions.extension,
		MemoryResumeReminder:          instructions.memoryResumeReminder,
		Attachments:                   instructions.attachments,
	}
}

func sessionTurnContextEventType(resuming bool, msg types.InboxMessage) promptcompiler.CompileEventType {
	if isWakeMessage(msg) {
		return promptcompiler.EventWake
	}
	if resuming {
		return promptcompiler.EventResume
	}
	return promptcompiler.EventFreshStart
}

func isWakeMessage(msg types.InboxMessage) bool {
	return strings.TrimSpace(msg.MessageID) != "" ||
		strings.TrimSpace(msg.Target) != "" ||
		strings.TrimSpace(msg.Content) != "" ||
		strings.TrimSpace(msg.MessageType) != "" ||
		len(msg.Attachments()) > 0 ||
		len(msg.Metadata) > 0
}

func promptSessionKindForBinding(binding conversationBinding) string {
	_ = binding
	return "interactive"
}

func isSensitivePromptKind(kind string) bool {
	switch strings.TrimSpace(kind) {
	case "base_instructions", "developer_instructions", "runtime_native_contract", "resume":
		return true
	default:
		return false
	}
}