package agent

import (
	"bytes"
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/core/domain"
	workflowpayload "github.com/tidbcloud/AgentTeam/internal/workflow/payload"
	"github.com/tidbcloud/AgentTeam/internal/workflow/publicerr"
	"github.com/tidbcloud/AgentTeam/internal/workflow/runartifact"
	runruntime "github.com/tidbcloud/AgentTeam/internal/workflow/orchestrator"
)

const (
	// The request envelope can carry both ordinary callback output and the
	// compiler-owned WorkItems value. Keep transport headroom above the
	// advertised 1 MiB WorkItems-only budget and align one callback request
	// with the durable runtime's single-payload boundary.
	workflowRunCallbackMaxBodyBytes        = 4 << 20
	workflowRunCallbackMaxJSONNesting      = 64
	workflowRunCallbackIdempotencyMaxLen   = 128
	workflowRunCallbackNotReadyCode        = "step_not_ready"
	workflowRunCallbackAlreadyResolvedCode = "step_already_resolved"
)

type WorkflowStepHandler struct {
	runtime runruntime.CallbackRuntime
	logger  *slog.Logger
}

func NewWorkflowStepHandler(runtime runruntime.CallbackRuntime, logger *slog.Logger) *WorkflowStepHandler {
	return &WorkflowStepHandler{runtime: runtime, logger: logger}
}

type workflowStepRequest struct {
	Action         string          `json:"action"`
	Section        string          `json:"section,omitempty"`
	Path           string          `json:"path,omitempty"`
	Output         json.RawMessage `json:"output,omitempty"`
	WorkScope      json.RawMessage `json:"workScope,omitempty"`
	WorkItems      json.RawMessage `json:"workItems,omitempty"`
	Code           string          `json:"code,omitempty"`
	Message        string          `json:"message,omitempty"`
	IdempotencyKey string          `json:"idempotencyKey,omitempty"`
	Cursor         string          `json:"cursor,omitempty"`
	Limit          int             `json:"limit,omitempty"`
	WorkflowRunID  string          `json:"workflowRunId,omitempty"`
	WorkflowStepID string          `json:"workflowStepId,omitempty"`
}

type workflowStepResponse struct {
	RunID             string                                `json:"runId"`
	StepID            string                                `json:"stepId"`
	Status            string                                `json:"status"`
	RunStatus         string                                `json:"runStatus"`
	Input             json.RawMessage                       `json:"input,omitempty"`
	OutputContract    json.RawMessage                       `json:"outputContract,omitempty"`
	WorkItemsContract *runruntime.CallbackWorkItemsContract `json:"workItemsContract,omitempty"`
	DeliveryScope     *domain.MessageScopeValue             `json:"deliveryScope,omitempty"`
	WorkScope         json.RawMessage                       `json:"workScope,omitempty"`
	RelatedScopes     *[]domain.MessageScopeValue           `json:"relatedScopes,omitempty"`
	Replayed          bool                                  `json:"replayed"`
	Payload           *workflowPayloadResponse              `json:"payload,omitempty"`
	InputRef          *workflowPayloadRef                   `json:"inputRef,omitempty"`
	OutputContractRef *workflowPayloadRef                   `json:"outputContractRef,omitempty"`
	WorkItemsRef      *workflowPayloadRef                   `json:"workItemsRef,omitempty"`
	// TeamMemberContext is the frozen minimal Workflow-specific member context
	// for a responder of any Run with a frozen Team ID (Team-linked or
	// dynamic-explicit). It is surfaced verbatim from the durable effect
	// receipt; it never contains the team roster, other members' duties, wake
	// configuration, or Team division hints.
	TeamMemberContext *runruntime.CallbackTeamMemberContext `json:"teamMemberContext,omitempty"`
}

type workflowStepListResponse struct {
	Steps      []workflowStepSummary `json:"steps"`
	NextCursor string                `json:"nextCursor,omitempty"`
}

type workflowStepSummary struct {
	RunID     string `json:"runId"`
	StepID    string `json:"stepId"`
	Status    string `json:"status"`
	RunStatus string `json:"runStatus"`
}

// registerAgentWorkflowStepRoutes exposes Step completion on an
// Agent-authenticated runtime route. Run and Step IDs are caller-selected
// resource locators. The route principal supplies Agent/workspace identity,
// and CallbackService reauthorizes the selector against the frozen responder
// receipt before returning or mutating callback state.
func registerAgentWorkflowStepRoutes(
	router gin.IRoutes,
	base string,
	auth AgentAuthenticator,
	runtime runruntime.CallbackRuntime,
	logger *slog.Logger,
	recovery ...WorkflowRecoveryResetter,
) {
	if router == nil || auth == nil || runtime == nil || strings.TrimSpace(base) == "" {
		return
	}
	handler := NewWorkflowStepHandler(runtime, logger)
	router.POST(
		strings.TrimSuffix(base, "/")+"/workflow/step",
		RequireAgentRuntimePrincipal(auth),
		api.HandleAgentScoped(
			handler.Execute,
			api.WithErrorClassifier(handler.classifyError),
			workflowRecoverySuccessOption(recovery...),
		),
	)
}

func (handler *WorkflowStepHandler) Execute(request api.AgentScopedRequest) (api.Response, error) {
	if handler == nil || handler.runtime == nil {
		return api.Response{}, publicError(
			http.StatusServiceUnavailable,
			"workflow_run_unavailable",
			"Workflow Steps are not enabled",
		)
	}
	body, err := decodeWorkflowCLIJSON[workflowStepRequest](
		request.Gin,
		workflowRunCallbackMaxBodyBytes,
		workflowRunCallbackMaxJSONNesting,
		workflowRunCallbackJSONError,
	)
	if err != nil {
		return api.Response{}, err
	}
	if strings.TrimSpace(request.Principal.MemberID) == "" {
		return api.Response{}, publicError(http.StatusBadRequest, "invalid_request_body", "authenticated Agent member identity is incomplete")
	}
	if strings.TrimSpace(body.Action) == "list" {
		command, err := normalizeWorkflowStepListRequestWithMember(
			request.Principal.WorkspaceID,
			request.Principal.AgentID,
			request.Principal.MemberID,
			body,
		)
		if err != nil {
			return api.Response{}, err
		}
		result, err := handler.runtime.ListCallbacks(request.Context(), command)
		if err != nil {
			return api.Response{}, err
		}
		steps := make([]workflowStepSummary, 0, len(result.Steps))
		for _, step := range result.Steps {
			steps = append(steps, workflowStepSummary{
				RunID: step.RunID, StepID: step.CallbackID,
				Status: step.Status, RunStatus: step.RunStatus,
			})
		}
		return api.OK(workflowStepListResponse{Steps: steps, NextCursor: result.NextCursor}), nil
	}
	command, err := normalizeWorkflowStepRequestWithMember(
		request.Principal.WorkspaceID,
		request.Principal.AgentID,
		request.Principal.MemberID,
		body,
	)
	if err != nil {
		return api.Response{}, err
	}
	if strings.TrimSpace(command.Section) != "" && strings.TrimSpace(command.Section) != "compact" {
		reader, ok := handler.runtime.(runruntime.CallbackPayloadReader)
		if !ok {
			return api.Response{}, publicError(http.StatusServiceUnavailable, "workflow_run_unavailable", "Workflow payload reads are not enabled")
		}
		payload, err := reader.ReadCallbackSection(request.Context(), command)
		if err != nil {
			return api.Response{}, err
		}
		return api.OK(workflowStepResponseFromPayload(payload)), nil
	}
	result, err := handler.runtime.ExecuteCallback(request.Context(), command)
	if err != nil {
		return api.Response{}, err
	}
	return api.OK(workflowStepResponseFromResult(result)), nil
}

func normalizeWorkflowStepListRequest(workspaceID, agentID string, request workflowStepRequest) (runruntime.CallbackListCommand, error) {
	return normalizeWorkflowStepListRequestWithMember(workspaceID, agentID, "", request)
}

func normalizeWorkflowStepListRequestWithMember(workspaceID, agentID, memberID string, request workflowStepRequest) (runruntime.CallbackListCommand, error) {
	command := runruntime.CallbackListCommand{
		WorkspaceID: strings.TrimSpace(workspaceID),
		AgentID:     strings.TrimSpace(agentID),
		MemberID:    strings.TrimSpace(memberID),
		Cursor:      strings.TrimSpace(request.Cursor),
		Limit:       request.Limit,
	}
	if command.Limit == 0 {
		command.Limit = 50
	}
	if command.WorkspaceID == "" || command.AgentID == "" || command.Limit < 1 || command.Limit > 100 {
		return runruntime.CallbackListCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "list requires authenticated identity and limit between 1 and 100")
	}
	if strings.TrimSpace(request.WorkflowRunID) != "" || strings.TrimSpace(request.WorkflowStepID) != "" ||
		strings.TrimSpace(request.Section) != "" || strings.TrimSpace(request.Path) != "" ||
		len(bytes.TrimSpace(request.Output)) != 0 || len(bytes.TrimSpace(request.WorkScope)) != 0 ||
		len(bytes.TrimSpace(request.WorkItems)) != 0 || strings.TrimSpace(request.Code) != "" ||
		strings.TrimSpace(request.Message) != "" || strings.TrimSpace(request.IdempotencyKey) != "" {
		return runruntime.CallbackListCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "list accepts only cursor and limit")
	}
	return command, nil
}

func workflowStepResponseFromPayload(result runruntime.CallbackPayloadRead) workflowStepResponse {
	return workflowStepResponse{
		RunID: result.RunID, StepID: result.CallbackID,
		Status: result.Status, RunStatus: result.RunStatus,
		Payload: &workflowPayloadResponse{
			Section: result.Page.Section, Encoding: result.Page.Encoding,
			Value: result.Page.Value, NextCursor: result.Page.NextCursor,
		},
	}
}

type workflowPayloadResponse struct {
	Section    string `json:"section"`
	Encoding   string `json:"encoding"`
	Value      string `json:"value"`
	NextCursor string `json:"nextCursor,omitempty"`
}

func workflowStepResponseFromResult(result runruntime.CallbackResult) workflowStepResponse {
	return workflowStepResponse{
		RunID: result.RunID, StepID: result.CallbackID,
		Status: result.Status, RunStatus: result.RunStatus,
		Input: result.Input, OutputContract: result.OutputContract, WorkItemsContract: result.WorkItemsContract,
		InputRef: workflowPayloadRefFromRuntime(result.InputRef), OutputContractRef: workflowPayloadRefFromRuntime(result.OutputContractRef), WorkItemsRef: workflowPayloadRefFromRuntime(result.WorkItemsRef),
		DeliveryScope: result.DeliveryScope, WorkScope: result.WorkScope, RelatedScopes: result.RelatedScopes,
		Replayed: result.Replayed, TeamMemberContext: result.TeamMemberContext,
	}
}

func workflowPayloadRefFromRuntime(value *workflowpayload.Ref) *workflowPayloadRef {
	if value == nil {
		return nil
	}
	return &workflowPayloadRef{Cursor: value.Cursor}
}

func normalizeWorkflowStepRequest(workspaceID, agentID string, request workflowStepRequest) (runruntime.CallbackCommand, error) {
	return normalizeWorkflowStepRequestWithMember(workspaceID, agentID, "", request)
}

func normalizeWorkflowStepRequestWithMember(workspaceID, agentID, memberID string, request workflowStepRequest) (runruntime.CallbackCommand, error) {
	publicAction := strings.TrimSpace(request.Action)
	switch publicAction {
	case "get", "complete", "fail", "pause", "park":
	default:
		return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "action must be list, get, complete, fail, pause, or park")
	}
	command := runruntime.CallbackCommand{
		WorkspaceID: strings.TrimSpace(workspaceID), AgentID: strings.TrimSpace(agentID), MemberID: strings.TrimSpace(memberID),
		RunID: strings.TrimSpace(request.WorkflowRunID), CallbackID: strings.TrimSpace(request.WorkflowStepID),
		Action: workflowStepRuntimeAction(publicAction), Section: strings.TrimSpace(request.Section), Path: strings.TrimSpace(request.Path), Cursor: strings.TrimSpace(request.Cursor), Limit: request.Limit, Output: append(json.RawMessage(nil), request.Output...),
		WorkItems:   append(json.RawMessage(nil), request.WorkItems...),
		FailureCode: strings.TrimSpace(request.Code), FailureMessage: strings.TrimSpace(request.Message),
		IdempotencyKey: strings.TrimSpace(request.IdempotencyKey),
	}
	workScopeJSON := bytes.TrimSpace(request.WorkScope)
	if command.WorkspaceID == "" || command.AgentID == "" || command.RunID == "" || command.CallbackID == "" {
		return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "trusted workflow Run context is incomplete")
	}
	switch command.Action {
	case runruntime.CallbackActionGetCurrent:
		if len(bytes.TrimSpace(command.Output)) != 0 || len(workScopeJSON) != 0 || len(bytes.TrimSpace(command.WorkItems)) != 0 || command.FailureCode != "" || command.FailureMessage != "" || command.IdempotencyKey != "" {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "get does not accept output, workScope, workItems, failure, or idempotencyKey")
		}
		if command.Section != "" && command.Section != "compact" && command.Section != "input" && command.Section != "output_contract" && command.Section != "work_items_contract" {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "get section is invalid")
		}
		if command.Section == "" && (command.Path != "" || command.Cursor != "" || command.Limit != 0) {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "path, cursor, and limit require a payload section")
		}
		if command.Section == "compact" && (command.Path != "" || command.Cursor != "" || command.Limit != 0) {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "compact get does not accept path, cursor, or limit")
		}
	case runruntime.CallbackActionPause:
		if command.Section != "" || command.Path != "" || command.Cursor != "" || command.Limit != 0 || len(bytes.TrimSpace(command.Output)) != 0 || len(workScopeJSON) != 0 || len(bytes.TrimSpace(command.WorkItems)) != 0 || command.FailureCode != "" || command.FailureMessage != "" {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "pause does not accept payload read fields")
		}
		if command.IdempotencyKey == "" || len(command.IdempotencyKey) > workflowRunCallbackIdempotencyMaxLen {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "pause requires a bounded idempotencyKey")
		}
	case runruntime.CallbackActionPark:
		if command.Section != "" || command.Path != "" || command.Cursor != "" || command.Limit != 0 || len(bytes.TrimSpace(command.Output)) != 0 || len(workScopeJSON) != 0 || len(bytes.TrimSpace(command.WorkItems)) != 0 || command.FailureCode != "" || command.FailureMessage != "" {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "park does not accept payload read fields")
		}
		// Park is a repeatable recovery-FSM cycle. A legacy key remains
		// accepted and bounded, but it is not required and is never a durable
		// replay identity.
		if command.IdempotencyKey != "" && len(command.IdempotencyKey) > workflowRunCallbackIdempotencyMaxLen {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "park idempotencyKey must be bounded when provided")
		}
	case runruntime.CallbackActionComplete:
		if command.Section != "" || command.Path != "" || command.Cursor != "" || command.Limit != 0 {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete does not accept payload read fields")
		}
		if command.IdempotencyKey == "" || len(command.IdempotencyKey) > workflowRunCallbackIdempotencyMaxLen {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete requires a bounded idempotencyKey")
		}
		if len(bytes.TrimSpace(command.Output)) == 0 {
			command.Output = json.RawMessage(`{}`)
		}
		if !json.Valid(command.Output) {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete requires valid JSON output")
		}
		if workItemsJSON := bytes.TrimSpace(command.WorkItems); len(workItemsJSON) != 0 {
			if !json.Valid(workItemsJSON) || workItemsJSON[0] != '[' {
				return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete workItems must be a JSON array")
			}
		}
		var object map[string]json.RawMessage
		if err := json.Unmarshal(command.Output, &object); err != nil || object == nil {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete output must be a JSON object")
		}
		if command.FailureCode != "" || command.FailureMessage != "" {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete does not accept failure fields")
		}
		if len(workScopeJSON) != 0 {
			var candidate any
			if err := json.Unmarshal(workScopeJSON, &candidate); err != nil {
				return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete workScope must be canonical")
			}
			scope, err := domain.NormalizeMessageScopeValue(candidate)
			if err != nil {
				return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "complete workScope must be canonical")
			}
			command.WorkScope = &scope
		}
	case runruntime.CallbackActionFail:
		if command.Section != "" || command.Path != "" || command.Cursor != "" || command.Limit != 0 {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "fail does not accept payload read fields")
		}
		if command.IdempotencyKey == "" || len(command.IdempotencyKey) > workflowRunCallbackIdempotencyMaxLen {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "fail requires a bounded idempotencyKey")
		}
		if command.FailureCode == "" {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "fail requires failureCode")
		}
		if len(bytes.TrimSpace(command.Output)) != 0 || len(workScopeJSON) != 0 || len(bytes.TrimSpace(command.WorkItems)) != 0 {
			return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "fail does not accept output, workScope, or workItems")
		}
	default:
		return runruntime.CallbackCommand{}, publicError(http.StatusBadRequest, "invalid_request_body", "action must be list, get, complete, fail, pause, or park")
	}
	return command, nil
}

func workflowStepRuntimeAction(action string) string {
	if action == "get" {
		return runruntime.CallbackActionGetCurrent
	}
	return action
}

func workflowRunCallbackJSONError(err error) error {
	var maxBytesError *http.MaxBytesError
	if errors.As(err, &maxBytesError) {
		return publicError(http.StatusRequestEntityTooLarge, "invalid_request_body", "invalid request body")
	}
	return publicError(http.StatusBadRequest, "invalid_request_body", "invalid request body")
}

func (handler *WorkflowStepHandler) classifyError(err error) (api.ErrorResponse, bool) {
	stage := string(publicerr.StageRun)
	// Typed callback validation failures project their stable code and typed
	// diagnostics through the canonical Workflow classifier instead of
	// collapsing to a single fixed step_output_invalid message. They carry a
	// storeerr.ErrBadRequest kind so they classify as 400; the classifier keeps
	// the code and projects the typed diagnostics from the details payload.
	var validation *runruntime.CallbackValidationError
	if errors.As(err, &validation) {
		classified, ok := publicerr.Classify(err, publicerr.Options{
			Stage:       publicerr.StageRun,
			DefaultCode: "step_output_invalid",
			Unredacted:  true,
		})
		if ok {
			// Carry every canonical contract field (stage, path, hint,
			// retryability, diagnostics) so the typed callback validation
			// failure survives the classifier→ErrorResponse→writer path.
			return workflowErrorResponse(classified, nil), true
		}
	}
	switch {
	case errors.Is(err, runruntime.ErrCallbackNotFound):
		return api.ErrorResponse{
			Status: http.StatusNotFound, Message: "Workflow Step was not found",
			Code: "step_not_found", Stage: stage,
		}, true
	case errors.Is(err, runruntime.ErrCallbackUnauthorized):
		return api.ErrorResponse{
			Status: http.StatusForbidden, Message: "Workflow Step action is forbidden",
			Code: "step_forbidden", Stage: stage,
		}, true
	case errors.Is(err, runruntime.ErrCallbackOutputInvalid):
		return api.ErrorResponse{Status: http.StatusBadRequest, Message: "Workflow Step output is invalid", Code: "step_output_invalid", Stage: stage}, true
	case errors.Is(err, runruntime.ErrCallbackNotReady):
		// The delivery race (step_not_ready) is always retryable at the top
		// level, including after the caller's bounded retry window is
		// exhausted: the callback may still converge, so the Agent must be
		// able to re-issue the exact same callback command.
		return api.ErrorResponse{
			Status: http.StatusConflict, Message: "Workflow Step delivery is not ready", Code: workflowRunCallbackNotReadyCode,
			Retryable: true, Stage: stage,
		}, true
	case errors.Is(err, runruntime.ErrCallbackNotWaiting), errors.Is(err, runruntime.ErrRunTerminal):
		return api.ErrorResponse{Status: http.StatusConflict, Message: "Workflow Step is no longer waiting", Code: "step_not_waiting", Stage: stage}, true
	case errors.Is(err, runruntime.ErrCallbackAlreadyResolved):
		return api.ErrorResponse{Status: http.StatusConflict, Message: "Workflow Step was already resolved with a different result", Code: workflowRunCallbackAlreadyResolvedCode, Stage: stage}, true
	case errors.Is(err, runruntime.ErrCallbackReplayConflict):
		return api.ErrorResponse{Status: http.StatusConflict, Message: "Workflow Step replay conflicts with the committed result", Code: "idempotency_conflict", Stage: stage}, true
	case errors.Is(err, runruntime.ErrCallbackPayloadChanged):
		return api.ErrorResponse{Status: http.StatusConflict, Message: "Workflow Step payload snapshot changed; restart the read", Code: "step_payload_conflict", Stage: stage}, true
	case errors.Is(err, runartifact.ErrUnauthorizedReference), errors.Is(err, runartifact.ErrInvalidReference):
		return api.ErrorResponse{Status: http.StatusForbidden, Message: "Workflow Step artifact reference is not authorized", Code: "step_artifact_forbidden", Stage: stage}, true
	case errors.Is(err, runartifact.ErrStaleReference), errors.Is(err, runartifact.ErrInvalidCursor), errors.Is(err, runartifact.ErrInvalidRead):
		return api.ErrorResponse{Status: http.StatusConflict, Message: "Workflow Step artifact reference or cursor is stale", Code: "step_artifact_conflict", Stage: stage}, true
	default:
		// Route canonical API-layer and store errors through the centralized
		// public classifier so strict-JSON, precondition, and runtime failures
		// project the same canonical contract. Preserve the distinction between
		// missing and forbidden instead of folding authorization failures into
		// not-found responses.
		if classified, ok := publicerr.Classify(err, publicerr.Options{
			Stage: publicerr.StageRun, DefaultCode: "run_request_failed",
			Unredacted: true,
		}); ok {
			return workflowErrorResponse(classified, nil), true
		}
		if handler != nil && handler.logger != nil {
			observed, internal := api.ObservedErrorWithCause(err)
			attrs := []any{slog.Any("error", observed)}
			if internal != nil {
				attrs = append(attrs, slog.Any("internal_cause", internal))
			}
			handler.logger.Warn("Workflow Step failed", attrs...)
		}
		return api.ErrorResponse{
			Status: http.StatusInternalServerError, Message: err.Error(), Code: "internal_error", Stage: stage,
		}, true
	}
}