package run

import (
	"encoding/json"
	"fmt"
	"math/big"
	"reflect"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/workflow/plan"
)

// LoopValues and JoinValues expose only the statically selected scope values
// referenced by compiled load instructions.
type LoopValues struct {
	State     Values
	Iteration int
	IsLast    bool
}

type JoinValues struct {
	Branches     map[string]Values
	WinnerBranch string
	WinnerValue  Values
}

// EvalContext contains runtime value spaces.  It intentionally has no source
// expression strings or dotted synthetic namespace.
type EvalContext struct {
	PlanInput        Values
	ActivationInput  Values
	ActivationOutput Values
	RunState         Values
	Loops            map[string]LoopValues
	Joins            map[string]JoinValues
	FanOutItems      map[string]Values
	JoinItems        map[string][]any
}

// EvalExpression evaluates one compiler-validated AST.  The evaluator still
// checks dynamic types so corrupt state and stale callbacks fail closed.
func EvalExpression(expression plan.ArtifactExpression, context EvalContext) (any, error) {
	evaluator := expressionEvaluator{cloneLoads: true}
	return evaluator.evaluate(expression, context)
}

// evalExpressionBorrowed evaluates the complete expression tree without
// cloning load results. It is intentionally limited to transient fan-out
// admission: callers must treat every returned composite as immutable and
// normalize it exactly once into the durable destination after the expanded
// payload has passed its resource budget. The mode is recursive so wrapping a
// large dependency load in equality, contains, length, in, or another
// compound expression cannot silently fall back to the cloning evaluator.
func evalExpressionBorrowed(expression plan.ArtifactExpression, context EvalContext) (any, error) {
	evaluator := expressionEvaluator{cloneLoads: false}
	return evaluator.evaluate(expression, context)
}

type expressionEvaluator struct {
	cloneLoads bool
}

func (evaluator expressionEvaluator) evaluate(
	expression plan.ArtifactExpression,
	context EvalContext,
) (any, error) {
	value, err := evalExpression(expression, context, evaluator.cloneLoads)
	if err != nil {
		return nil, err
	}
	if !evaluator.cloneLoads && value != nil {
		if err := validateFanOutPayloadBudget(
			"fan_out expression "+string(expression.Op),
			value,
		); err != nil {
			return nil, err
		}
	}
	// Absence is distinct from a JSON null in contracts, but nil is the safe
	// evaluator representation for nullable equality and template rendering.
	// Binding/state consumers retain the static MayBeAbsent bit so they can
	// omit the destination instead of attempting to normalize nil as a value.
	if value == nil && expression.MayBeAbsent() {
		return nil, nil
	}
	if err := validateExpressionResult(expression.ResultType, value); err != nil {
		return nil, fmt.Errorf("%s result: %w", expression.Op, err)
	}
	return value, nil
}

func evalExpression(expression plan.ArtifactExpression, context EvalContext, cloneLoads bool) (any, error) {
	switch expression.Op {
	case plan.ArtifactOpConst:
		return decodeLiteral(expression.Value)
	case plan.ArtifactOpLoad:
		if expression.Load == nil {
			return nil, fmt.Errorf("load has no address")
		}
		return evalLoadWithClone(*expression.Load, context, cloneLoads)
	case plan.ArtifactOpArray:
		values := make([]any, len(expression.Args))
		for index, argument := range expression.Args {
			value, err := (expressionEvaluator{cloneLoads: cloneLoads}).evaluate(argument, context)
			if err != nil {
				return nil, fmt.Errorf("array item %d: %w", index, err)
			}
			values[index] = value
		}
		return values, nil
	case plan.ArtifactOpNot:
		value, err := evalBooleanArg(expression, 0, context, cloneLoads)
		return !value, err
	case plan.ArtifactOpAnd:
		left, err := evalBooleanArg(expression, 0, context, cloneLoads)
		if err != nil || !left {
			return left, err
		}
		return evalBooleanArg(expression, 1, context, cloneLoads)
	case plan.ArtifactOpOr:
		left, err := evalBooleanArg(expression, 0, context, cloneLoads)
		if err != nil || left {
			return left, err
		}
		return evalBooleanArg(expression, 1, context, cloneLoads)
	case plan.ArtifactOpEqual, plan.ArtifactOpNotEqual:
		left, right, err := evalBinary(expression, context, cloneLoads)
		if err != nil {
			return nil, err
		}
		equal, err := valuesEqual(left, right)
		if expression.Op == plan.ArtifactOpNotEqual {
			equal = !equal
		}
		return equal, err
	case plan.ArtifactOpLess, plan.ArtifactOpLessEqual, plan.ArtifactOpGreater, plan.ArtifactOpGreaterEq:
		left, right, err := evalBinary(expression, context, cloneLoads)
		if err != nil {
			return nil, err
		}
		comparison, err := compareNumbers(left, right)
		if err != nil {
			return nil, err
		}
		switch expression.Op {
		case plan.ArtifactOpLess:
			return comparison < 0, nil
		case plan.ArtifactOpLessEqual:
			return comparison <= 0, nil
		case plan.ArtifactOpGreater:
			return comparison > 0, nil
		default:
			return comparison >= 0, nil
		}
	case plan.ArtifactOpIn:
		needle, haystack, err := evalBinary(expression, context, cloneLoads)
		if err != nil {
			return nil, err
		}
		items, ok := haystack.([]any)
		if !ok {
			return nil, fmt.Errorf("in haystack is %T, want array", haystack)
		}
		return arrayContains(items, needle)
	case plan.ArtifactOpContains:
		haystack, needle, err := evalBinary(expression, context, cloneLoads)
		if err != nil {
			return nil, err
		}
		if items, ok := haystack.([]any); ok {
			return arrayContains(items, needle)
		}
		left, leftOK := expressionString(haystack)
		right, rightOK := expressionString(needle)
		if !leftOK || !rightOK {
			return nil, fmt.Errorf("contains operands are %T and %T", haystack, needle)
		}
		return strings.Contains(left, right), nil
	case plan.ArtifactOpStartsWith, plan.ArtifactOpEndsWith:
		leftValue, rightValue, err := evalBinary(expression, context, cloneLoads)
		if err != nil {
			return nil, err
		}
		left, leftOK := expressionString(leftValue)
		right, rightOK := expressionString(rightValue)
		if !leftOK || !rightOK {
			return nil, fmt.Errorf("prefix/suffix operands are %T and %T", leftValue, rightValue)
		}
		if expression.Op == plan.ArtifactOpStartsWith {
			return strings.HasPrefix(left, right), nil
		}
		return strings.HasSuffix(left, right), nil
	case plan.ArtifactOpLength:
		value, err := (expressionEvaluator{cloneLoads: cloneLoads}).evaluate(expression.Args[0], context)
		if err != nil {
			return nil, err
		}
		switch typed := value.(type) {
		case string:
			return json.Number(fmt.Sprintf("%d", len([]rune(typed)))), nil
		case json.Number:
			return json.Number(fmt.Sprintf("%d", len([]rune(typed.String())))), nil
		case []any:
			return json.Number(fmt.Sprintf("%d", len(typed))), nil
		default:
			return nil, fmt.Errorf("length operand is %T", value)
		}
	default:
		return nil, fmt.Errorf("unsupported expression opcode %q", expression.Op)
	}
}

func evalLoadWithClone(load plan.ArtifactLoad, context EvalContext, clone bool) (any, error) {
	var root any
	switch load.Space {
	case plan.ArtifactLoadPlanInput:
		root = context.PlanInput
	case plan.ArtifactLoadActivationInput:
		root = context.ActivationInput
	case plan.ArtifactLoadActivationOutput:
		root = context.ActivationOutput
	case plan.ArtifactLoadRunState:
		root = context.RunState
	case plan.ArtifactLoadLoopState:
		loop, ok := context.Loops[load.ScopeOwnerID]
		if !ok {
			return nil, fmt.Errorf("loop scope %q is unavailable", load.ScopeOwnerID)
		}
		root = loop.State
	case plan.ArtifactLoadLoopIteration, plan.ArtifactLoadLoopIsLast:
		loop, ok := context.Loops[load.ScopeOwnerID]
		if !ok {
			return nil, fmt.Errorf("loop scope %q is unavailable", load.ScopeOwnerID)
		}
		if load.Space == plan.ArtifactLoadLoopIteration {
			return json.Number(fmt.Sprintf("%d", loop.Iteration)), nil
		}
		return loop.IsLast, nil
	case plan.ArtifactLoadJoinBranch:
		join, ok := context.Joins[load.ScopeOwnerID]
		if !ok {
			return nil, fmt.Errorf("join scope %q is unavailable", load.ScopeOwnerID)
		}
		if len(load.Path) == 0 && load.Branch == "" {
			root = valuesMap(join.Branches, clone)
		} else {
			branch, ok := join.Branches[load.Branch]
			if !ok {
				if load.MayBeAbsent() {
					return nil, nil
				}
				return nil, fmt.Errorf("join branch %q is unavailable", load.Branch)
			}
			root = branch
		}
	case plan.ArtifactLoadJoinWinnerBranch, plan.ArtifactLoadJoinWinnerValue:
		join, ok := context.Joins[load.ScopeOwnerID]
		if !ok {
			return nil, fmt.Errorf("join scope %q is unavailable", load.ScopeOwnerID)
		}
		if load.Space == plan.ArtifactLoadJoinWinnerBranch {
			return join.WinnerBranch, nil
		}
		root = join.WinnerValue
	case plan.ArtifactLoadFanOutItem:
		item, ok := context.FanOutItems[load.ScopeOwnerID]
		if !ok {
			return nil, fmt.Errorf("fan-out item scope %q is unavailable", load.ScopeOwnerID)
		}
		root = item
	case plan.ArtifactLoadJoinItems:
		items, ok := context.JoinItems[load.ScopeOwnerID]
		if !ok {
			return nil, fmt.Errorf("fan-out join scope %q is unavailable", load.ScopeOwnerID)
		}
		root = items
	default:
		return nil, fmt.Errorf("unsupported load space %q", load.Space)
	}
	if len(load.Path) == 0 {
		if clone {
			return cloneValue(root), nil
		}
		return root, nil
	}
	var object map[string]any
	switch typed := root.(type) {
	case map[string]any:
		object = typed
	case Values:
		object = map[string]any(typed)
	default:
		return nil, fmt.Errorf("load %s root is %T, want object", load.Space, root)
	}
	value, ok := object[load.Path[0]]
	if !ok {
		if load.MayBeAbsent() {
			return nil, nil
		}
		return nil, fmt.Errorf("load %s field %q is missing", load.Space, load.Path[0])
	}
	if clone {
		return cloneValue(value), nil
	}
	return value, nil
}

func valuesMap(values map[string]Values, clone bool) map[string]any {
	result := make(map[string]any, len(values))
	for key, value := range values {
		if clone {
			result[key] = cloneValues(value)
		} else {
			result[key] = value
		}
	}
	return result
}

func evalBooleanArg(expression plan.ArtifactExpression, index int, context EvalContext, cloneLoads bool) (bool, error) {
	value, err := (expressionEvaluator{cloneLoads: cloneLoads}).evaluate(expression.Args[index], context)
	if err != nil {
		return false, err
	}
	boolean, ok := value.(bool)
	if !ok {
		return false, fmt.Errorf("logical operand is %T", value)
	}
	return boolean, nil
}

func evalBinary(expression plan.ArtifactExpression, context EvalContext, cloneLoads bool) (any, any, error) {
	left, err := (expressionEvaluator{cloneLoads: cloneLoads}).evaluate(expression.Args[0], context)
	if err != nil {
		return nil, nil, err
	}
	right, err := (expressionEvaluator{cloneLoads: cloneLoads}).evaluate(expression.Args[1], context)
	if err == nil && !cloneLoads {
		if budgetErr := validateFanOutPayloadBudget(
			"fan_out binary expression "+string(expression.Op),
			[]any{left, right},
		); budgetErr != nil {
			return nil, nil, budgetErr
		}
	}
	return left, right, err
}

func arrayContains(items []any, needle any) (bool, error) {
	for _, item := range items {
		equal, err := valuesEqual(item, needle)
		if err != nil {
			return false, err
		}
		if equal {
			return true, nil
		}
	}
	return false, nil
}

func expressionString(value any) (string, bool) {
	switch typed := value.(type) {
	case string:
		return typed, true
	case json.Number:
		return typed.String(), true
	default:
		return "", false
	}
}

func valuesEqual(left, right any) (bool, error) {
	if _, leftOK := left.(json.Number); leftOK {
		if _, rightOK := right.(json.Number); rightOK {
			comparison, err := compareNumbers(left, right)
			return comparison == 0, err
		}
	}
	return reflect.DeepEqual(left, right), nil
}

func compareNumbers(left, right any) (int, error) {
	leftNumber, err := normalizeNumber(left)
	if err != nil {
		return 0, fmt.Errorf("left numeric operand: %w", err)
	}
	rightNumber, err := normalizeNumber(right)
	if err != nil {
		return 0, fmt.Errorf("right numeric operand: %w", err)
	}
	leftRat, leftOK := new(big.Rat).SetString(leftNumber.String())
	rightRat, rightOK := new(big.Rat).SetString(rightNumber.String())
	if !leftOK || !rightOK {
		return 0, fmt.Errorf("invalid numeric operand")
	}
	return leftRat.Cmp(rightRat), nil
}

func validateExpressionResult(valueType plan.ValueType, value any) error {
	if valueType == plan.ArtifactTypeNull {
		if value != nil {
			return fmt.Errorf("expected null, got %T", value)
		}
		return nil
	}
	schema := plan.ArtifactSchema{Type: valueType, Values: []string{}}
	if valueType == plan.TypeEnum {
		// An expression enum has been statically constrained; its runtime
		// representation is the same as string.  Membership is checked at the
		// eventual contract boundary.
		schema.Type = plan.TypeString
	}
	if valueType == plan.TypeArray {
		if _, ok := value.([]any); !ok {
			return fmt.Errorf("expected array, got %T", value)
		}
		return nil
	}
	if valueType == plan.TypeObject {
		_, mapOK := value.(map[string]any)
		_, valuesOK := value.(Values)
		if !mapOK && !valuesOK {
			return fmt.Errorf("expected object, got %T", value)
		}
		return nil
	}
	_, err := normalizeSchemaValue("expression", schema, value)
	return err
}

// EvalTemplate renders a compiled template without parsing authoring syntax.
func EvalTemplate(template plan.ArtifactTemplate, context EvalContext) (string, error) {
	var builder strings.Builder
	for index, segment := range template.Segments {
		switch segment.Kind {
		case plan.ArtifactTemplateSegmentText:
			builder.WriteString(segment.Text)
		case plan.ArtifactTemplateSegmentExpr:
			if segment.Expression == nil {
				return "", fmt.Errorf("template segment %d has no expression", index)
			}
			value, err := EvalExpression(*segment.Expression, context)
			if err != nil {
				return "", fmt.Errorf("template segment %d: %w", index, err)
			}
			text, err := formatTemplateValue(value)
			if err != nil {
				return "", fmt.Errorf("template segment %d: %w", index, err)
			}
			builder.WriteString(text)
		default:
			return "", fmt.Errorf("template segment %d has unknown kind %q", index, segment.Kind)
		}
	}
	return builder.String(), nil
}

func formatTemplateValue(value any) (string, error) {
	switch typed := value.(type) {
	case string:
		return typed, nil
	case json.Number:
		return typed.String(), nil
	case bool:
		if typed {
			return "true", nil
		}
		return "false", nil
	case nil:
		return "null", nil
	default:
		payload, err := json.Marshal(typed)
		if err != nil {
			return "", err
		}
		return string(payload), nil
	}
}