/*
Copyright 2025 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package engine provides the engine for executing the discovery rules.
package engine

import (
	"context"
	"log/slog"
	"os"
	"regexp"
	"runtime"
	"strings"

	"github.com/GoogleCloudPlatform/google-guest-agent/cmd/guest_telemetry_extension/isvdiscovery/commandlineexecutor"
	defpb "github.com/GoogleCloudPlatform/google-guest-agent/cmd/guest_telemetry_extension/isvdiscovery/definition/proto"
	"github.com/GoogleCloudPlatform/google-guest-agent/cmd/guest_telemetry_extension/isvdiscovery/engine/versioncommands"
)

// VMInfo contains discovered information about the VM to be used for rule evaluation.
type VMInfo struct {
	ProcessNames   []string
	ProcessPaths   []string
	ProcessArgs    []string
	ProcessEnvVars []string
	Usernames      []string
	OSName         string
}

// ProcessInfo contains discovered information about a specific process.
type ProcessInfo struct {
	Name     string
	Path     string
	Arg      string
	EnvVar   string
	Username string
	OSName   string
}

var versionNumberRegex = regexp.MustCompile(`\.?\d+(\.\d+)*`)
var envVarRegex = regexp.MustCompile(`\$([a-zA-Z_][a-zA-Z0-9_]*|\{([a-zA-Z_][a-zA-Z0-9_]*)\})`)
var safeShellCharsRegex = regexp.MustCompile(`^[a-zA-Z0-9_./=-]+$`)
var executeCommand = commandlineexecutor.ExecuteCommand

// ExecuteRules executes the discovery rules against the VM info and returns the discovery result.
func ExecuteRules(ctx context.Context, req *defpb.DiscoveryRules, vmInfo *VMInfo) *defpb.DiscoveryResult {
	rules := req.GetRules()
	var detectedData []*defpb.DetectedData
	for _, rule := range rules {
		if err := ctx.Err(); err != nil {
			slog.Info("ExecuteRules cancelled")
			break
		}
		foundMatch, processInfo := executeRule(rule, vmInfo)
		if foundMatch {
			version := executeVersionRules(ctx, rule, processInfo)
			detectedData = append(detectedData, defpb.DetectedData_builder{
				Name:    rule.GetDiscoveredWorkloadName(),
				Version: version,
			}.Build())
		}
	}
	return defpb.DiscoveryResult_builder{
		DetectedData: detectedData,
	}.Build()
}

func evalAllCondition(all *defpb.AllCondition, vmInfo *VMInfo) (bool, *ProcessInfo) {
	var processInfo *ProcessInfo
	for _, condition := range all.GetConditions() {
		result, pInfo := checkCondition(condition, vmInfo)
		if !result {
			return false, nil
		}
		if pInfo != nil && pInfo.Path != "" && processInfo == nil {
			processInfo = pInfo
		}
	}
	if all.HasAny() {
		result, pInfo := evalAnyCondition(all.GetAny(), vmInfo)
		if !result {
			return false, nil
		}
		if pInfo != nil && pInfo.Path != "" && processInfo == nil {
			processInfo = pInfo
		}
	}
	return true, processInfo
}

func evalAnyCondition(any *defpb.AnyCondition, vmInfo *VMInfo) (bool, *ProcessInfo) {
	for _, condition := range any.GetConditions() {
		result, pInfo := checkCondition(condition, vmInfo)
		if result {
			return true, pInfo
		}
	}
	if any.HasAll() {
		result, pInfo := evalAllCondition(any.GetAll(), vmInfo)
		if result {
			return true, pInfo
		}
	}
	return false, nil
}

// executeRule executes a single discovery rule.
// Returns true if the rule is satisfied, false otherwise.
func executeRule(rule *defpb.DiscoveryRule, vmInfo *VMInfo) (bool, *ProcessInfo) {
	switch rule.WhichRule() {
	case defpb.DiscoveryRule_Condition_case:
		return checkCondition(rule.GetCondition(), vmInfo)
	case defpb.DiscoveryRule_All_case:
		return evalAllCondition(rule.GetAll(), vmInfo)
	case defpb.DiscoveryRule_Any_case:
		return evalAnyCondition(rule.GetAny(), vmInfo)
	default:
		// This should never happen. Return false if it does.
		return false, nil
	}
}

func resolveCommand(command defpb.VersionCommand, extendedCommand defpb.ExtendedVersionCommand, processInfo *ProcessInfo) (string, bool) {
	var cmd string
	if command == defpb.VersionCommand_VERSION_COMMAND_UNSPECIFIED {
		if extendedCommand == defpb.ExtendedVersionCommand_EXTENDED_VERSION_COMMAND_UNSPECIFIED {
			slog.Debug("Version command is unspecified")
			return "", false
		}
		if int(extendedCommand) < 0 || int(extendedCommand) >= len(versioncommands.Commands.ExtendedCmd) {
			slog.Debug("Received unknown ExtendedVersionCommand", "command", extendedCommand)
			return "", false
		}
		cmd = versioncommands.Commands.ExtendedCmd[extendedCommand]
	} else {
		if int(command) < 0 || int(command) >= len(versioncommands.Commands.Cmd) {
			slog.Debug("Received unknown VersionCommand", "command", command)
			return "", false
		}
		cmd = versioncommands.Commands.Cmd[command]
	}
	if cmd == "USE_DISCOVERED_PROCESS_PATH" {
		if processInfo != nil && processInfo.Path != "" {
			cmd = processInfo.Path
		}
	}
	return cmd, true
}

// shellQuote safely quotes a string for use as a command-line argument in a shell execution.
func shellQuote(s string) string {
	if s == "" {
		return "''"
	}
	// Contains unresolved variables: use double quotes for safe shell expansion
	if strings.Contains(s, "$") {
		// 1. Escape backslashes, double quotes, and backticks for double-quoted string
		s = strings.ReplaceAll(s, `\`, `\\`)
		s = strings.ReplaceAll(s, `"`, `\"`)
		s = strings.ReplaceAll(s, "`", "\\`")

		// 2. Escape '$' unless it introduces a valid unbraced ($VAR) or braced (${VAR}) variable
		validVarRegex := regexp.MustCompile(`^\$([a-zA-Z_]\w*|\{[a-zA-Z_]\w*\})`)
		var buf strings.Builder
		for i := 0; i < len(s); i++ {
			if s[i] == '$' {
				if validVarRegex.MatchString(s[i:]) {
					buf.WriteByte('$')
				} else {
					buf.WriteString(`\$`)
				}
			} else {
				buf.WriteByte(s[i])
			}
		}
		return `"` + buf.String() + `"`
	}
	// Contains spaces or metacharacters: use single quotes for literal interpretation
	if !safeShellCharsRegex.MatchString(s) {
		return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
	}
	// Safe string: no quotes needed
	return s
}

func shellQuoteSlice(args []string) string {
	quoted := make([]string, len(args))
	for i, arg := range args {
		quoted[i] = shellQuote(arg)
	}
	return strings.Join(quoted, " ")
}

func buildCommandParams(cmd string, args []string, runAsUser bool, processInfo *ProcessInfo) commandlineexecutor.Params {
	return buildCommandParamsForOS(cmd, args, runAsUser, processInfo, runtime.GOOS)
}

func buildCommandParamsForOS(cmd string, args []string, runAsUser bool, processInfo *ProcessInfo, goos string) commandlineexecutor.Params {
	shouldRunAsUser := runAsUser && processInfo != nil && processInfo.Username != ""
	if !shouldRunAsUser {
		return commandlineexecutor.Params{
			Executable: cmd,
			Args:       args,
		}
	}
	if goos == "windows" {
		// On Windows (where 'su' is not available), populate the User field on Params
		// to convey the target process user context to the commandlineexecutor.
		return commandlineexecutor.Params{
			Executable: cmd,
			Args:       args,
			User:       processInfo.Username,
		}
	}
	fullCmd := shellQuote(cmd)
	cmdArgs := shellQuoteSlice(args)
	if cmdArgs != "" {
		fullCmd = fullCmd + " " + cmdArgs
	}
	// Note: User field must NOT be set when Executable is "su".
	// "su" must be launched as root so it can switch process credentials to processInfo.Username.
	// We pass -s /bin/sh to override disabled shells (like /sbin/nologin) for service accounts,
	// and -l to run as a login shell so profile environment variables are sourced.
	return commandlineexecutor.Params{
		Executable: "su",
		Args:       []string{"-s", "/bin/sh", "-l", processInfo.Username, "-c", fullCmd},
	}
}

// resolveEnvVars expands environment variables (e.g., $SPARK_HOME or $ORACLE_HOME) in string s
// using the captured environment block of the discovered process. If a variable is not present
// in processInfo, it falls back to the host operating system environment. If still unpopulated,
// the literal variable token (e.g., "$VAR" or "${VAR}") is preserved so subsequent shell executions (via su) can resolve it.
func resolveEnvVars(s string, processInfo *ProcessInfo) string {
	if !strings.Contains(s, "$") {
		return s
	}
	envMap := make(map[string]string)
	if processInfo != nil && processInfo.EnvVar != "" {
		// Split on newlines, carriage returns, and null characters to handle different line endings.
		// This is necessary because the environment block is a single string with these delimiters.
		for _, line := range strings.FieldsFunc(processInfo.EnvVar, func(r rune) bool {
			return r == '\n' || r == '\r' || r == '\x00'
		}) {
			if k, v, ok := strings.Cut(line, "="); ok {
				envMap[k] = v
			}
		}
	}
	return envVarRegex.ReplaceAllStringFunc(s, func(match string) string {
		name := strings.Trim(match[1:], "{}")
		if val, ok := envMap[name]; ok {
			return val
		}
		if val, ok := os.LookupEnv(name); ok {
			return val
		}
		return match
	})
}

// resolveEnvVarsSlice expands environment variables across each element in a command argument slice.
func resolveEnvVarsSlice(args []string, processInfo *ProcessInfo) []string {
	if len(args) == 0 {
		return args
	}
	res := make([]string, len(args))
	for i, arg := range args {
		res[i] = resolveEnvVars(arg, processInfo)
	}
	return res
}

func executeVersionRules(ctx context.Context, rule *defpb.DiscoveryRule, processInfo *ProcessInfo) string {
	for _, versionRule := range rule.GetVersionRules() {
		if err := ctx.Err(); err != nil {
			slog.Info("executeVersionRules cancelled")
			return ""
		}
		var versionRegex string
		if len(versionRule.GetSteps()) > 0 {
			var prevOutput string
			for _, step := range versionRule.GetSteps() {
				if err := ctx.Err(); err != nil {
					slog.Info("executeVersionRules step execution cancelled")
					break
				}
				versionRegex = step.GetRegexMatch()
				cmd, ok := resolveCommand(step.GetCommand(), step.GetExtendedCommand(), processInfo)
				if !ok {
					slog.Debug("Unable to resolve command", "command", step.GetCommand(), "extendedCommand", step.GetExtendedCommand())
					break
				}
				cmd = resolveEnvVars(cmd, processInfo)
				args := resolveEnvVarsSlice(step.GetCommandArgs(), processInfo)
				if step.GetCommand() == defpb.VersionCommand_USE_DISCOVERED_PROCESS_PATH && (processInfo == nil || processInfo.Username == "") {
					slog.Debug("Skipping USE_DISCOVERED_PROCESS_PATH execution: process username is missing", "command", cmd)
					break
				}
				runAsUser := step.GetRunAsDiscoveredProcessUser() || step.GetCommand() == defpb.VersionCommand_USE_DISCOVERED_PROCESS_PATH
				params := buildCommandParams(cmd, args, runAsUser, processInfo)
				if step.GetUsePreviousOutputAsStdin() {
					params.Stdin = prevOutput
				}
				res := executeCommand(ctx, params)
				if res.Error != nil || res.ExitCode != 0 || !res.ExecutableFound {
					slog.Debug("Step command failed", "executable", params.Executable, "args", params.Args, "error", res.Error,
						"exitCode", res.ExitCode, "executableFound", res.ExecutableFound)
					break
				}

				prevOutput = ""
				if versionRegex != "" {
					re, err := regexp.Compile(versionRegex)
					if err == nil {
						if re.MatchString(res.StdOut) {
							prevOutput = res.StdOut
						}
					}
				}
				// If we didn't get valid output, try the next version rule.
				if prevOutput == "" {
					slog.Debug("Step command did not produce valid output", "executable", params.Executable, "args", params.Args, "output", res.StdOut)
					break
				}
			}

			if version, found := extractVersionFromOutput(prevOutput, versionRegex, versionRule.GetVersionExtractPattern()); found {
				return version
			}
			continue
		}

		cmd, ok := resolveCommand(versionRule.GetCommand(), versionRule.GetExtendedCommand(), processInfo)
		if !ok {
			continue
		}
		cmd = resolveEnvVars(cmd, processInfo)
		args := resolveEnvVarsSlice(versionRule.GetCommandArgs(), processInfo)
		if versionRule.GetCommand() == defpb.VersionCommand_USE_DISCOVERED_PROCESS_PATH && (processInfo == nil || processInfo.Username == "") {
			slog.Debug("Skipping USE_DISCOVERED_PROCESS_PATH execution: process username is missing", "command", cmd)
			continue
		}
		runAsUser := versionRule.GetRunAsDiscoveredProcessUser() || versionRule.GetCommand() == defpb.VersionCommand_USE_DISCOVERED_PROCESS_PATH
		params := buildCommandParams(cmd, args, runAsUser, processInfo)
		res := executeCommand(ctx, params)

		if res.Error != nil || res.ExitCode != 0 || !res.ExecutableFound {
			slog.Debug("Command failed", "executable", params.Executable, "args", params.Args, "error", res.Error,
				"exitCode", res.ExitCode, "executableFound", res.ExecutableFound)
			continue
		}
		if version, found := extractVersionFromOutput(res.StdOut, versionRule.GetRegexMatch(), versionRule.GetVersionExtractPattern()); found {
			return version
		}
	}

	return ""
}

func extractVersionFromOutput(stdout, versionRegex, versionExtractPattern string) (string, bool) {
	re, err := regexp.Compile(versionRegex)
	if err != nil {
		slog.Debug("Failed to compile version regex", "regex", versionRegex, "error", err)
		return "", false
	}
	var extractRe *regexp.Regexp
	if versionExtractPattern != "" {
		extractRe, err = regexp.Compile(versionExtractPattern)
		if err != nil {
			slog.Debug("Failed to compile version extract regex", "pattern", versionExtractPattern, "error", err)
			return "", false
		}
	}
	lines := strings.Split(stdout, "\n")
	for _, line := range lines {
		if re.MatchString(line) {
			if extractRe != nil {
				if sub := extractRe.FindStringSubmatch(line); len(sub) > 1 {
					return sub[1], true
				}
			}
			if version := versionFromOutput(line); version != "" {
				return version, true
			}
		}
	}
	return "", false
}

func versionFromOutput(output string) string {
	return versionNumberRegex.FindString(output)
}

func checkCondition(condition *defpb.Condition, vmInfo *VMInfo) (bool, *ProcessInfo) {
	result := true
	var processInfo *ProcessInfo
	switch condition.WhichCondition() {
	case defpb.Condition_StringMatch_case:
		stringMatch := condition.GetStringMatch()
		switch stringMatch.WhichFields() {
		case defpb.StringMatchCondition_VmField_case:
			vmField := stringMatch.GetVmField()
			switch vmField {
			case defpb.StringMatchCondition_VM_PROCESS_NAME:
				result, processInfo = checkStringMatch(stringMatch.GetRegexMatch(), vmInfo.ProcessNames, vmInfo, true)
			case defpb.StringMatchCondition_VM_PROCESS_PATH:
				result, processInfo = checkStringMatch(stringMatch.GetRegexMatch(), vmInfo.ProcessPaths, vmInfo, true)
			case defpb.StringMatchCondition_VM_OS_NAME:
				result, processInfo = checkStringMatch(stringMatch.GetRegexMatch(), []string{vmInfo.OSName}, vmInfo, false)
			case defpb.StringMatchCondition_VM_CLI_ARGS:
				result, processInfo = checkStringMatch(stringMatch.GetRegexMatch(), vmInfo.ProcessArgs, vmInfo, true)
			case defpb.StringMatchCondition_VM_ENV_VARS:
				result, processInfo = checkStringMatch(stringMatch.GetRegexMatch(), vmInfo.ProcessEnvVars, vmInfo, true)
			default:
				// This should never happen. Return false if it does.
				return false, nil
			}
		default:
			// This should never happen. Return false if it does.
			return false, nil
		}
	default:
		// This should never happen. Return false if it does.
		return false, nil
	}

	if condition.GetNegated() {
		result = !result
	}
	return result, processInfo
}

func checkStringMatch(pattern string, values []string, vmInfo *VMInfo, isProcess bool) (bool, *ProcessInfo) {
	for i, value := range values {
		match, err := regexp.MatchString(pattern, value)
		if err == nil && match {
			if isProcess && vmInfo != nil {
				return true, &ProcessInfo{
					Name:     safeGet(vmInfo.ProcessNames, i),
					Path:     safeGet(vmInfo.ProcessPaths, i),
					Arg:      safeGet(vmInfo.ProcessArgs, i),
					EnvVar:   safeGet(vmInfo.ProcessEnvVars, i),
					Username: safeGet(vmInfo.Usernames, i),
					OSName:   vmInfo.OSName,
				}
			}
			if vmInfo != nil {
				return true, &ProcessInfo{OSName: vmInfo.OSName}
			}
			return true, nil
		}
	}
	return false, nil
}

func safeGet(s []string, i int) string {
	if i < len(s) {
		return s[i]
	}
	return ""
}
