package exec

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/core/domain"
)

const (
	workflowInvocationMaxInputBytes        = 1 << 20
	workflowActionMaxJSONNesting           = 64
	workflowRunMaxWorkItems                = 256
	workflowRunMaxWorkItemKeyBytes         = 256
	workflowRunMaxRelatedScopes            = 64
	workflowRunMaxDependencyEdges          = 2048
	workflowRunMaxRelatedScopeRefs         = 512
	workflowRunMaxWorkItemsJSONValues      = 6 * 1024
	workflowRunMaxWorkItemsJSONBytes       = 1 << 20
	workflowRunMaxInstructionRefs          = 64
	workflowRunMaxInstructionSegments      = 8
	workflowRunMaxInstructionSegmentBytes  = 64
	workflowRunMaxResolvedInstructionBytes = 32 * 1024
	workflowRunMaxInstructionScalarBytes   = 8 * 1024
	workflowRunMaxPayloadChunkBytes        = 32 * 1024
	// workflowRunMaxAllowedResponders mirrors plan.MaxAllowedResponders so
	// the public CLI schema and client-side validation share the same safe
	// upper bound for a responder.selectOne candidate pool.
	workflowRunMaxAllowedResponders = 64
)

// workflowRunContext combines the authenticated Agent session with the
// model-selected Run and Step. RunID and EffectID are selectors, not
// credentials: the Server independently verifies the frozen callback
// responder and the exact awaiting effect before exposing or mutating it.
type workflowRunContext struct {
	WorkspaceID string
	AgentID     string
	RunID       string
	EffectID    string
}

func (w workflowRunContext) normalize() workflowRunContext {
	w.WorkspaceID = strings.TrimSpace(w.WorkspaceID)
	w.AgentID = strings.TrimSpace(w.AgentID)
	w.RunID = strings.TrimSpace(w.RunID)
	w.EffectID = strings.TrimSpace(w.EffectID)
	return w
}

func (w workflowRunContext) valid() bool {
	w = w.normalize()
	return w.WorkspaceID != "" && w.AgentID != "" && w.RunID != "" && w.EffectID != ""
}

func (w workflowRunContext) validForList() bool {
	w = w.normalize()
	return w.WorkspaceID != "" && w.AgentID != ""
}

// workflowStepInput is the model-authored selector and business half of a
// callback command. Workspace and Agent identity still come only from the
// authenticated Exec principal.
type workflowStepInput struct {
	Action         string
	RunID          string
	StepID         string
	Section        string
	Path           string
	Cursor         string
	Limit          int
	Output         json.RawMessage
	WorkScope      *workflowMessageScopeInput
	WorkItems      *[]workflowWorkItemInput
	Code           string
	Message        string
	IdempotencyKey string
}

func workflowStepAction(action string) bool {
	switch action {
	case "list", "get", "complete", "fail", "pause", "park":
		return true
	default:
		return false
	}
}

func validateWorkflowStep(action string, raw json.RawMessage) (workflowStepInput, error) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(raw, &fields); err != nil {
		return workflowStepInput{}, NewError(CodeValidation, "invalid workflow step arguments")
	}
	if action == "list" {
		return validateWorkflowStepList(fields)
	}
	allowed := map[string]bool{}
	callback := workflowStepInput{Action: action}
	selectors := func() error {
		if err := json.Unmarshal(fields["runId"], &callback.RunID); err != nil {
			return NewError(CodeValidation, action+" requires runId")
		}
		callback.RunID = strings.TrimSpace(callback.RunID)
		if err := json.Unmarshal(fields["stepId"], &callback.StepID); err != nil {
			return NewError(CodeValidation, action+" requires stepId")
		}
		callback.StepID = strings.TrimSpace(callback.StepID)
		return nil
	}
	switch action {
	case "get":
		allowed["runId"] = true
		allowed["stepId"] = true
		allowed["section"] = true
		allowed["path"] = true
		allowed["cursor"] = true
		allowed["limit"] = true
		if rawSection, present := fields["section"]; present {
			if err := json.Unmarshal(rawSection, &callback.Section); err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "get section must be a string")
			}
			callback.Section = strings.TrimSpace(callback.Section)
		}
		if rawPath, present := fields["path"]; present {
			if err := json.Unmarshal(rawPath, &callback.Path); err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "get path must be a string")
			}
			callback.Path = strings.TrimSpace(callback.Path)
		}
		if rawCursor, present := fields["cursor"]; present {
			if err := json.Unmarshal(rawCursor, &callback.Cursor); err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "get cursor must be a string")
			}
			callback.Cursor = strings.TrimSpace(callback.Cursor)
		}
		if rawLimit, present := fields["limit"]; present {
			if err := json.Unmarshal(rawLimit, &callback.Limit); err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "get limit must be an integer")
			}
		}
		if callback.Limit < 0 || callback.Limit > workflowRunMaxPayloadChunkBytes {
			return workflowStepInput{}, NewError(CodeValidation, "get limit must be between 1 and 3072")
		}
		if callback.Section != "" && callback.Section != "compact" && callback.Section != "input" && callback.Section != "output_contract" && callback.Section != "work_items_contract" {
			return workflowStepInput{}, NewError(CodeValidation, "get section is invalid")
		}
		if len(callback.Section) > 64 || len(callback.Path) > 512 || len(callback.Cursor) > 2048 {
			return workflowStepInput{}, NewError(CodeValidation, "get payload selector is too large")
		}
		if callback.Section == "" && (callback.Path != "" || callback.Cursor != "" || callback.Limit != 0) {
			return workflowStepInput{}, NewError(CodeValidation, "get path, cursor, and limit require section")
		}
		if callback.Section == "compact" && (callback.Path != "" || callback.Cursor != "" || callback.Limit != 0) {
			return workflowStepInput{}, NewError(CodeValidation, "compact get does not accept path, cursor, or limit")
		}
		if err := selectors(); err != nil {
			return workflowStepInput{}, err
		}
	case "pause":
		// pause is a non-terminal action: it keeps the step awaiting and durably
		// stops automatic recovery permanently (no resume trigger). Its stable
		// idempotencyKey is required so a transport retry can replay the same
		// durable action without guessing whether the first request committed.
		allowed["runId"] = true
		allowed["stepId"] = true
		allowed["idempotencyKey"] = true
		if err := selectors(); err != nil {
			return workflowStepInput{}, err
		}
		if err := decodeWorkflowStepIdempotencyKey(fields, &callback.IdempotencyKey, "pause"); err != nil {
			return workflowStepInput{}, err
		}
	case "park":
		// park is a repeatable recovery-FSM cycle. A legacy idempotencyKey is
		// accepted for compatibility, but it is optional and never becomes a
		// lifecycle-wide replay identity.
		allowed["runId"] = true
		allowed["stepId"] = true
		allowed["idempotencyKey"] = true
		if err := selectors(); err != nil {
			return workflowStepInput{}, err
		}
		if _, present := fields["idempotencyKey"]; present {
			if err := decodeWorkflowStepIdempotencyKey(fields, &callback.IdempotencyKey, "park"); err != nil {
				return workflowStepInput{}, err
			}
		}
	case "complete":
		allowed["runId"] = true
		allowed["stepId"] = true
		allowed["output"] = true
		allowed["workScope"] = true
		allowed["workItems"] = true
		allowed["idempotencyKey"] = true
		callback.Output = append(json.RawMessage(nil), fields["output"]...)
		if err := selectors(); err != nil {
			return workflowStepInput{}, err
		}
		if len(bytes.TrimSpace(callback.Output)) == 0 {
			callback.Output = json.RawMessage(`{}`)
		}
		if err := json.Unmarshal(fields["idempotencyKey"], &callback.IdempotencyKey); err != nil {
			return workflowStepInput{}, NewError(CodeValidation, "complete requires idempotencyKey")
		}
		callback.IdempotencyKey = strings.TrimSpace(callback.IdempotencyKey)
		var output map[string]json.RawMessage
		if !json.Valid(callback.Output) || json.Unmarshal(callback.Output, &output) != nil || output == nil {
			return workflowStepInput{}, NewError(CodeValidation, "complete output must be a JSON object")
		}
		if rawScope, present := fields["workScope"]; present {
			scope, err := decodeWorkflowMessageScope(rawScope)
			if err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "complete workScope must be a canonical scope object")
			}
			callback.WorkScope = &scope
		}
		if rawWorkItems, present := fields["workItems"]; present {
			workItems, err := decodeWorkflowWorkItems(rawWorkItems)
			if err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "complete workItems must be a closed array of canonical WorkItems")
			}
			callback.WorkItems = &workItems
		}
		if callback.IdempotencyKey == "" || len(callback.IdempotencyKey) > 128 {
			return workflowStepInput{}, NewError(CodeValidation, "complete requires an idempotencyKey of at most 128 characters")
		}
	case "fail":
		allowed["runId"] = true
		allowed["stepId"] = true
		allowed["code"] = true
		allowed["message"] = true
		allowed["idempotencyKey"] = true
		var args struct {
			Code           string `json:"code"`
			Message        string `json:"message"`
			IdempotencyKey string `json:"idempotencyKey"`
		}
		if err := json.Unmarshal(raw, &args); err != nil {
			return workflowStepInput{}, NewError(CodeValidation, "invalid workflow step fail arguments")
		}
		callback.Code = strings.TrimSpace(args.Code)
		callback.Message = strings.TrimSpace(args.Message)
		callback.IdempotencyKey = strings.TrimSpace(args.IdempotencyKey)
		if err := selectors(); err != nil {
			return workflowStepInput{}, err
		}
		if callback.Code == "" {
			return workflowStepInput{}, NewError(CodeValidation, "fail requires code")
		}
		if callback.IdempotencyKey == "" || len(callback.IdempotencyKey) > 128 {
			return workflowStepInput{}, NewError(CodeValidation, "fail requires an idempotencyKey of at most 128 characters")
		}
	default:
		return workflowStepInput{}, NewError(CodeValidation, "action must be list, get, complete, fail, pause, or park")
	}
	if action != "list" && callback.RunID == "" {
		return workflowStepInput{}, NewError(CodeValidation, action+" requires runId")
	}
	if action != "list" && callback.StepID == "" {
		return workflowStepInput{}, NewError(CodeValidation, action+" requires stepId")
	}
	for field := range fields {
		if !allowed[field] {
			return workflowStepInput{}, NewError(CodeValidation, action+" does not accept "+field)
		}
	}
	return callback, nil
}

func decodeWorkflowStepIdempotencyKey(fields map[string]json.RawMessage, destination *string, action string) error {
	if destination == nil {
		return NewError(CodeValidation, action+" requires an idempotencyKey")
	}
	if err := json.Unmarshal(fields["idempotencyKey"], destination); err != nil {
		return NewError(CodeValidation, action+" requires an idempotencyKey")
	}
	*destination = strings.TrimSpace(*destination)
	if *destination == "" || len(*destination) > 128 {
		return NewError(CodeValidation, action+" requires an idempotencyKey of at most 128 characters")
	}
	return nil
}

func validateWorkflowStepList(fields map[string]json.RawMessage) (workflowStepInput, error) {
	callback := workflowStepInput{Action: "list", Limit: 50}
	for field, raw := range fields {
		switch field {
		case "cursor":
			if err := json.Unmarshal(raw, &callback.Cursor); err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "list cursor must be a string")
			}
			callback.Cursor = strings.TrimSpace(callback.Cursor)
		case "limit":
			if err := json.Unmarshal(raw, &callback.Limit); err != nil {
				return workflowStepInput{}, NewError(CodeValidation, "list limit must be an integer")
			}
		default:
			return workflowStepInput{}, NewError(CodeValidation, "list does not accept "+field)
		}
	}
	if callback.Limit < 1 || callback.Limit > 100 {
		return workflowStepInput{}, NewError(CodeValidation, "list limit must be between 1 and 100")
	}
	return callback, nil
}

func decodeWorkflowWorkItems(raw json.RawMessage) ([]workflowWorkItemInput, error) {
	if len(raw) > workflowRunMaxWorkItemsJSONBytes {
		return nil, fmt.Errorf("workItems is too large")
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	var encodedItems []json.RawMessage
	if err := decoder.Decode(&encodedItems); err != nil || encodedItems == nil {
		return nil, fmt.Errorf("invalid workItems")
	}
	if err := decoder.Decode(new(struct{})); err != io.EOF {
		if err == nil {
			return nil, fmt.Errorf("trailing workItems JSON")
		}
		return nil, err
	}
	if len(encodedItems) > workflowRunMaxWorkItems {
		return nil, fmt.Errorf("too many workItems")
	}
	workItems := make([]workflowWorkItemInput, len(encodedItems))
	dependencyEdges := 0
	relatedScopeReferences := 0
	structuralValues := 1
	for index, encoded := range encodedItems {
		var fields map[string]json.RawMessage
		if err := json.Unmarshal(encoded, &fields); err != nil || fields == nil {
			return nil, fmt.Errorf("workItems[%d] must be an object", index)
		}
		for _, required := range []string{"key", "workScope", "relatedScopes", "responder", "instruction"} {
			if _, present := fields[required]; !present {
				return nil, fmt.Errorf("workItems[%d] is missing %s", index, required)
			}
		}
		if raw := bytes.TrimSpace(fields["relatedScopes"]); len(raw) == 0 || raw[0] != '[' {
			return nil, fmt.Errorf("workItems[%d].relatedScopes must be an array", index)
		}
		if raw, present := fields["dependsOn"]; present {
			raw = bytes.TrimSpace(raw)
			if len(raw) == 0 || raw[0] != '[' {
				return nil, fmt.Errorf("workItems[%d].dependsOn must be an array", index)
			}
		}
		itemDecoder := json.NewDecoder(bytes.NewReader(encoded))
		itemDecoder.DisallowUnknownFields()
		item := &workItems[index]
		if err := itemDecoder.Decode(item); err != nil {
			return nil, err
		}
		if err := itemDecoder.Decode(new(struct{})); err != io.EOF {
			if err == nil {
				return nil, fmt.Errorf("trailing workItems[%d] JSON", index)
			}
			return nil, err
		}
		if !validWorkflowWorkItemKey(item.Key) ||
			!validWorkflowWorkItemResponder(item.Responder) ||
			strings.TrimSpace(item.Instruction) == "" {
			return nil, fmt.Errorf("workItems[%d] contains non-canonical text", index)
		}
		for dependencyIndex, dependency := range item.DependsOn {
			if !validWorkflowWorkItemKey(dependency) {
				return nil, fmt.Errorf("workItems[%d].dependsOn[%d] is not canonical", index, dependencyIndex)
			}
		}
		if err := validateWorkflowWorkItemInstructionReferences(item.Instruction, item.DependsOn); err != nil {
			return nil, fmt.Errorf("workItems[%d].instruction: %w", index, err)
		}
		scope, err := decodeWorkflowMessageScopeValue(item.WorkScope)
		if err != nil {
			return nil, fmt.Errorf("workItems[%d].workScope: %w", index, err)
		}
		item.WorkScope = scope
		for relatedIndex := range item.RelatedScopes {
			scope, err := decodeWorkflowMessageScopeValue(item.RelatedScopes[relatedIndex])
			if err != nil {
				return nil, fmt.Errorf("workItems[%d].relatedScopes[%d]: %w", index, relatedIndex, err)
			}
			item.RelatedScopes[relatedIndex] = scope
		}
		if len(item.RelatedScopes) > workflowRunMaxRelatedScopes {
			return nil, fmt.Errorf("workItems[%d].relatedScopes is too large", index)
		}
		dependencyEdges += len(item.DependsOn)
		if dependencyEdges > workflowRunMaxDependencyEdges {
			return nil, fmt.Errorf("workItems has too many dependency edges")
		}
		relatedScopeReferences += len(item.RelatedScopes)
		if relatedScopeReferences > workflowRunMaxRelatedScopeRefs {
			return nil, fmt.Errorf("workItems has too many related scope references")
		}
		structuralValues += 9 + len(item.DependsOn) + 3*len(item.RelatedScopes)
		if structuralValues > workflowRunMaxWorkItemsJSONValues {
			return nil, fmt.Errorf("workItems has too many deterministic JSON values")
		}
	}
	return workItems, nil
}

func decodeWorkflowMessageScope(raw json.RawMessage) (workflowMessageScopeInput, error) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(raw, &fields); err != nil || len(fields) != 2 {
		return workflowMessageScopeInput{}, fmt.Errorf("invalid scope")
	}
	for field := range fields {
		if field != "kind" && field != "scopeId" {
			return workflowMessageScopeInput{}, fmt.Errorf("invalid scope field")
		}
	}
	var scope workflowMessageScopeInput
	if err := json.Unmarshal(fields["kind"], &scope.Kind); err != nil {
		return workflowMessageScopeInput{}, err
	}
	if err := json.Unmarshal(fields["scopeId"], &scope.ScopeID); err != nil {
		return workflowMessageScopeInput{}, err
	}
	return decodeWorkflowMessageScopeValue(scope)
}

func decodeWorkflowMessageScopeValue(scope workflowMessageScopeInput) (workflowMessageScopeInput, error) {
	normalized, err := domain.NormalizeMessageScopeValue(domain.MessageScopeValue{
		Kind:    domain.MessageScopeKind(scope.Kind),
		ScopeID: scope.ScopeID,
	})
	if err != nil || string(normalized.Kind) != scope.Kind || normalized.ScopeID != scope.ScopeID {
		return workflowMessageScopeInput{}, fmt.Errorf("scope is not canonical")
	}
	return scope, nil
}

// decodeWorkflowAction rejects duplicate fields and excessive nesting before
// splitting the action discriminator from the exact server payload. Descriptor
// schema validation closes every branch at the route boundary; this decoder
// preserves the same structural boundary for direct handler tests as well.
func decodeWorkflowAction(raw json.RawMessage) (string, json.RawMessage, error) {
	if len(bytes.TrimSpace(raw)) == 0 {
		raw = json.RawMessage(`{}`)
	}
	if err := validateWorkflowActionJSON(raw); err != nil {
		return "", nil, err
	}
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(raw, &fields); err != nil {
		return "", nil, err
	}
	var action string
	if err := json.Unmarshal(fields["action"], &action); err != nil || strings.TrimSpace(action) == "" || action != strings.TrimSpace(action) {
		return "", nil, fmt.Errorf("action is required")
	}
	delete(fields, "action")
	payload, err := json.Marshal(fields)
	return action, payload, err
}