package exec

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

type workflowRemote struct {
	client            *apiclient.Client
	receipt           workflowReceiptTransport
	callbackRetryWait func(context.Context, int, time.Time) error
}

func (remote workflowRemote) receiptTransport() workflowReceiptTransport {
	transport := remote.receipt
	transport.client = remote.client
	return transport
}

const (
	workflowRunCallbackMaxAttempts    = 64
	workflowRunCallbackRetryWindow    = 45 * time.Second
	workflowRunCallbackInitialBackoff = 100 * time.Millisecond
	workflowRunCallbackMaxBackoff     = 2 * time.Second
)

var errWorkflowRunCallbackRetryWindowElapsed = errors.New("workflow run callback retry window elapsed")

type workflowCLIHiddenContext struct {
	AgentID     string `json:"agentId"`
	WorkspaceID string `json:"workspaceId"`
	CommandID   string `json:"commandId"`
}

type workflowCLIRequest[I any] struct {
	Context workflowCLIHiddenContext `json:"context"`
	Input   I                        `json:"input"`
}

type workflowTransportResponse[O any] struct {
	Value *O
}

func (response *workflowTransportResponse[O]) UnmarshalJSON(data []byte) error {
	value, err := decodeWorkflowOperationResult[O](data)
	if err != nil {
		return err
	}
	response.Value = value
	return nil
}

type workflowStepTransportRequest struct {
	Action         string                     `json:"action"`
	Section        string                     `json:"section,omitempty"`
	Path           string                     `json:"path,omitempty"`
	Output         json.RawMessage            `json:"output,omitempty"`
	WorkScope      *workflowMessageScopeInput `json:"workScope,omitempty"`
	WorkItems      *[]workflowWorkItemInput   `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 workflowStepListTransportResponse struct {
	Steps      []workflowStepSummary `json:"steps"`
	NextCursor string                `json:"nextCursor,omitempty"`
}

func (a workflowRemote) plan(ctx context.Context, input workflowPlanRequest) (workflowPlanResult, error) {
	switch input.Action {
	case "create":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanCreateResponse](ctx, a, "plan", input)
		return workflowPlanResult{create: value}, err
	case "list":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanListResponse](ctx, a, "plan", input)
		return workflowPlanResult{list: value}, err
	case "archive":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanResponse](ctx, a, "plan", input)
		return workflowPlanResult{plan: value}, err
	case "show":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanAgentViewResponse](ctx, a, "plan", input)
		return workflowPlanResult{agentView: value}, err
	case "version":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanVersionResponse](ctx, a, "plan", input)
		return workflowPlanResult{version: value}, err
	case "validate":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanValidationResponse](ctx, a, "plan", input)
		return workflowPlanResult{validation: value}, err
	case "update":
		if input.TeamID != "" {
			value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanPublishResponse](ctx, a, "plan", input)
			return workflowPlanResult{publish: value}, err
		}
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanDraftResponse](ctx, a, "plan", input)
		return workflowPlanResult{draft: value}, err
	case "publish":
		value, err := callAgentWorkflowObject[workflowPlanRequest, workflowPlanPublishResponse](ctx, a, "plan", input)
		return workflowPlanResult{publish: value}, err
	default:
		return workflowPlanResult{}, NewError(CodeValidation, "unsupported workflow plan operation")
	}
}

func (a workflowRemote) trigger(ctx context.Context, input workflowTriggerRequest) (workflowTriggerResult, error) {
	switch input.Action {
	case "create", "update", "enable":
		value, err := callAgentWorkflowObject[workflowTriggerRequest, workflowTriggerMutationResponse](ctx, a, "trigger", input)
		return workflowTriggerResult{mutation: value}, err
	case "list":
		value, err := callAgentWorkflowObject[workflowTriggerRequest, workflowTriggerListResponse](ctx, a, "trigger", input)
		return workflowTriggerResult{list: value}, err
	case "show":
		value, err := callAgentWorkflowObject[workflowTriggerRequest, workflowTriggerShowResponse](ctx, a, "trigger", input)
		return workflowTriggerResult{show: value}, err
	case "disable", "archive":
		value, err := callAgentWorkflowObject[workflowTriggerRequest, workflowTriggerResponse](ctx, a, "trigger", input)
		return workflowTriggerResult{trigger: value}, err
	case "validate":
		value, err := callAgentWorkflowObject[workflowTriggerRequest, workflowTriggerValidationResponse](ctx, a, "trigger", input)
		return workflowTriggerResult{validation: value}, err
	default:
		return workflowTriggerResult{}, NewError(CodeValidation, "unsupported workflow trigger operation")
	}
}

func (a workflowRemote) run(ctx context.Context, input workflowRunRequest) (workflowRunResult, error) {
	switch input.Action {
	case "start":
		value, err := callAgentWorkflowObject[workflowRunRequest, workflowRunStartResponse](ctx, a, "run", input)
		return workflowRunResult{start: value}, err
	case "list":
		value, err := callAgentWorkflowObject[workflowRunRequest, workflowRunListResponse](ctx, a, "run", input)
		return workflowRunResult{list: value}, err
	case "show":
		value, err := callAgentWorkflowObject[workflowRunRequest, workflowRunShowResponse](ctx, a, "run", input)
		return workflowRunResult{show: value}, err
	case "logs":
		if input.Section == "details" {
			value, err := callAgentWorkflowObject[workflowRunRequest, workflowRunLogsDetailsResponse](ctx, a, "run", input)
			return workflowRunResult{logDetails: value}, err
		}
		value, err := a.runLogs(ctx, input)
		return workflowRunResult{logs: value}, err
	case "cancel":
		value, err := callAgentWorkflowObject[workflowRunRequest, workflowRunCancelResponse](ctx, a, "run", input)
		return workflowRunResult{cancel: value}, err
	default:
		return workflowRunResult{}, NewError(CodeValidation, "unsupported workflow run operation")
	}
}

func (a workflowRemote) runLogs(ctx context.Context, input workflowRunRequest) (*workflowRunLogsResponse, error) {
	maxItems := input.Limit
	if input.Tail > 0 {
		maxItems = input.Tail
	}
	value, err := callAgentWorkflowObject[workflowRunRequest, workflowRunLogsResponse](ctx, a, "run", input)
	if err != nil || !input.Follow {
		return value, err
	}
	result := &workflowRunLogsResponse{}
	for {
		result.Items = append(result.Items, value.Items...)
		result.NextCursor = value.NextCursor
		result.RunStatus = value.RunStatus
		result.Terminal = value.Terminal
		if value.Terminal || value.NextCursor == "" || (maxItems > 0 && len(result.Items) >= maxItems) {
			if maxItems > 0 && len(result.Items) > maxItems {
				result.Items = result.Items[:maxItems]
			}
			return result, nil
		}
		input.Cursor = value.NextCursor
		input.Tail = 0
		if maxItems > 0 {
			input.Limit = maxItems - len(result.Items)
		}
		value, err = callAgentWorkflowObject[workflowRunRequest, workflowRunLogsResponse](ctx, a, "run", input)
		if err != nil {
			return nil, err
		}
	}
}

func (a workflowRemote) step(
	ctx context.Context,
	input workflowStepInput,
	runCtx workflowRunContext,
) (*workflowStepView, *workflowStepListResponse, error) {
	if err := requireRuntimeClient(a.client); err != nil {
		return nil, nil, &Error{Code: CodePermissionDenied, Problem: "workflow run callbacks require an authenticated workspace-scoped Agent session", Cause: err}
	}
	if strings.TrimSpace(a.client.LocalAgentID) == "" {
		return nil, nil, &Error{Code: CodePermissionDenied, Problem: "workflow run callbacks must be invoked from an authenticated Agent session"}
	}
	runCtx = runCtx.normalize()
	if (input.Action == "list" && !runCtx.validForList()) || (input.Action != "list" && !runCtx.valid()) {
		return nil, nil, &Error{Code: CodeValidation, Problem: "workflow run callback selector is incomplete"}
	}
	if runCtx.WorkspaceID != strings.TrimSpace(a.client.RouteWorkspaceID) {
		return nil, nil, &Error{Code: CodePermissionDenied, Problem: "workflow run workspace does not match the Agent session"}
	}
	if runCtx.AgentID != strings.TrimSpace(a.client.AgentID) {
		return nil, nil, &Error{Code: CodePermissionDenied, Problem: "workflow run Agent does not match the Agent session"}
	}
	body := workflowStepTransportRequest{
		Action: input.Action, Section: input.Section, Path: input.Path, Output: input.Output, WorkScope: input.WorkScope, WorkItems: input.WorkItems,
		Code: input.Code, Message: input.Message, IdempotencyKey: input.IdempotencyKey,
		Cursor: input.Cursor, Limit: input.Limit,
		WorkflowRunID: runCtx.RunID, WorkflowStepID: runCtx.EffectID,
	}
	callbackPath := a.client.AgentRuntimePath("/workflow/step")
	stepMutation := workflowOperationIsMutation("step", input.Action)
	stepIdempotent := stepMutation && strings.TrimSpace(input.IdempotencyKey) != ""
	stepHeaders := map[string]string{}
	if stepIdempotent {
		stepHeaders["Idempotency-Key"] = strings.TrimSpace(input.IdempotencyKey)
	}
	postStep := func(result any) error {
		if stepIdempotent {
			return a.receiptTransport().post(ctx, callbackPath, body, result, stepHeaders)
		}
		return a.receiptTransport().postOnce(ctx, callbackPath, body, result, nil)
	}
	retryDeadline := time.Now().Add(workflowRunCallbackRetryWindow)
	retryWait := a.callbackRetryWait
	if retryWait == nil {
		retryWait = waitWorkflowRunCallbackRetry
	}
	for attempt := 1; attempt <= workflowRunCallbackMaxAttempts; attempt++ {
		var err error
		if input.Action == "list" {
			var response workflowTransportResponse[workflowStepListTransportResponse]
			err = postStep(&response)
			if err == nil {
				steps := []workflowStepSummary{}
				if response.Value != nil && response.Value.Steps != nil {
					steps = response.Value.Steps
				}
				nextCursor := ""
				if response.Value != nil {
					nextCursor = response.Value.NextCursor
				}
				return nil, &workflowStepListResponse{Steps: steps, NextCursor: nextCursor}, nil
			}
		} else {
			var response workflowTransportResponse[workflowStepView]
			err = postStep(&response)
			if err == nil {
				return response.Value, nil, nil
			}
		}
		if err == nil {
			return nil, nil, fmt.Errorf("workflow step response is empty")
		}
		if !isWorkflowRunCallbackNotReadyError(err) || attempt == workflowRunCallbackMaxAttempts {
			return nil, nil, err
		}
		if waitErr := retryWait(ctx, attempt, retryDeadline); waitErr != nil {
			if errors.Is(waitErr, errWorkflowRunCallbackRetryWindowElapsed) {
				return nil, nil, err
			}
			// The outer transport deadline can fire before the retry window
			// elapses (for example the CLI's bounded HTTP request timeout).
			// Retry exhaustion must preserve the canonical delivery-race code
			// and top-level retryable so the caller can re-issue the exact same
			// callback command, instead of failing closed to an opaque internal
			// error that hides the retryable race.
			return nil, nil, err
		}
	}
	return nil, nil, fmt.Errorf("workflow run callback retry budget exhausted")
}

func waitWorkflowRunCallbackRetry(ctx context.Context, attempt int, deadline time.Time) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	remaining := time.Until(deadline)
	if remaining <= 0 {
		return errWorkflowRunCallbackRetryWindowElapsed
	}
	backoff := workflowRunCallbackInitialBackoff
	for step := 1; step < attempt && backoff < workflowRunCallbackMaxBackoff; step++ {
		backoff *= 2
		if backoff > workflowRunCallbackMaxBackoff {
			backoff = workflowRunCallbackMaxBackoff
		}
	}
	if backoff > remaining {
		backoff = remaining
	}
	timer := time.NewTimer(backoff)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		if !time.Now().Before(deadline) {
			return errWorkflowRunCallbackRetryWindowElapsed
		}
		return nil
	}
}

func callAgentWorkflowObject[I, O any](ctx context.Context, api workflowRemote, object string, input I) (*O, error) {
	if err := requireRuntimeClient(api.client); err != nil {
		return nil, &Error{Code: CodePermissionDenied, Problem: "Workflow commands require an authenticated workspace-scoped Agent session", Cause: err}
	}
	agentID := strings.TrimSpace(api.client.AgentID)
	workspaceID := strings.TrimSpace(api.client.RouteWorkspaceID)
	if strings.TrimSpace(api.client.LocalAgentID) == "" {
		return nil, &Error{Code: CodePermissionDenied, Problem: "Workflow commands must be invoked from an authenticated Agent session"}
	}
	if agentID == "" {
		return nil, &Error{Code: CodePermissionDenied, Problem: "Workflow commands require a routed Agent identity"}
	}
	if workspaceID == "" {
		return nil, &Error{Code: CodePermissionDenied, Problem: "Workflow commands require a workspace-scoped Agent session"}
	}
	commandID, err := newWorkflowCLICommandID()
	if err != nil {
		return nil, &Error{Code: CodeInternal, Problem: "cannot allocate Workflow command identity", Cause: err}
	}
	hidden := workflowCLIHiddenContext{AgentID: agentID, WorkspaceID: workspaceID, CommandID: commandID}
	operation := workflowOperation(input)
	mutation := workflowOperationIsMutation(object, operation)
	headers := map[string]string{}
	if mutation {
		headers["Idempotency-Key"] = workflowCLIIdempotencyKey(agentID, workspaceID, object, operation, commandID)
	}
	var response workflowTransportResponse[O]
	path := api.client.AgentRuntimePath("/workflow/" + object)
	request := workflowCLIRequest[I]{Context: hidden, Input: input}
	var postErr error
	if mutation {
		postErr = api.receiptTransport().post(ctx, path, request, &response, headers)
	} else {
		postErr = api.receiptTransport().postOnce(ctx, path, request, &response, headers)
	}
	if postErr != nil {
		return nil, postErr
	}
	return response.Value, nil
}

func workflowOperation[I any](input I) string {
	encoded, _ := json.Marshal(input)
	var discriminator struct {
		Action string `json:"action"`
	}
	_ = json.Unmarshal(encoded, &discriminator)
	return discriminator.Action
}

func newWorkflowCLICommandID() (string, error) {
	var entropy [16]byte
	if _, err := rand.Read(entropy[:]); err != nil {
		return "", err
	}
	return "cli:" + hex.EncodeToString(entropy[:]), nil
}

func workflowCLIIdempotencyKey(agentID, workspaceID, object, operation, commandID string) string {
	digest := sha256.Sum256([]byte(strings.Join([]string{strings.TrimSpace(agentID), strings.TrimSpace(workspaceID), strings.TrimSpace(object), strings.TrimSpace(operation), strings.TrimSpace(commandID)}, "\x00")))
	return "workflow-cli:" + hex.EncodeToString(digest[:])
}