package expressions

import (
	"fmt"
	"strings"

	"github.com/projectdiscovery/govaluate"

	"github.com/projectdiscovery/nuclei/v3/pkg/operators/common/dsl"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/marker"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/replacer"
	"github.com/projectdiscovery/nuclei/v3/pkg/types"
	stringsutil "github.com/projectdiscovery/utils/strings"
)

// Eval compiles the given expression and evaluate it with the given values preserving the return type
func Eval(expression string, values map[string]interface{}) (interface{}, error) {
	compiled, err := govaluate.NewEvaluableExpressionWithFunctions(expression, dsl.HelperFunctions)
	if err != nil {
		return nil, err
	}
	return compiled.Evaluate(values)
}

// Evaluate checks if the match contains a dynamic variable, for each
// found one we will check if it's an expression and can
// be compiled, it will be evaluated and the results will be returned.
//
// The provided keys from finalValues will be used as variable names
// for substitution inside the expression.
func Evaluate(data string, base map[string]interface{}) (string, error) {
	return evaluate(data, base)
}

// EvaluateByte checks if the match contains a dynamic variable, for each
// found one we will check if it's an expression and can
// be compiled, it will be evaluated and the results will be returned.
//
// The provided keys from finalValues will be used as variable names
// for substitution inside the expression.
func EvaluateByte(data []byte, base map[string]interface{}) ([]byte, error) {
	finalData, err := evaluate(string(data), base)
	return []byte(finalData), err
}

func evaluate(data string, base map[string]interface{}) (string, error) {
	expressions := FindExpressions(data, marker.ParenthesisOpen, marker.ParenthesisClose, base)

	// replace simple placeholders (key => value) MarkerOpen + key + MarkerClose and General + key + General to value
	data = replacer.Replace(data, base)

	// expressions can be:
	// - simple: containing base values keys (variables)
	// - complex: containing helper functions [ + variables]
	// literals like {{2+2}} are not considered expressions
	for _, expression := range expressions {
		originalExpression := expression
		// data has already had simple placeholders replaced; keep the same
		// marker shape for output replacement, but never compile this string.
		replacedExpression := replacer.Replace(expression, base)
		expression = replaceStringPlaceholders(expression, base)

		// turns expressions (either helper functions+base values or base values)
		compiled, err := govaluate.NewEvaluableExpressionWithFunctions(expression, dsl.HelperFunctions)
		if err != nil {
			return data, fmt.Errorf("failed to compile expression %q: %w", originalExpression, err)
		}

		result, err := compiled.Evaluate(base)
		if err != nil {
			return data, fmt.Errorf("failed to evaluate expression %q: %w", originalExpression, err)
		}

		replacement := result
		// Preserve unresolved markers only when a helper call would otherwise
		// hide them from downstream validation. Plain expressions such as
		// comparisons should evaluate normally.
		if markers := unresolvedVarMarkers(compiled.Vars(), base); markers != "" {
			usesFunctions := false
			for _, token := range compiled.Tokens() {
				if token.Kind == govaluate.FUNCTION {
					usesFunctions = true
					break
				}
			}
			if usesFunctions && ContainsUnresolvedVariables(fmt.Sprint(result)) == nil {
				replacement = markers
			}
		}

		// replace incrementally
		data = replacer.ReplaceOne(data, replacedExpression, replacement)
	}
	return data, nil
}

func replaceStringPlaceholders(expression string, base map[string]interface{}) string {
	if len(base) == 0 || (!strings.Contains(expression, marker.ParenthesisOpen) && !strings.Contains(expression, marker.General)) {
		return expression
	}

	var builder strings.Builder
	builder.Grow(len(expression))

	var quote byte
	escaped := false
	for i := 0; i < len(expression); {
		char := expression[i]

		if escaped {
			builder.WriteByte(char)
			escaped = false
			i++
			continue
		}

		if quote != 0 {
			if char == '\\' {
				builder.WriteByte(char)
				escaped = true
				i++
				continue
			}

			if char == quote {
				builder.WriteByte(char)
				quote = 0
				i++
				continue
			}

			replacement, next, ok := stringPlaceholderReplacement(expression, i, quote, base, marker.ParenthesisOpen, marker.ParenthesisClose)
			if ok {
				builder.WriteString(replacement)
				i = next
				continue
			}

			replacement, next, ok = stringPlaceholderReplacement(expression, i, quote, base, marker.General, marker.General)
			if ok {
				builder.WriteString(replacement)
				i = next
				continue
			}

			builder.WriteByte(char)
			i++
			continue
		}

		if char == '\'' || char == '"' {
			quote = char
		}

		builder.WriteByte(char)
		i++
	}

	return builder.String()
}

func stringPlaceholderReplacement(expression string, index int, quote byte, base map[string]interface{}, open, close string) (string, int, bool) {
	if !strings.HasPrefix(expression[index:], open) {
		return "", 0, false
	}

	markerStart := index + len(open)
	markerEnd := strings.Index(expression[markerStart:], close)
	if markerEnd < 0 {
		return "", 0, false
	}

	key := expression[markerStart : markerStart+markerEnd]
	value, ok := base[key]
	if !ok {
		return "", 0, false
	}

	replacement := EscapeStringValue(types.ToString(value), quote)
	next := markerStart + markerEnd + len(close)
	return replacement, next, true
}

// EscapeStringValue escapes value for insertion into a govaluate string literal
// that is already delimited by quote.
func EscapeStringValue(value string, quote byte) string {
	var builder strings.Builder
	builder.Grow(len(value))

	for i := 0; i < len(value); i++ {
		char := value[i]
		if char == '\\' || char == quote {
			builder.WriteByte('\\')
		}
		builder.WriteByte(char)
	}

	return builder.String()
}

// maxIterations to avoid infinite loop
const maxIterations = 250

func FindExpressions(data, OpenMarker, CloseMarker string, base map[string]interface{}) []string {
	var (
		iterations int
		exps       []string
	)
	for iterations <= maxIterations {
		// check if we reached the maximum number of iterations

		iterations++
		// attempt to find open markers
		indexOpenMarker := strings.Index(data, OpenMarker)
		// exits if not found
		if indexOpenMarker < 0 {
			break
		}

		indexOpenMarkerOffset := indexOpenMarker + len(OpenMarker)

		shouldSearchCloseMarker := true
		closeMarkerFound := false
		innerData := data
		var potentialMatch string
		var indexCloseMarker, indexCloseMarkerOffset int
		skip := indexOpenMarkerOffset
		for shouldSearchCloseMarker {
			// attempt to find close marker
			indexCloseMarker = stringsutil.IndexAt(innerData, CloseMarker, skip)
			// if no close markers are found exit
			if indexCloseMarker < 0 {
				shouldSearchCloseMarker = false
				continue
			}
			indexCloseMarkerOffset = indexCloseMarker + len(CloseMarker)

			potentialMatch = innerData[indexOpenMarkerOffset:indexCloseMarker]
			if isExpression(potentialMatch, base) {
				closeMarkerFound = true
				shouldSearchCloseMarker = false
				exps = append(exps, potentialMatch)
			} else {
				skip = indexCloseMarkerOffset
			}
		}

		if closeMarkerFound {
			// move after the close marker
			data = data[indexCloseMarkerOffset:]
		} else {
			// move after the open marker
			data = data[indexOpenMarkerOffset:]
		}
	}
	return exps
}

func isExpression(data string, base map[string]interface{}) bool {
	if _, ok := base[data]; ok {
		return false
	}

	if _, err := govaluate.NewEvaluableExpression(data); err == nil {
		if stringsutil.ContainsAny(data, getFunctionsNames(base)...) {
			return true
		} else if stringsutil.ContainsAny(data, dsl.FunctionNames...) {
			return true
		}
		return false
	}
	_, err := govaluate.NewEvaluableExpressionWithFunctions(data, dsl.HelperFunctions)
	return err == nil
}

// unresolvedVarMarkers returns concatenated {{...}} markers found in the
// string values of the given variable names. Returns "" if none.
func unresolvedVarMarkers(vars []string, base map[string]any) string {
	seen := make(map[string]struct{})
	var markers []string
	for _, varName := range vars {
		val, ok := base[varName]
		if !ok {
			continue
		}
		valStr, ok := val.(string)
		if !ok {
			continue
		}
		for _, match := range unresolvedVariablesRegex.FindAllStringSubmatch(valStr, -1) {
			if len(match) < 2 {
				continue
			}
			if numericalExpressionRegex.MatchString(match[1]) || hasLiteralsOnly(match[1]) {
				continue
			}
			full := marker.ParenthesisOpen + match[1] + marker.ParenthesisClose
			if _, exists := seen[full]; !exists {
				seen[full] = struct{}{}
				markers = append(markers, full)
			}
		}
	}
	return strings.Join(markers, "")
}

func getFunctionsNames(m map[string]interface{}) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}
