package cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	agentmodel "github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agentview"
	agentteamexec "github.com/tidbcloud/AgentTeam/desktop/internal/api/exec"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

type errorPolicyAgentReader struct {
	snapshot agentview.Snapshot
	listErr  error
	pageErr  error
	found    bool
	getErr   error
}

func (reader errorPolicyAgentReader) ListAgents(context.Context) ([]agentview.Snapshot, error) {
	if reader.listErr != nil {
		return nil, reader.listErr
	}
	return []agentview.Snapshot{reader.snapshot}, nil
}

func (reader errorPolicyAgentReader) ListAgentsPageWithAnchor(context.Context, string, string, string, int) ([]agentview.Snapshot, bool, string, string, error) {
	if reader.pageErr != nil {
		return nil, false, "", "", reader.pageErr
	}
	return []agentview.Snapshot{reader.snapshot}, false, "revision-1", "", nil
}

func (reader errorPolicyAgentReader) GetAgent(context.Context, string) (agentview.Snapshot, bool, error) {
	if reader.getErr != nil {
		return agentview.Snapshot{}, false, reader.getErr
	}
	return reader.snapshot, reader.found, nil
}

func startErrorPolicyRuntime(
	t *testing.T,
	reader agentteamexec.AgentReader,
	status int,
	body string,
) (agentteamexec.RuntimeCredentialLease, *agentteamexec.RuntimeAuthorizationStore) {
	t.Helper()
	gin.SetMode(gin.TestMode)
	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("Content-Type", "application/json")
		writer.WriteHeader(status)
		_, _ = writer.Write([]byte(body))
	}))
	t.Cleanup(upstream.Close)

	keys, err := token.NewKeyManager(t.TempDir() + "/config.yaml")
	require.NoError(t, err)
	authorizations := agentteamexec.NewRuntimeAuthorizationStore(keys)
	lease, err := authorizations.Create(agentteamexec.RuntimeAuthorizationConfig{
		ServerID: "server-1", AgentID: "agent-1", SessionID: "session-1", RunID: "run-1",
		WorkDir: t.TempDir(),
		Server: &apiclient.Client{
			ServerURL: upstream.URL, LocalAgentID: "agent-1", AgentID: "server-agent-1",
			AuthToken: "server-agent-token", RouteWorkspaceID: "workspace-1", HTTPClient: upstream.Client(),
		},
		TokenTTL: time.Hour,
	})
	require.NoError(t, err)
	router := gin.New()
	agentteamexec.RegisterRoutes(router.Group("/api/internal/exec"), agentteamexec.RouteDependencies{
		Authorizations: authorizations,
		Agents:         reader,
	})
	daemon := httptest.NewServer(router)
	t.Cleanup(daemon.Close)
	setExecCommandTokenFile(t, daemon.URL+"/api/internal/exec/servers/server-1/agents/agent-1", lease.Token)
	return lease, authorizations
}

func assertErrorPolicyCLIByIdentity(t *testing.T, identity, domain, action, input string, expected agentteamexec.ErrorCode) {
	t.Helper()
	stdout, stderr, err := executeCobra(newExecCommand(), domain, action, input)
	require.Error(t, err)
	require.Empty(t, stdout)
	var envelope struct {
		Error agentteamexec.ErrorEnvelope `json:"error"`
	}
	require.NoError(t, json.Unmarshal([]byte(stderr), &envelope))
	assertClosedErrorEnvelope(t, stderr, envelope.Error, reviewedErrorEnvelopeKeys(identity, domain, action, expected))
	require.Equal(t, expected, envelope.Error.Code, "policy identity %s route %s %s", identity, domain, action)
	if expected == agentteamexec.CodeResultBudgetExceeded {
		assertReviewedResultBudgetValues(t, envelope.Error, domain, action)
	} else if problem, ok := reviewedRouteProblem(identity, domain, action); ok {
		require.Equal(t, problem, envelope.Error.Problem)
	}
	require.NotEmpty(t, envelope.Error.Problem)
	exitCode, ok := ExecExitCode(err)
	require.True(t, ok)
	require.NotZero(t, exitCode)
	require.Equal(t, agentteamexec.ErrorExitCode(expected), exitCode)
}

func reviewedErrorEnvelopeKeys(identity, domain, action string, code agentteamexec.ErrorCode) []string {
	keys := []string{"code", "problem"}
	if code == agentteamexec.CodeResultBudgetExceeded {
		keys = []string{"code", "hint", "next", "problem"}
		if domain == "project" || (domain == "agent" && action == "start") {
			keys = []string{"code", "hint", "metadata", "problem"}
		}
	}
	if code == agentteamexec.CodeValidation && identity == "exec.local-read.v1" {
		keys = []string{"code", "hint", "problem"}
	}
	return keys
}

func assertClosedErrorEnvelope(t *testing.T, stderr string, envelope agentteamexec.ErrorEnvelope, expectedKeys []string) {
	t.Helper()
	var raw map[string]json.RawMessage
	require.NoError(t, json.Unmarshal([]byte(stderr), &raw))
	require.Equal(t, []string{"error"}, sortedJSONKeys(raw))
	var fields map[string]json.RawMessage
	require.NoError(t, json.Unmarshal(raw["error"], &fields))
	require.Equal(t, expectedKeys, sortedJSONKeys(fields), "the reviewed error class must emit an exact key set")
	require.Contains(t, fields, "code")
	require.Contains(t, fields, "problem")
	var code agentteamexec.ErrorCode
	require.NoError(t, json.Unmarshal(fields["code"], &code))
	require.Equal(t, envelope.Code, code)
	var problem string
	require.NoError(t, json.Unmarshal(fields["problem"], &problem))
	require.Equal(t, envelope.Problem, problem)
	if envelope.Retryable {
		var retryable bool
		require.NoError(t, json.Unmarshal(fields["retryable"], &retryable))
		require.True(t, retryable)
	} else {
		_, present := fields["retryable"]
		require.False(t, present, "false retryable is omitted by the canonical projector")
	}
	for key, expectedValue := range map[string]any{
		"hint":        envelope.Hint,
		"next":        envelope.Next,
		"stage":       envelope.Stage,
		"path":        envelope.Path,
		"retryable":   envelope.Retryable,
		"serverCode":  envelope.ServerCode,
		"diagnostics": envelope.Diagnostics,
		"truncation":  envelope.Truncation,
		"truncated":   envelope.Truncated,
		"metadata":    envelope.Metadata,
	} {
		actual, present := fields[key]
		if !present {
			switch value := expectedValue.(type) {
			case bool:
				require.False(t, value, "omitted %s must be false", key)
			case nil:
				require.Nil(t, value, "omitted %s must be nil", key)
			default:
				require.Empty(t, value, "omitted %s must be empty", key)
			}
			continue
		}
		encoded, err := json.Marshal(expectedValue)
		require.NoError(t, err)
		require.JSONEq(t, string(encoded), string(actual), "canonical %s value must survive daemon -> LocalClient -> Cobra", key)
	}
	if diagnostics, ok := fields["diagnostics"]; ok {
		var values []agentteamexec.Diagnostic
		var rawValues []map[string]json.RawMessage
		require.NoError(t, json.Unmarshal(diagnostics, &values))
		require.NoError(t, json.Unmarshal(diagnostics, &rawValues))
		require.Len(t, rawValues, len(values))
		for index, rawDiagnostic := range rawValues {
			expectedDiagnostic := []string{"code", "message", "severity"}
			if values[index].Path != "" {
				expectedDiagnostic = append(expectedDiagnostic, "path")
			}
			if values[index].Field != "" {
				expectedDiagnostic = append(expectedDiagnostic, "field")
			}
			if values[index].Index != nil {
				expectedDiagnostic = append(expectedDiagnostic, "index")
			}
			slices.Sort(expectedDiagnostic)
			require.Equal(t, expectedDiagnostic, sortedJSONKeys(rawDiagnostic))
			require.NotEmpty(t, values[index].Code)
			require.NotEmpty(t, values[index].Message)
		}
	}
	require.NotContains(t, stderr, `"cause"`)
}

func reviewedRouteProblem(identity, domain, action string) (string, bool) {
	switch identity {
	case "exec.remote-read.v1":
		return "list workspace failed", domain == "workspace" && action == "list"
	case "exec.natural-mutation.v1":
		return "create project failed", domain == "project" && action == "create"
	case "exec.receipt-mutation.v1":
		return "start agent failed", domain == "agent" && action == "start"
	case "exec.workflow-mutation.v1":
		return "workflow plan show failed", domain == "workflow" && action == "plan"
	default:
		return "", false
	}
}

func assertMutationAppliedEnvelope(t *testing.T, stderr string) {
	t.Helper()
	var rendered struct {
		Error agentteamexec.ErrorEnvelope `json:"error"`
	}
	require.NoError(t, json.Unmarshal([]byte(stderr), &rendered))
	assertClosedErrorEnvelope(t, stderr, rendered.Error, []string{"code", "hint", "metadata", "problem"})
	require.Equal(t, agentteamexec.CodeResultBudgetExceeded, rendered.Error.Code)
	require.Equal(t, "Mutation outcome is recorded; inspect the receipt and retrieval path before issuing another command.", rendered.Error.Hint)
	require.Empty(t, rendered.Error.Next)
	require.False(t, rendered.Error.Retryable)
	receiptJSON, err := json.Marshal(rendered.Error.Metadata)
	require.NoError(t, err)
	require.JSONEq(t, `{"mutation_applied":true,"mutation_receipt":{"applied":true,"action":"channel update","identifiers":{"channelId":"chn_1"},"revision":"3","idempotencyKey":"req_1","retrievalPath":"agentteam exec channel get {\"channelId\":\"chn_1\"}"}}`, string(receiptJSON))
}

func assertReviewedResultBudgetValues(t *testing.T, envelope agentteamexec.ErrorEnvelope, domain, action string) {
	t.Helper()
	require.Equal(t, fmt.Sprintf("Exec action %s %s result exceeds the result budget (64 KiB success boundary)", domain, action), envelope.Problem)
	type continuation struct {
		policy string
		next   string
	}
	policies := map[string]continuation{
		"workspace list": {policy: "cursor", next: "agentteam exec workspace list with a narrower request"},
		"agent list":     {policy: "cursor", next: "agentteam exec agent list with a narrower request"},
		"agent get":      {policy: "narrow", next: "agentteam exec agent get with a narrower request"},
		"project create": {policy: "bounded", next: "agentteam exec project create with a narrower request"},
		"agent start":    {policy: "bounded", next: "agentteam exec agent start with a narrower request"},
		"workflow plan":  {policy: "workflow", next: "agentteam exec workflow plan with a narrower request"},
	}
	selected, ok := policies[domain+" "+action]
	require.True(t, ok, "reviewed result-budget route %s %s", domain, action)
	if envelope.Metadata != nil && envelope.Metadata["mutation_applied"] == true {
		require.Equal(t, "Mutation outcome is recorded; inspect the receipt and retrieval path before issuing another command.", envelope.Hint)
		require.Empty(t, envelope.Next)
		require.False(t, envelope.Retryable)
		return
	}
	require.Equal(t, fmt.Sprintf("No partial result was returned (result_policy=%s); %s.", selected.policy, map[string]string{
		"workspace list": "repeat workspace list with a smaller limit or the next cursor",
		"agent list":     "repeat agent list with a smaller limit or the next cursor",
		"agent get":      "request one agent by its stable agentId",
		"workflow plan":  "use the workflow plan action's operation-specific section or pagination input",
	}[domain+" "+action]), envelope.Hint)
	require.Equal(t, selected.next, envelope.Next)
	require.False(t, envelope.Retryable)
}

func sortedJSONKeys(values map[string]json.RawMessage) []string {
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	slices.Sort(keys)
	return keys
}

func TestExecBuiltinErrorPolicyMatrixUsesProductionDaemonAndCobra(t *testing.T) {
	statusCases := []struct {
		name   string
		status int
		code   agentteamexec.ErrorCode
	}{
		{name: "validation", status: http.StatusBadRequest, code: agentteamexec.CodeValidation},
		{name: "permission", status: http.StatusUnauthorized, code: agentteamexec.CodePermissionDenied},
		{name: "not found", status: http.StatusNotFound, code: agentteamexec.CodeNotFound},
		{name: "conflict", status: http.StatusConflict, code: agentteamexec.CodeConflict},
		{name: "rate limited", status: http.StatusTooManyRequests, code: agentteamexec.CodeRateLimited},
		{name: "internal", status: http.StatusInternalServerError, code: agentteamexec.CodeInternal},
	}
	body := `{"error":"bounded production error"}`
	cases := []struct {
		name, identity, domain, action, input string
	}{
		{name: "remote read", identity: "exec.remote-read.v1", domain: "workspace", action: "list", input: `{}`},
		{name: "natural mutation", identity: "exec.natural-mutation.v1", domain: "project", action: "create", input: `{"rootPath":"/tmp","name":"project"}`},
		{name: "receipt mutation", identity: "exec.receipt-mutation.v1", domain: "agent", action: "start", input: `{"requestId":"matrix-request","agentId":"agt-target"}`},
		{name: "workflow mutation", identity: "exec.workflow-mutation.v1", domain: "workflow", action: "plan", input: `{"action":"show","planId":"pln_1"}`},
	}
	for _, policyCase := range cases {
		t.Run(policyCase.name, func(t *testing.T) {
			for _, statusCase := range statusCases {
				t.Run(statusCase.name, func(t *testing.T) {
					startErrorPolicyRuntime(t, errorPolicyAgentReader{}, statusCase.status, body)
					assertErrorPolicyCLIByIdentity(t, policyCase.identity, policyCase.domain, policyCase.action, policyCase.input, statusCase.code)
				})
			}
		})
	}
}

func TestExecErrorPolicyMatrixProjectsCanonicalDiagnostics(t *testing.T) {
	body := `{"error":"bounded diagnostic","code":"plan_request_failed","stage":"plan","path":"$.input","hint":"retry after review","retryable":true,"diagnostics":[{"severity":"error","code":"input.invalid","path":"$.input.name","field":"name","index":2,"message":"bounded diagnostic"}]}`
	startErrorPolicyRuntime(t, errorPolicyAgentReader{}, http.StatusBadRequest, body)
	stdout, stderr, err := executeCobra(newExecCommand(), "workflow", "plan", `{"action":"show","planId":"pln_1"}`)
	require.Error(t, err)
	require.Empty(t, stdout)
	var response struct {
		Error agentteamexec.ErrorEnvelope `json:"error"`
	}
	require.NoError(t, json.Unmarshal([]byte(stderr), &response))
	assertClosedErrorEnvelope(t, stderr, response.Error, []string{"code", "diagnostics", "hint", "path", "problem", "retryable", "serverCode", "stage"})
	require.Equal(t, agentteamexec.CodeValidation, response.Error.Code)
	require.Equal(t, "plan", response.Error.Stage)
	require.Equal(t, "$.input", response.Error.Path)
	require.Equal(t, "plan_request_failed", response.Error.ServerCode)
	require.True(t, response.Error.Retryable)
	require.Len(t, response.Error.Diagnostics, 1)
	require.Equal(t, "name", response.Error.Diagnostics[0].Field)
	require.NotNil(t, response.Error.Diagnostics[0].Index)
	require.Equal(t, 2, *response.Error.Diagnostics[0].Index)
	require.Equal(t, 64, agentteamexec.ErrorExitCode(response.Error.Code))
}

func TestExecLocalErrorPolicyMatrixUsesProductionAgentBoundary(t *testing.T) {
	large := strings.Repeat("x", 96*1024)
	cases := []struct {
		name, identity, action, input string
		reader                        errorPolicyAgentReader
		revoke                        bool
		expected                      agentteamexec.ErrorCode
	}{
		{name: "local read permission", identity: "exec.local-read.v1", action: "list", input: `{}`, revoke: true, expected: agentteamexec.CodePermissionDenied},
		{name: "local read validation", identity: "exec.local-read.v1", action: "list", input: `{"limit":101}`, expected: agentteamexec.CodeValidation},
		{name: "local read conflict", identity: "exec.local-read.v1", action: "list", input: `{}`, reader: errorPolicyAgentReader{pageErr: agentmodel.NewPageError(agentmodel.PageErrorStaleRevision, "stale")}, expected: agentteamexec.CodeConflict},
		{name: "local read internal", identity: "exec.local-read.v1", action: "list", input: `{}`, reader: errorPolicyAgentReader{pageErr: errors.New("source failed")}, expected: agentteamexec.CodeInternal},
		{name: "local read result budget", identity: "exec.local-read.v1", action: "list", input: `{}`, reader: errorPolicyAgentReader{snapshot: agentview.Snapshot{ID: "agt-1", Name: large}}, expected: agentteamexec.CodeResultBudgetExceeded},
		{name: "local get permission", identity: "exec.local-get.v1", action: "get", input: `{"agentId":"agt-1"}`, revoke: true, expected: agentteamexec.CodePermissionDenied},
		{name: "local get not found", identity: "exec.local-get.v1", action: "get", input: `{"agentId":"missing"}`, reader: errorPolicyAgentReader{found: false}, expected: agentteamexec.CodeNotFound},
		{name: "local get internal", identity: "exec.local-get.v1", action: "get", input: `{"agentId":"agt-1"}`, reader: errorPolicyAgentReader{getErr: errors.New("source failed")}, expected: agentteamexec.CodeInternal},
		{name: "local get result budget", identity: "exec.local-get.v1", action: "get", input: `{"agentId":"agt-1"}`, reader: errorPolicyAgentReader{found: true, snapshot: agentview.Snapshot{ID: "agt-1", Name: large}}, expected: agentteamexec.CodeResultBudgetExceeded},
	}
	for _, policyCase := range cases {
		t.Run(policyCase.name, func(t *testing.T) {
			lease, authorizations := startErrorPolicyRuntime(t, policyCase.reader, http.StatusOK, `{}`)
			if policyCase.revoke {
				<-authorizations.Revoke(lease.AuthorizationID)
			}
			assertErrorPolicyCLIByIdentity(t, policyCase.identity, "agent", policyCase.action, policyCase.input, policyCase.expected)
		})
	}
}

func TestExecBuiltinErrorPolicyResultBudgetMatrixUsesProductionRoutes(t *testing.T) {
	large := strings.Repeat("x", 96*1024)
	cases := []struct {
		name, identity, domain, action, input, body string
	}{
		{
			name: "remote read", identity: "exec.remote-read.v1", domain: "workspace", action: "list", input: `{"kind":"humans"}`,
			body: `{"kind":"humans","humans":[{"name":"` + large + `"}]}`,
		},
		{
			name: "natural mutation", identity: "exec.natural-mutation.v1", domain: "project", action: "create", input: `{"rootPath":"/tmp","name":"project"}`,
			body: `{"project":{"id":"prj-1","rootPath":"/tmp","name":"` + large + `"}}`,
		},
		{
			name: "receipt mutation", identity: "exec.receipt-mutation.v1", domain: "agent", action: "start", input: `{"requestId":"matrix-budget","agentId":"agt-target"}`,
			body: `{"status":"` + large + `"}`,
		},
		{
			name: "workflow mutation", identity: "exec.workflow-mutation.v1", domain: "workflow", action: "plan", input: `{"action":"validate","source":"apiVersion: loop.plan/v1\nkind: PlanVersion\n"}`,
			body: `{"diagnostics":[{"severity":"error","code":"invalid","path":"$.input","message":"` + large + `"}]}`,
		},
	}
	for _, policyCase := range cases {
		t.Run(policyCase.name, func(t *testing.T) {
			startErrorPolicyRuntime(t, errorPolicyAgentReader{}, http.StatusOK, policyCase.body)
			assertErrorPolicyCLIByIdentity(t, policyCase.identity, policyCase.domain, policyCase.action, policyCase.input, agentteamexec.CodeResultBudgetExceeded)
		})
	}
}