package promptcompiler

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

const updateGoldenEnv = "UPDATE_GOLDEN"

func TestPromptCompilerGoldens_ThreadContractArtifacts(t *testing.T) {
	workDir := t.TempDir()
	require.NoError(t, os.MkdirAll(filepath.Join(workDir, "memory", "system"), 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(workDir, "memory", "system", "identity.md"), []byte("# Role\nPersistent reviewer\n"), 0o644))
	require.NoError(t, os.WriteFile(filepath.Join(workDir, "AGENTS.md"), []byte("Always validate migrations before shipping."), 0o644))

	toolName := func(name string) string { return "mcp__agentteam__" + name }
	compiler := New(DefaultPolicy())
	env, err := compiler.CompileThreadContract(context.Background(), EventThreadStart, "evt-golden-thread-contract", CompileInput{
		AgentID:                  "agt-1",
		AgentName:                "reviewer",
		AgentDescription:         "Review backend changes carefully",
		Runtime:                  "codex",
		WorkDir:                  workDir,
		ExternalSkillCorePayload: AgentTeamSkillDocParts(toolName).Core,
	})
	require.NoError(t, err)

	rendered, err := compiler.RenderRuntimeInput(context.Background(), env, "codex", RenderOptions{})
	require.NoError(t, err)

	assertGoldenJSON(t, "thread_contract.runtime_input.json", env)
	assertGoldenJSON(t, "thread_contract.prompt_bundle.json", env.Bundle)
	assertGoldenJSON(t, "thread_contract.rendered_payload.json", rendered)
	assertGoldenText(t, "thread_contract.rendered_turn.md", renderedPayloadMarkdown(rendered))
}

func TestPromptCompilerGoldens_TurnContextArtifacts(t *testing.T) {
	workDir := t.TempDir()
	require.NoError(t, identitynotice.Save(workDir, identitynotice.Pending{
		CurrentName: "reviewer",
		OldName:     "planner",
	}))

	compiler := New(DefaultPolicy())
	env, err := compiler.CompileTurnContext(context.Background(), EventDeferredContinuation, "evt-golden-turn-context", "turn-golden-1", CompileInput{
		AgentID: "agt-1",
		Runtime: "codex",
		WorkDir: workDir,
		Binding: ConversationBinding{
			ChannelID:   "agentteam",
			ChannelType: "channel",
			TargetID:    "#ops",
		},
		Message: types.InboxMessage{
			MessageID:     "msg-1",
			AuthorProfile: types.AuthorProfile{ActorType: "human", DisplayName: "alice"},
			Content:       "Continue the deployment audit.",
		},
		ExternalSkillReferencePayload: "## Reference\n\nRead this when needed.",
		MemoryResumeReminder:          "Re-read MEMORY.md before replying.",
		Delivery: DeliveryCompileInput{
			PendingBehind: 3,
		},
	})
	require.NoError(t, err)

	rendered, err := compiler.RenderRuntimeInput(context.Background(), env, "codex", RenderOptions{})
	require.NoError(t, err)

	assertGoldenJSON(t, "turn_context.runtime_input.json", env)
	assertGoldenJSON(t, "turn_context.prompt_bundle.json", env.Bundle)
	assertGoldenJSON(t, "turn_context.rendered_payload.json", rendered)
	assertGoldenText(t, "turn_context.rendered_turn.md", renderedPayloadMarkdown(rendered))
}

func TestPromptCompilerGoldens_WakeAttachmentOnlyTurnContext(t *testing.T) {
	compiler := New(DefaultPolicy())
	env, err := compiler.CompileTurnContext(context.Background(), EventWake, "evt-golden-wake-attachment", "", CompileInput{
		AgentID: "agt-1",
		Runtime: "codex",
		Binding: ConversationBinding{
			ChannelID:   "agentteam",
			ChannelType: "channel",
			TargetID:    "#ops",
		},
		Message: types.InboxMessage{
			MessageID:       "msg-attachment-only",
			Target:          "#ops",
			AuthorProfile:   types.AuthorProfile{ActorType: "human", DisplayName: "alice"},
			AttachmentIDs:   "att-1",
			AttachmentNames: "deploy.log",
		},
		UnreadSummary: map[string]int{"#ops": 2},
	})
	require.NoError(t, err)

	rendered, err := compiler.RenderRuntimeInput(context.Background(), env, "codex", RenderOptions{})
	require.NoError(t, err)

	assertGoldenText(t, "wake_attachment_only.rendered_turn.md", renderedPayloadMarkdown(rendered))
}

func TestPromptCompilerGoldens_DeliveryArtifacts(t *testing.T) {
	compiler := New(DefaultPolicy())
	env, err := compiler.CompileDeliveryEnvelope(context.Background(), EventRoutedMessageDelivery, "evt-golden-delivery-current", CompileInput{
		AgentID: "agt-1",
		Runtime: "codex",
		Delivery: DeliveryCompileInput{
			MessageID:         "msg-100",
			Target:            "#ops",
			ChannelType:       "channel",
			AuthorIdentity:    types.AuthorIdentity{ActorType: "human", DisplayName: "alice"},
			Content:           "Second delivery",
			Timestamp:         "2026-04-20T12:00:05Z",
			MessageType:       "text",
			DeliveryMode:      "turn_direct",
			RuntimeSessionID:  "sess-1",
			ThreadID:          "thread-1",
			ActiveTurnID:      "turn-1",
			WakeToken:         "wake-1",
			Attempt:           2,
			DeferredBatchID:   "batch-1",
			ReplayState:       "replayed",
			AcceptedAt:        time.Date(2026, 4, 20, 12, 0, 5, 0, time.UTC),
			OriginalOrderKey:  "order-1",
			PriorAttemptCount: 1,
			EvidenceClass:     EvidenceReplayStarted,
			Metadata: map[string]any{
				"reaction": map[string]any{
					"targetMessageId": "msg-root-2",
				},
			},
		},
	})
	require.NoError(t, err)

	rendered, err := compiler.RenderRuntimeInput(context.Background(), env, "codex", RenderOptions{})
	require.NoError(t, err)

	assertGoldenJSON(t, "delivery.runtime_input.json", env)
	assertGoldenJSON(t, "delivery.rendered_payload.json", rendered)
	assertGoldenText(t, "delivery.rendered_turn.md", renderedPayloadMarkdown(rendered))
}

func TestPromptCompilerGoldens_DeliveryWithAttachments(t *testing.T) {
	compiler := New(DefaultPolicy())
	env, err := compiler.CompileDeliveryEnvelope(context.Background(), EventRoutedMessageDelivery, "evt-golden-delivery-attachment", CompileInput{
		AgentID: "agt-1",
		Runtime: "codex",
		Delivery: DeliveryCompileInput{
			MessageID:      "msg-attachment",
			Target:         "#ops",
			ChannelType:    "channel",
			AuthorIdentity: types.AuthorIdentity{ActorType: "human", DisplayName: "alice"},
			Timestamp:      "2026-04-20T12:00:00Z",
			MessageType:    "text",
			Attachments: []types.MessageAttachment{
				{ID: "att-1", Filename: "deploy.log"},
				{ID: "att-2", Filename: "trace.json"},
			},
			DeliveryMode: "turn_direct",
		},
	})
	require.NoError(t, err)

	rendered, err := compiler.RenderRuntimeInput(context.Background(), env, "codex", RenderOptions{})
	require.NoError(t, err)

	assertGoldenJSON(t, "delivery_attachment.runtime_input.json", env)
	assertGoldenText(t, "delivery_attachment.rendered_turn.md", renderedPayloadMarkdown(rendered))
}

func renderedPayloadMarkdown(rendered RenderedPromptPayload) string {
	parts := make([]string, 0, len(rendered.RenderOrder))
	for _, slot := range rendered.RenderOrder {
		value, ok := rendered.Slots[slot]
		if !ok {
			continue
		}
		body := ""
		if text, ok := value.(string); ok {
			body = text
		} else {
			data, err := json.MarshalIndent(value, "", "  ")
			if err != nil {
				continue
			}
			body = string(data)
		}
		parts = append(parts, "## "+slot+"\n\n"+strings.TrimRight(body, "\n"))
	}
	if len(parts) == 0 {
		return ""
	}
	return strings.Join(parts, "\n\n") + "\n"
}

func assertGoldenJSON(t *testing.T, name string, value any) {
	t.Helper()
	data, err := json.MarshalIndent(value, "", "  ")
	require.NoError(t, err)
	assertGoldenText(t, name, string(append(data, '\n')))
}

func assertGoldenText(t *testing.T, name, got string) {
	t.Helper()
	path := filepath.Join("testdata", name)
	if os.Getenv(updateGoldenEnv) != "" {
		require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o755))
		require.NoError(t, os.WriteFile(path, []byte(got), 0o644))
	}
	want, err := os.ReadFile(path)
	require.NoError(t, err)
	require.Equal(t, string(want), got)
}