// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2026-present Datadog, Inc.

package tests_test

import (
	"bytes"
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"mvdan.cc/sh/v3/syntax"

	"github.com/DataDog/rshell/internal/interpoption"
	"github.com/DataDog/rshell/interp"
)

// whileRun runs a script with all commands allowed and no path restrictions.
func whileRun(t *testing.T, script string) (string, string, int) {
	t.Helper()
	return whileRunCtx(context.Background(), t, script)
}

// whileRunCtx is whileRun with an explicit context — used by tests that need
// to assert ctx-cancellation behaviour. All commands are allowed; no path
// restrictions are configured.
func whileRunCtx(ctx context.Context, t *testing.T, script string) (string, string, int) {
	t.Helper()
	parser := syntax.NewParser()
	prog, err := parser.Parse(strings.NewReader(script), "")
	require.NoError(t, err)

	var outBuf, errBuf bytes.Buffer
	runner, err := interp.New(
		interp.StdIO(nil, &outBuf, &errBuf),
		interpoption.AllowAllCommands().(interp.RunnerOption),
	)
	require.NoError(t, err)
	defer runner.Close()

	err = runner.Run(ctx, prog)
	exitCode := 0
	if err != nil {
		var es interp.ExitStatus
		if errors.As(err, &es) {
			exitCode = int(es)
		} else if ctx.Err() == nil {
			t.Fatalf("unexpected error: %v", err)
		}
	}
	return outBuf.String(), errBuf.String(), exitCode
}

// --- Basic semantics ---

func TestWhileBasicCounter(t *testing.T) {
	stdout, stderr, code := whileRun(t, `i=; while [ "$i" != aaa ]; do i="${i}a"; echo "$i"; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\naaa\n", stdout)
	assert.Equal(t, "", stderr)
}

func TestUntilBasicCounter(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; until [ "$i" = aaa ]; do i="${i}a"; echo "$i"; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\naaa\n", stdout)
}

func TestWhileFalseZeroIterations(t *testing.T) {
	stdout, _, code := whileRun(t, `while false; do echo never; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=0\n", stdout)
}

func TestUntilTrueZeroIterations(t *testing.T) {
	stdout, _, code := whileRun(t, `until true; do echo never; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=0\n", stdout)
}

// --- Exit-status semantics ---

// Per POSIX 1003.1-2024 §2.9.4.1, the while loop's exit status is the exit
// status of the last command of the last body iteration.
func TestWhileExitStatusFromLastBody(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; false; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=1\n", stdout)
}

// Cond's exit status must NOT leak as the loop's exit status — once the cond
// fails (status 1), $? after the loop should still be 0 if the body never ran.
func TestWhileCondFailDoesNotLeak(t *testing.T) {
	stdout, _, code := whileRun(t, `while false; do :; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=0\n", stdout)
}

func TestUntilCondTrueDoesNotLeak(t *testing.T) {
	stdout, _, code := whileRun(t, `until true; do :; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=0\n", stdout)
}

// Symmetric with TestWhileExitStatusFromLastBody: until's exit status is the
// last body command's exit (POSIX 2.9.4.2).
func TestUntilExitStatusFromLastBody(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; until [ "$i" = aa ]; do i="${i}a"; false; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=1\n", stdout)
}

// --- Pipeline / list condition ---

func TestWhilePipelineCondition(t *testing.T) {
	stdout, _, code := whileRun(t, `while echo x | grep -q x; do echo body; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "body\n", stdout)
}

func TestWhileMultiStmtConditionLastWins(t *testing.T) {
	// Cond is "true; false" — last cmd is false → don't enter body.
	stdout, _, code := whileRun(t, `while true; false; do echo body; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=0\n", stdout)
}

func TestWhileMultiStmtConditionLastTrueRuns(t *testing.T) {
	// Cond runs an assignment then a test. Loop iterates while assignment-then-test
	// returns 0.
	stdout, _, code := whileRun(t, `i=; while i="${i}a"; [ "$i" != aaa ]; do echo "$i"; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\n", stdout)
}

// --- break / continue from body ---

func TestWhileBreakSimple(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do echo a; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\n", stdout)
}

func TestWhileBreakTwoLevels(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do while true; do echo a; break 2; done; echo unreachable; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\ndone\n", stdout)
}

func TestWhileExcessBreakClampedAtOutermost(t *testing.T) {
	// "break 99" at outermost loop → clamped to break-out-of-this-loop.
	stdout, _, code := whileRun(t, `while true; do echo a; break 99; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\ndone\n", stdout)
}

func TestWhileExcessContinueClampedAtOutermost(t *testing.T) {
	// "continue 99" at outermost loop should keep iterating (bash clamp).
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aaa ]; do i="${i}a"; echo "$i"; continue 99; echo unreachable; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\naaa\n", stdout)
}

func TestWhileContinueSkipsRestOfBody(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; continue; echo unreachable; done; echo "i=$i"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "i=aa\n", stdout)
}

// --- Nesting ---

func TestWhileInsideUntilWithBreak2(t *testing.T) {
	stdout, _, code := whileRun(t, `until false; do while true; do echo inner; break 2; done; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "inner\ndone\n", stdout)
}

func TestUntilInsideWhileWithContinue2(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; until false; do echo "iter $i"; continue 2; done; echo unreachable; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "iter a\niter aa\ndone\n", stdout)
}

func TestDeeplyNestedWhile(t *testing.T) {
	// 50 levels of nested while-true; innermost breaks all the way out.
	var b strings.Builder
	for i := 0; i < 50; i++ {
		b.WriteString("while true; do ")
	}
	b.WriteString("echo deep; break 50; ")
	for i := 0; i < 50; i++ {
		b.WriteString("done; ")
	}
	b.WriteString("echo done")
	stdout, _, code := whileRun(t, b.String())
	assert.Equal(t, 0, code)
	assert.Equal(t, "deep\ndone\n", stdout)
}

// --- Interaction with if / brace group ---

func TestWhileBreakInsideElse(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do if false; then :; else break; fi; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "done\n", stdout)
}

func TestWhileContinueInsideElse(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; if false; then :; else continue; fi; echo unreachable; done; echo "i=$i"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "i=aa\n", stdout)
}

func TestWhileBraceBody(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do { echo foo; break; }; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "foo\n", stdout)
}

// --- Pipes feeding the loop / loop feeding pipes ---

func TestWhilePipeOutputToHead(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aaaaa ]; do i="${i}a"; echo "$i"; done | head -3`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\naaa\n", stdout)
}

// --- break/continue inside subshells nested within a loop ---
//
// Bash distinguishes between subshell flavors when validating break/continue:
//
//   - Pipeline stages and command substitutions ($(...) and `…`) silently
//     no-op break/continue when nested inside a loop (no diagnostic, the
//     break/continue counters cannot escape the subshell anyway).
//   - Bare subshells `(...)` always print the diagnostic, even when the bare
//     subshell itself runs inside a loop in the parent shell.
//
// Verified against bash 5.2 (debian:bookworm-slim) — see PR #223 review.

func TestWhileBreakInPipelineSubshellNoDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t, `while break | cat; do echo body; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "body\n", stdout)
	assert.Empty(t, stderr, "break in pipeline subshell should not print a diagnostic when inside a loop")
}

func TestWhileBreakInCmdSubstNoDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t, `while true; do x=$(break); echo body; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "body\n", stdout)
	assert.Empty(t, stderr, "break in $(...) should not print a diagnostic when the substitution runs inside a loop")
}

// Bare subshells `(...)` print the diagnostic even when nested inside a loop —
// matching bash 5.2.
func TestWhileBreakInBareSubshellPrintsDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t, `while true; do (break); echo after; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "after\n", stdout)
	assert.Contains(t, stderr, "break")
	assert.Contains(t, stderr, "loop")
}

// Outside any loop, a bare subshell's break must also print the diagnostic.
func TestBreakInBareSubshellOutsideLoopStillDiagnoses(t *testing.T) {
	_, stderr, _ := whileRun(t, `(break)`)
	assert.Contains(t, stderr, "break is only useful in a loop")
}

// Bash treats a `break` invoked inside a `{...}` pipeline stage as
// out-of-context: it prints the "only useful in a loop" diagnostic AND keeps
// running the rest of the stage's statements. Pipeline subshells must not
// inherit the parent's loop context for compound stages — otherwise the
// stage-internal break/continue counter aborts the rest of the stage's
// statements and silently drops commands that bash still runs.
func TestWhileBreakInGroupedPipelineStageContinuesStage(t *testing.T) {
	stdout, stderr, code := whileRun(t,
		`while true; do echo before; true | { break; echo bothered; }; echo afterpipe; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "before\nbothered\nafterpipe\n", stdout)
	assert.Contains(t, stderr, "break")
	assert.Contains(t, stderr, "loop")
}

// Same property for `continue` inside a `{...}` pipeline stage. We use a
// `for` loop because the while-loop test helper here does not support `((…))`
// arithmetic commands; the property under test is the same regardless of
// loop kind.
func TestForContinueInGroupedPipelineStageContinuesStage(t *testing.T) {
	stdout, stderr, code := whileRun(t,
		`for i in 1 2; do echo before $i; { continue; echo bothered $i; } | cat; echo afterpipe $i; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "before 1\nbothered 1\nafterpipe 1\nbefore 2\nbothered 2\nafterpipe 2\n", stdout)
	assert.Contains(t, stderr, "continue")
	assert.Contains(t, stderr, "loop")
}

// Nested pipelines must propagate loop context to non-rightmost stages. The
// AST for `a | b | c` is left-associative: Pipe(Pipe(a, b), c). Without
// inheriting inLoop into the intermediate Pipe node, the inner pipeline
// runner sees inLoop=false and bare `break`/`continue` in `a` or `b` would
// emit a spurious diagnostic. Verified against bash 5.2.
func TestWhileBreakInNestedPipelineFirstStageNoDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t,
		`while true; do break | cat | cat; echo after; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "after\n", stdout)
	assert.Empty(t, stderr, "break in first stage of nested pipeline must not print a diagnostic when inside a loop")
}

func TestWhileContinueInNestedPipelineFirstStageNoDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t,
		`while true; do continue | cat | cat; echo after; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "after\n", stdout)
	assert.Empty(t, stderr, "continue in first stage of nested pipeline must not print a diagnostic when inside a loop")
}

func TestWhileBreakInFourStagePipelineNoDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t,
		`while true; do break | cat | cat | cat; echo after; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "after\n", stdout)
	assert.Empty(t, stderr)
}

// Compound stage inside a nested pipeline still prints the diagnostic and
// keeps running the rest of the stage's statements — propagation only kicks
// in for simple commands and pipelines, never for compound stages.
func TestWhileBreakInGroupedFirstStageOfNestedPipelinePrintsDiagnostic(t *testing.T) {
	stdout, stderr, code := whileRun(t,
		`while true; do { break; echo never; } | cat | cat; echo after; break; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "never\nafter\n", stdout)
	assert.Contains(t, stderr, "break")
	assert.Contains(t, stderr, "loop")
}

// --- exit propagates out of the loop ---

func TestExitInsideWhileBody(t *testing.T) {
	_, _, code := whileRun(t, `while true; do exit 5; done`)
	assert.Equal(t, 5, code)
}

func TestExitInsideWhileCondition(t *testing.T) {
	_, _, code := whileRun(t, `while exit 7; do echo body; done`)
	assert.Equal(t, 7, code)
}

// --- break / continue from inside the condition list ---

// `break` inside the condition list of a while exits the loop (the body is
// not entered for this iteration).
func TestWhileBreakInCondition(t *testing.T) {
	stdout, _, code := whileRun(t, `while echo cond; break; do echo body; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "cond\nexit=0\n", stdout)
}

// When a body iteration ran (with a non-zero last command), and a SUBSEQUENT
// condition triggers `break`, the loop's exit status must be the break's
// (0) — NOT the stale previous-body status. Bash matches this; an earlier
// version of this code overwrote the break status with the stale body
// status (POSIX §2.9.4 reading too literally).
func TestWhileCondBreakDoesNotInheritStaleBodyExit(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != a ] || break; do i=a; false; done; echo "exit=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "exit=0\n", stdout)
}

// `continue` inside the condition list of a while re-evaluates the condition
// (the body is not entered).
func TestWhileContinueInCondition(t *testing.T) {
	// Use a counter-based cond so we don't hang. Because `continue` inside the
	// cond skips body and re-evaluates cond, we'll spin until [ "$i" = aa ]
	// (cond fails, loop exits).
	stdout, _, code := whileRun(t, `i=; while i="${i}a"; [ "$i" != aa ] && true || break; continue; do echo body; done; echo "i=$i"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "i=aa\n", stdout)
}

// `break N` inside the condition of an inner while propagates outward.
func TestWhileBreak2InCondition(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do while echo inner; break 2; do echo body; done; echo unreachable; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "inner\ndone\n", stdout)
}

// `continue 2` inside the condition of an inner while propagates outward to
// the enclosing loop.
func TestWhileContinue2InCondition(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; while echo "iter $i"; continue 2; do echo body; done; echo unreachable; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "iter a\niter aa\ndone\n", stdout)
}

// `continue N` inside the condition at the OUTERMOST loop must be clamped to
// "continue 1" (re-evaluate cond) — matching bash's behaviour for excess
// continue levels. Without the clamp we would propagate out of the only loop
// and exit cleanly, which is wrong (bash spins forever on this; we simulate
// finite progress via a counter cond).
func TestWhileExcessContinueInCondClampedAtOutermost(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while i="${i}a"; [ "$i" != aaa ] && true || break; continue 99; do echo body; done; echo "i=$i"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "i=aaa\n", stdout)
}

// `continue` invoked inside an `until` condition list short-circuits the rest
// of the cond and gives the cond list a status of 0 (continue's own status).
// For `until`, status 0 means "exit the loop", so the loop must exit rather
// than re-evaluate the cond. An earlier version of this code unconditionally
// re-evaluated cond on continue, which made `until` execute body commands
// that bash skips.
func TestUntilContinueInCondExitsLoop(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; until i=${i}a; [ "$i" != aa ] && continue; do echo body; break; done; echo "i=$i status=$?"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "i=a status=0\n", stdout)
}

// Symmetric check: in a `while` cond, `continue` gives status 0, which for
// `while` means "enter body" — the continue flag then re-evaluates cond
// (skipping body). The loop iterates until the cond evaluates non-zero.
func TestWhileContinueInCondReEvaluates(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while i=${i}a; [ "$i" != aaa ] && continue; do echo body; break; done; echo "i=$i"`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "i=aaa\n", stdout)
}

// `continue 2` invoked inside an inner `until` cond: the inner loop exits
// (regardless of cond status) and the remaining level (1) propagates to the
// enclosing loop, which continues its own iteration.
func TestUntilContinue2InCondPropagates(t *testing.T) {
	stdout, _, code := whileRun(t, `o=; while [ "$o" != aaa ]; do o=${o}a; i=; until i=${i}a; [ "$i" != aa ] && continue 2; do echo body; break; done; echo "post-inner o=$o"; done; echo "after o=$o"`)
	assert.Equal(t, 0, code)
	// The inner until's cond runs `continue 2` on the first iteration, which
	// peels one level (now contnEnclosing=1) and breaks out of inner; outer's
	// `loopStmtsBroken` then sees contnEnclosing=1 and continues outer. So
	// the "post-inner" line is never printed.
	assert.Equal(t, "after o=aaa\n", stdout)
}

// --- ctx cancellation ---

// Long-running while-true terminates within the ctx deadline rather than
// hanging indefinitely. (The ctx-cancel path returns context.Canceled, which
// the helper does not surface as a non-zero exit code; the test passes if it
// returns at all within a small fraction of the test's overall timeout.)
func TestWhileTrueRespectsContextCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	timer := time.AfterFunc(50*time.Millisecond, cancel)
	defer timer.Stop()
	defer cancel()
	start := time.Now()
	_, _, _ = whileRunCtx(ctx, t, `while true; do :; done`)
	assert.Less(t, time.Since(start), 5*time.Second, "loop did not terminate after ctx cancel")
}

// A pre-cancelled context exits the loop on the next iteration's top-of-loop
// ctx check (loopCtx.Err()) — this exercises a path that only fires when the
// parent ctx is already done at the moment we enter a new iteration.
func TestWhileExitsOnPreCancelledContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	stdout, _, _ := whileRunCtx(ctx, t, `while true; do echo x; done`)
	assert.Empty(t, stdout, "no body iterations should run when ctx is pre-cancelled")
}

// Infinite-output while loop must respect the runner's stdout cap rather than
// growing memory unbounded. We use a small ctx deadline as the outer bound
// and assert stdout size stays under a generous upper bound that catches an
// unbounded-growth regression.
func TestWhileTrueOutputRespectsStdoutCap(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	stdout, _, _ := whileRunCtx(ctx, t, `while true; do echo x; done`)
	// 32 MiB is well above any plausible cap, so any value past this points
	// at unbounded buffer growth.
	const generousUpperBound = 1 << 25
	assert.Less(t, len(stdout), generousUpperBound, "stdout grew past the cap; runaway loop?")
}

// vuln-hunt S1 (2026-05-18-initial-audit): `until false` is the symmetric twin
// of `while true`. The runner_exec.go loopCtx.Err() check at the top of every
// iteration must propagate the parent ctx cancel through `until` exactly the
// same way it does through `while` — TestWhileTrueRespectsContextCancellation
// covers `while`, this covers `until`.
func TestUntilFalseRespectsContextCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	timer := time.AfterFunc(50*time.Millisecond, cancel)
	defer timer.Stop()
	defer cancel()
	start := time.Now()
	_, _, _ = whileRunCtx(ctx, t, `until false; do :; done`)
	assert.Less(t, time.Since(start), 5*time.Second, "until-loop did not terminate after ctx cancel")
}

// vuln-hunt S2 (2026-05-18-initial-audit): an infinite while loop placed
// inside a paren subshell must still honour the outer ctx cancel. The
// subshell uses a forked runner with shared ctx, so the cancel must
// propagate through the subshell into the inner while's loopCtx.Err() check.
func TestSubshellInfiniteWhileRespectsContextCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	timer := time.AfterFunc(50*time.Millisecond, cancel)
	defer timer.Stop()
	defer cancel()
	start := time.Now()
	_, _, _ = whileRunCtx(ctx, t, `( while true; do :; done )`)
	assert.Less(t, time.Since(start), 5*time.Second, "subshell-wrapped while did not terminate after ctx cancel")
}

// vuln-hunt P1 (2026-05-18-initial-audit): `while ; do :; done` (empty
// condition list) is a parser error and must never reach the runner — if it
// did, the loop would have nothing to evaluate and could iterate forever.
// This guards against a parser regression that would accept a syntactically
// degenerate while.
func TestWhileEmptyCondIsParserError(t *testing.T) {
	parser := syntax.NewParser()
	_, err := parser.Parse(strings.NewReader(`while ; do :; done`), "")
	require.Error(t, err, "empty while-cond must be rejected by the parser")
	assert.Contains(t, err.Error(), "while", "parser error should mention 'while'")
}
