package exploit

import (
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/pipeline"
	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/scope"
	"github.com/google/uuid"
)

// Executor safely runs exploitation commands with scope validation and cleanup.
type Executor struct {
	scopeDef *scope.ScopeDefinition
	cleanup  CleanupRegisterer
	dryRun   bool
	safeMode bool
	confirm  ConfirmFunc

	// allowedExecutables, when non-empty, restricts the binaries the
	// executor is willing to launch to a known tool surface. Empty means
	// no allowlist gate (preserves behaviour for tests / callers that
	// don't opt in via WithAllowedExecutables). See #43.
	allowedExecutables map[string]struct{}
}

// ConfirmFunc is the human-in-the-loop hook for assist mode (4.6.4).
// Called BEFORE every executed step (after scope + safe-mode pass).
// Return false to skip the step (treated as a benign no-op, not a
// failure). Return an error to abort the whole campaign.
//
// CLI wires a TTY-backed implementation in cli/scan.go; agents pass
// nil to opt out.
type ConfirmFunc func(step pipeline.AttackStep) (proceed bool, err error)

// CleanupRegisterer is the interface for registering cleanup actions.
type CleanupRegisterer interface {
	Register(ctx context.Context, campaignID uuid.UUID, command, target string) error
}

// NewExecutor creates a new command executor.
func NewExecutor(scopeDef *scope.ScopeDefinition, cleanup CleanupRegisterer, dryRun bool) *Executor {
	return &Executor{
		scopeDef: scopeDef,
		cleanup:  cleanup,
		dryRun:   dryRun,
	}
}

// WithSafeMode enables the safe-mode allowlist on a new executor.
// Use via: exploit.NewExecutor(scope, cleanup, dryRun).WithSafeMode(true).
// Safe-mode rejects commands containing destructive tokens (rm, DROP,
// TRUNCATE, kill, chmod, chown, shutdown) with an actionable error —
// the researcher can still run them if they really want by dropping the
// flag, but the swarm itself will never fire them.
func (e *Executor) WithSafeMode(on bool) *Executor {
	e.safeMode = on
	return e
}

// WithAllowedExecutables restricts the binaries the executor will
// launch to the given set. The engine derives this list from
// tools.Coordinator.RegisteredToolNames(), so adding a new tool
// adapter automatically widens the allowlist — no hand-maintained
// list to drift. An empty / nil slice disables the gate entirely
// (used by tests). See #43.
func (e *Executor) WithAllowedExecutables(names []string) *Executor {
	if len(names) == 0 {
		e.allowedExecutables = nil
		return e
	}
	e.allowedExecutables = make(map[string]struct{}, len(names))
	for _, n := range names {
		e.allowedExecutables[strings.ToLower(n)] = struct{}{}
	}
	return e
}

// WithConfirm wires assist mode (4.6.4). When non-nil, the executor
// pauses before every command and asks the function to approve. A
// false return skips the step; an error aborts the campaign.
func (e *Executor) WithConfirm(fn ConfirmFunc) *Executor {
	e.confirm = fn
	return e
}

// destructiveTokens is the deny-list safe-mode uses. Case-insensitive
// whole-word match at parse time (parseCommand turns the command into
// []string so 'search' won't match 'sh' for example).
var destructiveTokens = map[string]struct{}{
	"rm":       {},
	"rmdir":    {},
	"shred":    {},
	"kill":     {},
	"killall":  {},
	"chmod":    {},
	"chown":    {},
	"shutdown": {},
	"reboot":   {},
	"halt":     {},
	"mkfs":     {},
	"dd":       {},
	// SQL DDL keywords show up in sqlmap/manual curl payloads; safe mode
	// blocks these to prevent accidental data loss during PoC validation.
	"drop":     {},
	"truncate": {},
	"delete":   {},
}

// Execute runs an attack step with full safety checks.
func (e *Executor) Execute(ctx context.Context, step pipeline.AttackStep, campaignID uuid.UUID) (*pipeline.ExecutionResult, error) {
	start := time.Now()

	// 1. Scope validation — HARD STOP if fails
	if err := scope.ValidateCommand(step.Command, *e.scopeDef); err != nil {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          fmt.Sprintf("BLOCKED: %s", err),
			Success:         false,
			ExecutedAt:      start,
			DurationMs:      0,
		}, fmt.Errorf("scope violation: %w", err)
	}

	// 1b. Allowlist gate (#43) — when enabled, the LLM is only allowed
	// to invoke binaries that match a registered tool adapter. Blocks
	// shell wrappers (bash, sh -c, python -c) and arbitrary interpreters
	// that a prompt injection could try to reach as a bridge to RCE.
	// Both the primary command and the cleanup command must pass.
	if err := e.checkExecutableAllowed(step.Command); err != nil {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          fmt.Sprintf("BLOCKED: %s", err),
			Success:         false,
			ExecutedAt:      start,
			DurationMs:      0,
		}, err
	}
	if step.CleanupCommand != "" {
		if err := e.checkExecutableAllowed(step.CleanupCommand); err != nil {
			return nil, fmt.Errorf("cleanup command rejected: %w", err)
		}
	}

	// 2. Register cleanup BEFORE execution — only after the cleanup
	// command has passed the same executable policy as the primary step.
	if step.CleanupCommand != "" && e.cleanup != nil {
		if err := e.cleanup.Register(ctx, campaignID, step.CleanupCommand, step.Name); err != nil {
			return nil, fmt.Errorf("failed to register cleanup: %w", err)
		}
	}

	// 3. Dry run mode — return command without executing
	if e.dryRun {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          "[DRY RUN] Command would be: " + step.Command,
			Success:         true,
			ExecutedAt:      start,
			DurationMs:      0,
		}, nil
	}

	// 4. Execute with timeout
	timeout := 120 * time.Second
	execCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// Quote-aware parsing; rejects shell metacharacters (|, >, <, &, ;, $(),
	// backticks, newlines) unless they appear inside quotes. An attack step
	// that needs a real shell must wrap its command in `sh -c "..."` explicitly.
	parts, err := parseCommand(step.Command)
	if err != nil {
		return &pipeline.ExecutionResult{
			StepID:          step.ID,
			CampaignID:      campaignID,
			CommandExecuted: step.Command,
			Output:          fmt.Sprintf("BLOCKED: %s", err),
			Success:         false,
			ExecutedAt:      start,
			DurationMs:      0,
		}, fmt.Errorf("unsafe command: %w", err)
	}

	// Safe-mode: reject destructive tokens before execution. This is
	// a belt-and-braces check — programs that disallow automated
	// scanning would NEVER want rm / DROP / etc. to fire accidentally.
	if e.safeMode {
		for _, part := range parts {
			// Split each token into words so quoted args like 'drop table users'
			// are checked word-by-word, not as a single opaque string.
			for _, word := range strings.Fields(part) {
				if _, bad := destructiveTokens[strings.ToLower(word)]; bad {
					msg := fmt.Sprintf("safe-mode blocked destructive token %q (drop --safe-mode to override)", word)
					return &pipeline.ExecutionResult{
						StepID: step.ID, CampaignID: campaignID, CommandExecuted: step.Command,
						Output: "BLOCKED: " + msg, Success: false, ExecutedAt: start,
					}, fmt.Errorf("%s", msg)
				}
			}
		}
	}

	// Assist mode (4.6.4) — pause for human y/N before firing the
	// command. Skips are recorded as a non-fatal step result so the
	// rest of the swarm can keep planning around the gap. An error
	// from confirm aborts the whole campaign (e.g. EOF on stdin →
	// scripted run with no terminal).
	if e.confirm != nil {
		proceed, cerr := e.confirm(step)
		if cerr != nil {
			return nil, fmt.Errorf("assist confirm: %w", cerr)
		}
		if !proceed {
			return &pipeline.ExecutionResult{
				StepID:          step.ID,
				CampaignID:      campaignID,
				CommandExecuted: step.Command,
				Output:          "[SKIPPED by assist mode]",
				Success:         false,
				ExecutedAt:      start,
				DurationMs:      int(time.Since(start).Milliseconds()),
			}, nil
		}
	}

	cmd := exec.CommandContext(execCtx, parts[0], parts[1:]...)
	output, cmdErr := cmd.CombinedOutput()
	err = cmdErr

	result := &pipeline.ExecutionResult{
		StepID:          step.ID,
		CampaignID:      campaignID,
		CommandExecuted: step.Command,
		Output:          string(output),
		ExecutedAt:      start,
		DurationMs:      int(time.Since(start).Milliseconds()),
	}

	if err != nil {
		result.Success = false
		result.Output += "\nError: " + err.Error()
		return result, nil
	}

	// 5. Check output against expected pattern
	if step.ExpectedOutputPattern != "" {
		matched, _ := regexp.MatchString(step.ExpectedOutputPattern, string(output))
		result.Success = matched
	} else {
		result.Success = true
	}

	// 6. Build evidence
	result.Evidence = []pipeline.Evidence{
		{
			Type:        "command_output",
			Content:     string(output),
			Timestamp:   time.Now(),
			Description: fmt.Sprintf("Output from: %s", step.Name),
		},
	}

	return result, nil
}

// DryRun returns the command string with all substitutions applied, without executing.
func (e *Executor) DryRun(step pipeline.AttackStep) string {
	return fmt.Sprintf("[DRY RUN] %s\n  Cleanup: %s\n  Expected: %s",
		step.Command, step.CleanupCommand, step.ExpectedOutputPattern)
}

// checkExecutableAllowed enforces the per-binary allowlist. Returns nil
// when (a) no allowlist is configured (preserves test behaviour and
// callers that don't opt in) or (b) the first token's basename is in the
// allowed set. Uses shellsafe.Parse so quoted-arg edge cases and
// metachar rejection are consistent with the main parse path.
func (e *Executor) checkExecutableAllowed(cmd string) error {
	if len(e.allowedExecutables) == 0 {
		return nil
	}
	parts, err := parseCommand(cmd)
	if err != nil {
		// Parse failure here just means the main parse will also fail
		// with the same error; returning nil keeps a single source of
		// blame for shell-metachar rejections.
		return nil
	}
	exe := strings.ToLower(filepath.Base(strings.TrimSpace(parts[0])))
	if exe == "" {
		return fmt.Errorf("empty executable")
	}
	if _, ok := e.allowedExecutables[exe]; !ok {
		return fmt.Errorf("executable %q is not in the allowed tool surface (use a registered tool adapter or extend the allowlist)", exe)
	}
	return nil
}
