package agent

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/api/workflow/binding"
	"github.com/tidbcloud/AgentTeam/internal/api/workflow/plan"
	storeagent "github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/workflow/orchestrator"
	"github.com/tidbcloud/AgentTeam/internal/workflow/planmanage"
)

type workflowRecoveryResetProbe struct {
	calls     []workflowRecoveryResetCall
	nextError error
}

type workflowRecoveryResetCall struct {
	workspaceID string
	agentID     string
}

func (probe *workflowRecoveryResetProbe) ResetAgentRecovery(_ context.Context, workspaceID, agentID string) error {
	probe.calls = append(probe.calls, workflowRecoveryResetCall{workspaceID: workspaceID, agentID: agentID})
	err := probe.nextError
	probe.nextError = nil
	return err
}

type workflowRecoveryPlanService struct {
	plan.Service
}

func (workflowRecoveryPlanService) List(context.Context, planmanage.ListInput) (planmanage.PlanPage, error) {
	return planmanage.PlanPage{}, nil
}

type workflowRecoveryBindingService struct {
	binding.Service
}

type recoveryStatusCommandRuntime struct {
	calls       []orchestrator.CallbackCommand
	sideEffects map[string]int
	seenKeys    map[string]bool
}

func (runtime *recoveryStatusCommandRuntime) ExecuteCallback(_ context.Context, command orchestrator.CallbackCommand) (orchestrator.CallbackResult, error) {
	runtime.calls = append(runtime.calls, command)
	key := strings.TrimSpace(command.IdempotencyKey)
	if key == "" {
		key = command.Action
	}
	wasSeen := runtime.seenKeys[key]
	if !wasSeen {
		runtime.seenKeys[key] = true
		runtime.sideEffects[key]++
	}
	status, runStatus := "waiting", "running"
	switch command.Action {
	case orchestrator.CallbackActionComplete:
		status, runStatus = "completed", "succeeded"
	case orchestrator.CallbackActionFail:
		status, runStatus = "failed", "failed"
	}
	return orchestrator.CallbackResult{
		RunID: command.RunID, CallbackID: command.CallbackID,
		Status: status, RunStatus: runStatus, Replayed: wasSeen,
	}, nil
}

func (runtime *recoveryStatusCommandRuntime) ListCallbacks(context.Context, orchestrator.CallbackListCommand) (orchestrator.CallbackListResult, error) {
	return orchestrator.CallbackListResult{}, nil
}

func (workflowRecoveryBindingService) List(context.Context, binding.ListInput) (binding.BindingPage, error) {
	return binding.BindingPage{}, nil
}

func (workflowRecoveryBindingService) ReplayCommand(context.Context, binding.ReplayCommandInput) (binding.ReplayCommandResult, bool, error) {
	return binding.ReplayCommandResult{}, false, nil
}

func TestWorkflowRecoverySuccessExitResetsAllFourWorkflowRoutes(t *testing.T) {
	gin.SetMode(gin.TestMode)
	probe := &workflowRecoveryResetProbe{}
	auth := fakeAgentAuthenticator{agent: storeagent.AuthenticatedAgent{
		AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test",
	}}
	base := "/api/v1/workspaces/:workspaceId/agents/:agentId/runtime"
	router := gin.New()
	registerAgentWorkflowPlanRoutes(router, base, auth, workflowRecoveryPlanService{}, nil, nil, probe)
	registerAgentWorkflowBindingRoutes(router, base, auth, workflowRecoveryBindingService{}, publicRouteRunOwner{}, publicRouteRunWorkspace{}, probe)
	registerAgentWorkflowRunRoutes(router, base, auth, publicRouteRunService{}, nil, nil, nil, probe)
	registerAgentWorkflowStepRoutes(router, base, auth, &publicRouteStepRuntime{}, nil, probe)

	for _, route := range []string{"/workflow/plan", "/workflow/trigger", "/workflow/run"} {
		response := postWorkflowPublicCommand(t, router, route, map[string]any{"action": "list"})
		require.Equal(t, http.StatusOK, response.Code, route+": "+response.Body.String())
	}
	response := postWorkflowStepCommand(t, router, map[string]any{"action": "list"})
	require.Equal(t, http.StatusOK, response.Code, response.Body.String())
	require.Len(t, probe.calls, 4)
}

func TestWorkflowRecoverySuccessExitResetsOnlySuccessfulRunAndStepActions(t *testing.T) {
	gin.SetMode(gin.TestMode)
	probe := &workflowRecoveryResetProbe{}
	auth := fakeAgentAuthenticator{agent: storeagent.AuthenticatedAgent{
		AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test",
	}}
	base := "/api/v1/workspaces/:workspaceId/agents/:agentId/runtime"
	router := gin.New()
	registerAgentWorkflowRunRoutes(router, base, auth, publicRouteRunService{}, nil, nil, nil, probe)
	runtime := &publicRouteStepRuntime{}
	registerAgentWorkflowStepRoutes(router, base, auth, runtime, nil, probe)

	response := postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{
		"action": "list",
	})
	require.Equal(t, http.StatusOK, response.Code, response.Body.String())
	require.Len(t, probe.calls, 1)
	require.Equal(t, workflowRecoveryResetCall{workspaceID: "wsp_test", agentID: "agt_0000000000test"}, probe.calls[0])

	// Request validation/dispatch failures never reach the shared success exit.
	response = postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{
		"action": "unsupported",
	})
	require.NotEqual(t, http.StatusOK, response.Code)
	require.Len(t, probe.calls, 1)

	response = postWorkflowStepCommand(t, router, map[string]any{"action": "list"})
	require.Equal(t, http.StatusOK, response.Code, response.Body.String())
	require.Len(t, probe.calls, 2)

	response = postWorkflowStepCommand(t, router, map[string]any{"action": "complete", "workflowRunId": "wri_1", "workflowStepId": "wfx_1"})
	require.Equal(t, http.StatusBadRequest, response.Code, response.Body.String())
	require.Len(t, probe.calls, 2)
}

func TestWorkflowRecoverySuccessExitFailsTheRouteWhenResetFails(t *testing.T) {
	gin.SetMode(gin.TestMode)
	probe := &workflowRecoveryResetProbe{nextError: errors.New("temporary state outage")}
	auth := fakeAgentAuthenticator{agent: storeagent.AuthenticatedAgent{
		AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test",
	}}
	router := gin.New()
	registerAgentWorkflowRunRoutes(
		router, "/api/v1/workspaces/:workspaceId/agents/:agentId/runtime", auth,
		publicRouteRunService{}, nil, nil, nil, probe,
	)
	response := postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{"action": "list"})
	require.Equal(t, http.StatusInternalServerError, response.Code, response.Body.String())
	require.Len(t, probe.calls, 1)
	var failure struct {
		Code      string `json:"code"`
		Stage     string `json:"stage"`
		Retryable bool   `json:"retryable"`
		Error     string `json:"error"`
	}
	require.NoError(t, json.Unmarshal(response.Body.Bytes(), &failure))
	require.Equal(t, workflowRecoveryResetFailedCode, failure.Code)
	require.Equal(t, "run", failure.Stage)
	require.True(t, failure.Retryable)
	require.Equal(t, "Workflow recovery state could not be persisted; retry the request", failure.Error)

	// A retry replays the same business read and can complete the reset.
	response = postWorkflowPublicCommand(t, router, "/workflow/run", map[string]any{"action": "list"})
	require.Equal(t, http.StatusOK, response.Code, response.Body.String())
	require.Len(t, probe.calls, 2)
}

func TestWorkflowRecoverySuccessExitCoversParkAndTerminalReplay(t *testing.T) {
	gin.SetMode(gin.TestMode)
	probe := &workflowRecoveryResetProbe{}
	runtime := &recoveryStatusCommandRuntime{
		sideEffects: make(map[string]int), seenKeys: make(map[string]bool),
	}
	auth := fakeAgentAuthenticator{agent: storeagent.AuthenticatedAgent{
		AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test",
	}}
	router := gin.New()
	registerAgentWorkflowStepRoutes(router, "/api/v1/workspaces/:workspaceId/agents/:agentId/runtime", auth, runtime, nil, probe)

	park := postWorkflowStepCommand(t, router, map[string]any{
		"action": "park", "workflowRunId": "wri_1", "workflowStepId": "wfx_1",
	})
	require.Equal(t, http.StatusOK, park.Code, park.Body.String())
	require.Equal(t, 1, len(probe.calls))
	require.Equal(t, orchestrator.CallbackActionPark, runtime.calls[0].Action)
	require.Equal(t, 1, runtime.sideEffects[orchestrator.CallbackActionPark])

	probe.nextError = errors.New("reset persistence unavailable")
	completeInput := map[string]any{
		"action": "complete", "workflowRunId": "wri_1", "workflowStepId": "wfx_1",
		"idempotencyKey": "complete-key", "output": map[string]any{"ok": true},
	}
	// The first callback commit succeeds, but the final reset hook fails. The
	// client must replay the same business key; the callback side effect is not
	// repeated and the reset succeeds on the replay.
	firstComplete := postWorkflowStepCommand(t, router, completeInput)
	require.Equal(t, http.StatusInternalServerError, firstComplete.Code, firstComplete.Body.String())
	require.Equal(t, 1, runtime.sideEffects["complete-key"])
	require.Equal(t, 2, len(probe.calls))
	secondComplete := postWorkflowStepCommand(t, router, completeInput)
	require.Equal(t, http.StatusOK, secondComplete.Code, secondComplete.Body.String())
	require.Equal(t, 1, runtime.sideEffects["complete-key"], "idempotent replay must not repeat the business side effect")
	require.Equal(t, 3, len(probe.calls))
	var replay workflowStepResponse
	require.NoError(t, json.Unmarshal(secondComplete.Body.Bytes(), &replay))
	require.True(t, replay.Replayed)

	fail := postWorkflowStepCommand(t, router, map[string]any{
		"action": "fail", "workflowRunId": "wri_1", "workflowStepId": "wfx_1",
		"idempotencyKey": "fail-key", "code": "blocked", "message": "blocked by test",
	})
	require.Equal(t, http.StatusOK, fail.Code, fail.Body.String())
	require.Equal(t, 1, runtime.sideEffects["fail-key"])
	require.Equal(t, 4, len(probe.calls))
	require.Equal(t, orchestrator.CallbackActionFail, runtime.calls[3].Action)
}

func TestWorkflowRecoverySuccessOptionIgnoresNon2xxResponses(t *testing.T) {
	gin.SetMode(gin.TestMode)
	probe := &workflowRecoveryResetProbe{}
	auth := fakeAgentAuthenticator{agent: storeagent.AuthenticatedAgent{
		AgentID: "agt_0000000000test", WorkspaceID: "wsp_test", MemberID: "00000000test",
	}}
	router := gin.New()
	router.POST(
		"/api/v1/workspaces/:workspaceId/agents/:agentId/runtime/workflow/test",
		RequireAgentRuntimePrincipal(auth),
		api.HandleAgentScoped(
			func(api.AgentScopedRequest) (api.Response, error) {
				return api.Response{Status: http.StatusBadRequest, Body: map[string]string{"error": "expected"}}, nil
			},
			workflowRecoverySuccessOption(probe),
		),
	)
	request := postWorkflowPublicCommand(t, router, "/workflow/test", map[string]any{})
	require.Equal(t, http.StatusBadRequest, request.Code, request.Body.String())
	require.Empty(t, probe.calls)
}