package taskengine

import (
	"fmt"
	"strings"
	"time"

	avsproto "github.com/AvaProtocol/EigenLayer-AVS/protobuf"
	"google.golang.org/protobuf/types/known/structpb"
)

type BranchProcessor struct {
	*CommonProcessor
}

func NewBranchProcessor(vm *VM) *BranchProcessor {
	return &BranchProcessor{
		CommonProcessor: &CommonProcessor{
			vm: vm,
		},
	}
}

// Helper function to safely extract string values from map[string]interface{}
func getString(m map[string]interface{}, key string) string {
	if val, ok := m[key]; ok {
		if str, ok := val.(string); ok {
			return str
		}
	}
	return ""
}

func (r *BranchProcessor) Validate(node *avsproto.BranchNode) error {
	if node == nil {
		return fmt.Errorf("invalid node data")
	}

	if node.Config == nil {
		return fmt.Errorf("BranchNode Config is nil")
	}

	conditions := node.Config.Conditions
	if len(conditions) == 0 {
		return fmt.Errorf("no conditions defined")
	}

	// Validate each condition
	for i, condition := range conditions {
		if condition == nil {
			return fmt.Errorf("condition at index %d is nil", i)
		}

		if strings.TrimSpace(condition.Id) == "" {
			return fmt.Errorf("condition at index %d has empty ID", i)
		}

		if strings.TrimSpace(condition.Type) == "" {
			return fmt.Errorf("condition at index %d has empty type", i)
		}

		// Validate condition type
		if condition.Type != "if" && condition.Type != "else" {
			return fmt.Errorf("condition at index %d has invalid type: %s (must be 'if' or 'else')", i, condition.Type)
		}

		// First condition must be 'if'
		if i == 0 && condition.Type != "if" {
			return fmt.Errorf("first condition must be 'if', got '%s'", condition.Type)
		}

		// For 'if' conditions, expression can be empty (treated as false)
		// For 'else' conditions, expression should be empty
		if condition.Type == "else" && strings.TrimSpace(condition.Expression) != "" {
			return fmt.Errorf("condition at index %d is 'else' type but has non-empty expression", i)
		}
	}

	return nil
}

func (r *BranchProcessor) Execute(stepID string, node *avsproto.BranchNode) (*avsproto.Execution_Step, *Step, error) {
	// Use shared function to create execution step
	executionStep := createNodeExecutionStep(stepID, avsproto.NodeType_NODE_TYPE_BRANCH, r.vm)

	// Initialize success to false for branches - only set to true when condition matches
	executionStep.Success = false

	var log strings.Builder
	log.WriteString("Start branch execution for node " + stepID + " at " + time.Now().Format(time.RFC3339) + "\n")

	// Get conditions from Config message (static configuration)
	if node.Config == nil {
		err := fmt.Errorf("BranchNode Config is nil")
		log.WriteString("Error: " + err.Error() + "\n")
		finalizeExecutionStep(executionStep, false, err.Error(), log.String())
		return executionStep, nil, err
	}

	conditions := node.Config.Conditions
	if len(conditions) == 0 {
		err := fmt.Errorf("there is no condition to evaluate")
		log.WriteString("Error: " + err.Error() + "\n")
		finalizeExecutionStep(executionStep, false, err.Error(), log.String())
		return executionStep, nil, err
	}

	if conditions[0].Type != "if" {
		err := fmt.Errorf("the first condition need to be an if but got: " + conditions[0].Type)
		log.WriteString("Error: " + err.Error() + "\n")
		finalizeExecutionStep(executionStep, false, err.Error(), log.String())
		return executionStep, nil, err
	}

	// Initialize JavaScript VM for this branch evaluation
	jsvm := NewGojaVM()

	// Set variables from VM context
	r.vm.mu.Lock()
	for key, value := range r.vm.vars {
		if err := jsvm.Set(key, value); err != nil {
			r.vm.mu.Unlock()
			log.WriteString("Error setting variable '" + key + "' in JS VM: " + err.Error() + "\n")
			finalizeExecutionStep(executionStep, false, err.Error(), log.String())
			return executionStep, nil, err
		}
	}
	r.vm.mu.Unlock()

	// Evaluate conditions in order
	for _, condition := range conditions {
		if condition.Type == "else" {
			// Found an else condition, execute it
			log.WriteString("Executing else condition '" + condition.Id + "'\n")
			executionStep.Success = true
			// Create standardized branch data
			branchData := map[string]interface{}{
				"conditionId": fmt.Sprintf("%s.%s", stepID, condition.Id),
			}

			// Convert to protobuf Value
			dataValue, err := structpb.NewValue(branchData)
			if err != nil {
				// Fallback to empty data on error
				dataValue, _ = structpb.NewValue(map[string]interface{}{})
			}

			branchOutput := &avsproto.BranchNode_Output{
				Data: dataValue,
			}

			executionStep.OutputData = &avsproto.Execution_Step_Branch{
				Branch: branchOutput,
			}
			log.WriteString("Branching to else condition '" + condition.Id + "'\n")

			// Find the actual target node from the edges
			conditionId := fmt.Sprintf("%s.%s", stepID, condition.Id)
			var targetNodeId string
			r.vm.mu.Lock()
			if r.vm.task != nil && r.vm.task.Edges != nil {
				for _, edge := range r.vm.task.Edges {
					if edge.Source == conditionId {
						targetNodeId = edge.Target
						break
					}
				}
			}
			r.vm.mu.Unlock()

			// Set the output variable for the branch node
			branchVarOutput := map[string]interface{}{
				"condition_results": []map[string]interface{}{
					{
						"id":           condition.Id,
						"result":       true,
						"next_node_id": targetNodeId,
					},
				},
			}
			// Use shared function to set output variable for this step
			setNodeOutputData(r.CommonProcessor, stepID, branchVarOutput)

			// Find the next step in the plan based on this condition ID
			r.vm.mu.Lock() // Lock for reading vm.plans
			nextStepInPlan, exists := r.vm.plans[fmt.Sprintf("%s.%s", stepID, condition.Id)]
			r.vm.mu.Unlock()
			if !exists {
				// If no plan exists, create a simple step for unit testing purposes
				nextStepInPlan = &Step{NodeID: fmt.Sprintf("%s.%s", stepID, condition.Id), Next: []string{}}
			}

			// Use shared function to finalize execution step
			finalizeExecutionStep(executionStep, true, "", log.String())
			return executionStep, nextStepInPlan, nil
		}

		// Regular if condition
		expression := condition.Expression
		if expression == "" {
			log.WriteString("Condition '" + condition.Id + "' has empty expression, treating as false\n")
			continue
		}

		// Preprocess the expression for template variables
		processedExpression := r.vm.preprocessTextWithVariableMapping(expression)
		log.WriteString("Condition '" + condition.Id + "' original expression: " + expression + "\n")
		log.WriteString("Condition '" + condition.Id + "' processed expression: " + processedExpression + "\n")

		// Trim whitespace
		trimmedProcessed := strings.TrimSpace(processedExpression)

		// Check if the expression is empty after processing
		if trimmedProcessed == "" {
			log.WriteString("Condition '" + condition.Id + "' expression is empty after processing, treating as false\n")
			continue
		}

		// SECURITY: Validate the processed expression before execution
		validationResult := ValidateCodeInjection(trimmedProcessed)
		if !validationResult.Valid {
			log.WriteString("Dangerous expression detected in condition '" + condition.Id + "': " + validationResult.Error + "\n")
			log.WriteString("Original expression: " + condition.Expression + "\n")
			log.WriteString("Processed expression: " + trimmedProcessed + "\n")
			log.WriteString("Condition '" + condition.Id + "' failed security validation, treating as false and continuing\n")
			continue // Skip this condition (treat as false)
		}

		// Evaluate the expression
		wrappedExpression := fmt.Sprintf("(%s)", trimmedProcessed)
		log.WriteString("Final wrapped expression for '" + condition.Id + "': " + wrappedExpression + "\n")
		value, err := jsvm.RunString(wrappedExpression)
		if err != nil {
			log.WriteString("Error evaluating expression for '" + condition.Id + "': " + err.Error() + "\n")
			log.WriteString("Original expression: " + condition.Expression + "\n")
			log.WriteString("Processed expression: " + processedExpression + "\n")
			log.WriteString("Wrapped expression: " + wrappedExpression + "\n")
			// Log the error but silently skip this condition instead of failing the entire branch
			log.WriteString("Condition '" + condition.Id + "' failed to evaluate, treating as false and continuing\n")
			continue // Skip this condition (treat as false)
		}

		boolValue, ok := value.Export().(bool)
		if !ok {
			log.WriteString("Expression for '" + condition.Id + "' did not evaluate to a boolean value, got: " + fmt.Sprintf("%T %v\n", value.Export(), value.Export()))
			continue
		}

		log.WriteString("Condition '" + condition.Id + "' evaluated to: " + fmt.Sprintf("%t\n", boolValue))
		if boolValue {
			executionStep.Success = true

			// Create standardized branch data
			branchData := map[string]interface{}{
				"conditionId": fmt.Sprintf("%s.%s", stepID, condition.Id),
			}

			// Convert to protobuf Value
			dataValue, err := structpb.NewValue(branchData)
			if err != nil {
				// Fallback to empty data on error
				dataValue, _ = structpb.NewValue(map[string]interface{}{})
			}

			branchOutput := &avsproto.BranchNode_Output{
				Data: dataValue,
			}

			executionStep.OutputData = &avsproto.Execution_Step_Branch{
				Branch: branchOutput,
			}
			log.WriteString("Branching to condition '" + condition.Id + "'\n")

			// Find the actual target node from the edges
			conditionId := fmt.Sprintf("%s.%s", stepID, condition.Id)
			var targetNodeId string
			r.vm.mu.Lock()
			if r.vm.task != nil && r.vm.task.Edges != nil {
				for _, edge := range r.vm.task.Edges {
					if edge.Source == conditionId {
						targetNodeId = edge.Target
						break
					}
				}
			}
			r.vm.mu.Unlock()

			// Set the output variable for the branch node
			branchVarOutput := map[string]interface{}{
				"condition_results": []map[string]interface{}{
					{
						"id":           condition.Id,
						"result":       true,
						"next_node_id": targetNodeId,
					},
				},
			}
			// Use shared function to set output variable for this step
			setNodeOutputData(r.CommonProcessor, stepID, branchVarOutput)

			// Find the next step in the plan based on this condition ID
			r.vm.mu.Lock() // Lock for reading vm.plans
			nextStepInPlan, exists := r.vm.plans[fmt.Sprintf("%s.%s", stepID, condition.Id)]
			r.vm.mu.Unlock()
			if !exists {
				// If no plan exists, create a simple step for unit testing purposes
				nextStepInPlan = &Step{NodeID: fmt.Sprintf("%s.%s", stepID, condition.Id), Next: []string{}}
			}

			// Use shared function to finalize execution step
			finalizeExecutionStep(executionStep, true, "", log.String())
			return executionStep, nextStepInPlan, nil
		}
	}

	// If no condition evaluated to true
	log.WriteString("No branch condition evaluated to true.\n")

	// Check if there are any 'else' conditions defined
	hasElseCondition := false
	for _, condition := range conditions {
		if condition.Type == "else" {
			hasElseCondition = true
			break
		}
	}

	// Check if we have any conditions at all
	if len(conditions) == 0 {
		// No conditions at all - this is an error
		noConditionMetError := "no branch condition met"
		log.WriteString(noConditionMetError + "\n")
		finalizeExecutionStep(executionStep, false, noConditionMetError, log.String())
		return executionStep, nil, fmt.Errorf(noConditionMetError)
	} else if hasElseCondition {
		// If there's an else condition but we reached here, it means the else condition failed to execute
		// This should be an error because else conditions should always execute if reached
		noConditionMetError := "no branch condition met"
		log.WriteString(noConditionMetError + "\n")
		finalizeExecutionStep(executionStep, false, noConditionMetError, log.String())
		return executionStep, nil, fmt.Errorf(noConditionMetError)
	} else {
		// If there are only 'if' conditions and none matched, this is a valid "no-op" scenario
		log.WriteString("No conditions matched and no else condition defined - this is a valid no-op.\n")
		executionStep.Success = true
		executionStep.OutputData = nil // No branch action taken

		// Set the output variable for the branch node with empty results
		branchVarOutput := map[string]interface{}{
			"condition_results": []map[string]interface{}{},
		}
		// Use shared function to set output variable for this step
		setNodeOutputData(r.CommonProcessor, stepID, branchVarOutput)

		// Use shared function to finalize execution step
		finalizeExecutionStep(executionStep, true, "", log.String())
		return executionStep, nil, nil // Success with no next step
	}
}
