package promptcompiler

import (
	"context"
	"embed"
	"encoding/json"
	"os"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// promptMatrixGolden is intentionally checked in: changing semantic text,
// ordering, template provenance, or carrier shape requires an explicit golden
// review rather than a non-empty marker assertion.
//
//go:embed testdata/prompt_matrix.json
var promptMatrixGolden embed.FS

type promptMatrixArtifact struct {
	Cases []promptMatrixCase `json:"cases"`
}

type promptMatrixCase struct {
	Name      string               `json:"name"`
	Runtime   string               `json:"runtime"`
	Event     string               `json:"event"`
	Adapter   string               `json:"adapterBoundary"`
	Surfaces  promptMatrixSurfaces `json:"surfaces,omitempty"`
	Blocks    []promptMatrixBlock  `json:"blocks"`
	Generated bool                 `json:"generated,omitempty"`
}

type promptMatrixSurfaces struct {
	Prompt         string `json:"prompt,omitempty"`
	FreshPrompt    string `json:"freshPrompt,omitempty"`
	ResumePrompt   string `json:"resumePrompt,omitempty"`
	Base           string `json:"baseInstructions,omitempty"`
	Developer      string `json:"developerInstructions,omitempty"`
	System         string `json:"systemPrompt,omitempty"`
	RuntimeCarrier string `json:"runtimeCarrier,omitempty"`
	TurnInput      string `json:"turnInput,omitempty"`
}

type promptMatrixBlock struct {
	ID           string `json:"id"`
	Slot         string `json:"slot"`
	Priority     int    `json:"priority"`
	Durability   string `json:"durability"`
	Template     string `json:"template"`
	TemplateVer  string `json:"templateVersion"`
	RenderedText string `json:"renderedText"`
}

func TestPromptMatrixMatchesCheckedInGoldenArtifact(t *testing.T) {
	actual := buildPromptMatrixArtifact(t)
	actualJSON, err := json.MarshalIndent(actual, "", "  ")
	require.NoError(t, err)
	golden, err := promptMatrixGolden.ReadFile("testdata/prompt_matrix.json")
	require.NoError(t, err)
	if os.Getenv("UPDATE_PROMPT_GOLDENS") == "1" {
		require.NoError(t, os.WriteFile(filepath.Join("testdata", "prompt_matrix.json"), append(actualJSON, '\n'), 0o644))
		golden = append(actualJSON, '\n')
	}
	require.JSONEq(t, string(golden), string(actualJSON))
}

func buildPromptMatrixArtifact(t *testing.T) promptMatrixArtifact {
	t.Helper()
	var artifact promptMatrixArtifact
	runtimes := []string{"codex", "claude", "opencode", "kimi", "pi"}
	for _, runtime := range runtimes {
		for _, event := range []CompileEventType{EventFreshStart, EventWake, EventResume, EventDeferredContinuation} {
			input := promptMatrixInput(runtime, "message")
			carrier, err := New(DefaultPolicy()).CompileLaunchPrompt(context.Background(), LaunchCompileDescriptor{
				EventID: "golden-event", TurnEventType: event, SessionID: launchMatrixSessionID(event), Runtime: runtime,
				ContractMode: LaunchContractModeForRuntime(runtime), ContractCarrierFile: LaunchContractFileForRuntime(runtime),
				FreshFallbackOnResume: LaunchNeedsFreshFallbackOnResume(runtime),
			}, LaunchCompileInput{Input: input})
			require.NoError(t, err)
			artifact.Cases = append(artifact.Cases, matrixCaseFromCarrier("lifecycle-"+runtime+"-"+string(event), runtime, event, carrier))
		}
	}
	for _, runtime := range runtimes {
		for _, tc := range []struct {
			name  string
			event CompileEventType
			input CompileInput
		}{
			{"normal-message", EventFreshStart, promptMatrixInput("codex", "message")},
			{"task-dispatch", EventFreshStart, promptMatrixInput("codex", "task")},
			{"team-context", EventWake, promptMatrixInput("codex", "team")},
			{"workflow-callback", EventWake, promptMatrixInput("codex", "workflow")},
			{"workflow-edit", EventWake, promptMatrixInput("codex", "workflow-edit")},
			{"attachment-success", EventFreshStart, promptMatrixInput("codex", "attachment-success")},
			{"attachment-failure", EventFreshStart, promptMatrixInput("codex", "attachment-failure")},
			{"exec-guidance", EventFreshStart, promptMatrixInput("codex", "message")},
			{"extension-guidance", EventFreshStart, promptMatrixInput("codex", "message")},
			{"unread-summary", EventWake, promptMatrixInput("codex", "message")},
		} {
			input := tc.input
			input.Runtime = runtime
			switch tc.name {
			case "exec-guidance":
				input.Exec = ExecCompileInput{Enabled: true, Guidance: ExecGuidanceTemplateData{Domains: []string{"agent", "task"}, CommonActions: []ExecGuidanceAction{{Domain: "task", Name: "get", ExampleJSON: `{"number":1}`}}}}
			case "extension-guidance":
				input.Extension = ExtensionCompileInput{Enabled: true, Guidance: ExtensionCatalogGuidanceTemplateData{Actions: []ExtensionCatalogGuidanceAction{{ExtensionName: "github", ActionName: "list-issues"}}}}
			case "unread-summary":
				input.Message = types.InboxMessage{}
				input.UnreadSummary = map[string]int{"#zeta": 2, "#alpha": 1}
			}
			carrier, err := New(DefaultPolicy()).CompileLaunchPrompt(context.Background(), LaunchCompileDescriptor{
				EventID: "golden-domain-event", TurnEventType: tc.event, Runtime: runtime,
				ContractMode: LaunchContractModeForRuntime(runtime), ContractCarrierFile: LaunchContractFileForRuntime(runtime),
				FreshFallbackOnResume: LaunchNeedsFreshFallbackOnResume(runtime),
			}, LaunchCompileInput{Input: input})
			require.NoError(t, err)
			artifact.Cases = append(artifact.Cases, matrixCaseFromCarrier(tc.name+"-"+runtime, runtime, tc.event, carrier))
		}
	}
	for _, tc := range []struct {
		name  string
		kind  GeneratedPromptKind
		event CompileEventType
	}{
		{"recovery-restart", GeneratedRecoveryRestart, EventDeferredContinuation},
		{"final-visible-reply", GeneratedFinalReplyCheck, EventFinalVisibleReplyCheck},
		{"context-compaction", GeneratedContextCompaction, EventDeferredContinuation},
		{"output-token-truncation", GeneratedOutputTruncation, EventDeferredContinuation},
	} {
		for _, runtime := range runtimes {
			carrier, err := New(DefaultPolicy()).CompileGeneratedPrompt(context.Background(), tc.event, "generated-golden", GeneratedPromptInput{
				Kind:       tc.kind,
				FinalReply: FinalVisibleReplyCheckTemplateData{Wakes: []FinalVisibleReplyCheckTemplateWake{{Target: "#matrix", HasTarget: true}}, SendMessageTool: "mcp__agentteam__message_send"},
			}, runtime)
			require.NoError(t, err)
			artifact.Cases = append(artifact.Cases, matrixCaseFromGenerated(tc.name+"-"+runtime, runtime, tc.event, carrier))
		}
	}
	return artifact
}

func promptMatrixInput(runtime, kind string) CompileInput {
	input := CompileInput{AgentID: "agt-matrix", AgentName: "matrix", Runtime: runtime, ExternalSkillCorePayload: "core skill", ExternalSkillReferencePayload: "reference skill", Message: types.InboxMessage{MessageID: "msg-matrix", Target: "#matrix", Content: "continue the matrix", AuthorProfile: types.AuthorProfile{ActorType: "human", DisplayName: "alice"}}}
	switch kind {
	case "task":
		input.Message.MessageType = "task_dispatch"
		input.Message.Metadata = map[string]any{"task": map[string]any{"taskNumber": 42, "title": "ship the feature", "channelRef": "#matrix", "status": "todo", "acceptanceCriteria": []any{"use the task scope"}, "taskScope": "#tsk_matrix"}, "taskDiscussion": map[string]any{"body": "Please do the task.", "createdBy": map[string]any{"type": "human", "name": "alice"}}}
	case "team":
		input.Message.TeamContexts = []types.TeamContext{{TeamID: "team-1", TeamHandle: "matrix", Role: "leader", RoleLabel: "Lead", TeamDescription: "Matrix team", TeamMemberID: "member-1", MemberRoster: []types.TeamRosterMember{{TeamMemberID: "member-1", Handle: "matrix", DisplayName: "Matrix", Role: "lead", Self: true}}}}
	case "workflow":
		input.Message.WorkflowContext = types.WorkflowContext{RunID: "run-matrix", EffectID: "effect-matrix", CallbackExpected: true}
	case "workflow-edit":
		input.Message.Metadata = map[string]any{"workflowEdit": map[string]any{"kind": "request", "planId": "plan-matrix", "responderAgentId": "agt-matrix", "baseDraftRevision": 3, "target": map[string]any{"kind": "node", "id": "notify", "nodeType": "send_message"}, "diagnostics": []any{map[string]any{"severity": "warning", "code": "W1", "path": "nodes[0]", "message": "check target"}}}}
	case "attachment-success":
		input.Attachments = AttachmentResolution{Prepared: []types.RuntimeAttachment{{ID: "att-matrix", Filename: "matrix.txt", URI: "file:///tmp/matrix.txt", MIME: "text/plain"}}}
	case "attachment-failure":
		input.Attachments = AttachmentResolution{Failure: &AttachmentPreparationFailure{Attachments: []types.MessageAttachment{{ID: "att-matrix", Filename: "matrix.txt"}}, Error: "offline"}}
	}
	return input
}

func matrixCaseFromCarrier(name, runtime string, event CompileEventType, carrier CompiledPrompt) promptMatrixCase {
	return promptMatrixCase{Name: name, Runtime: runtime, Event: string(event), Adapter: string(LaunchContractModeForRuntime(runtime)), Surfaces: promptMatrixSurfaces{Prompt: carrier.LaunchSurface("prompt"), FreshPrompt: carrier.LaunchSurface("freshPrompt"), ResumePrompt: carrier.LaunchSurface("resumePrompt"), Base: carrier.LaunchSurface("baseInstructions"), Developer: carrier.LaunchSurface("developerInstructions"), System: carrier.LaunchSurface("systemPrompt"), RuntimeCarrier: carrier.LaunchSurface("runtimeCarrier"), TurnInput: carrier.SlotText("turnInput")}, Blocks: matrixBlocks(carrier)}
}

func matrixCaseFromGenerated(name, runtime string, event CompileEventType, carrier CompiledPrompt) promptMatrixCase {
	return promptMatrixCase{Name: name, Runtime: runtime, Event: string(event), Adapter: "generated-runtime-input", Generated: true, Surfaces: promptMatrixSurfaces{TurnInput: carrier.SlotText("turnInput")}, Blocks: matrixBlocks(carrier)}
}

func matrixBlocks(carrier CompiledPrompt) []promptMatrixBlock {
	var blocks []promptMatrixBlock
	components := carrier.LaunchComponents()
	if len(components) == 0 && carrier.Envelope().Bundle != nil {
		components = []CompiledPromptComponent{{Envelope: carrier.Envelope()}}
	}
	for _, component := range components {
		if component.Envelope.Bundle == nil {
			continue
		}
		for _, block := range component.Envelope.Bundle.Blocks {
			blocks = append(blocks, promptMatrixBlock{ID: block.ID, Slot: string(block.Slot), Priority: block.Priority, Durability: string(block.Durability), Template: block.Source.Path, TemplateVer: block.Source.Version, RenderedText: block.RenderedText()})
		}
	}
	return blocks
}