// 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"
	"fmt"
	"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"
)

func ifRun(t *testing.T, script string) (string, string, int) {
	t.Helper()
	return ifRunCtx(context.Background(), t, script)
}

func ifRunCtx(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
}

// --- Deeply nested if chains ---

func TestIfPentestDeeplyNestedIf(t *testing.T) {
	// 100 levels of nested if should work without stack overflow.
	var b strings.Builder
	for i := 0; i < 100; i++ {
		b.WriteString("if true; then ")
	}
	b.WriteString("echo deep")
	for i := 0; i < 100; i++ {
		b.WriteString("; fi")
	}
	stdout, _, code := ifRun(t, b.String())
	assert.Equal(t, "deep\n", stdout)
	assert.Equal(t, 0, code)
}

func TestIfPentestDeeplyNestedElif(t *testing.T) {
	// 50 elif branches — tests the recursive Else chain.
	var b strings.Builder
	b.WriteString("if false; then echo never")
	for i := 0; i < 50; i++ {
		b.WriteString(fmt.Sprintf("; elif false; then echo never%d", i))
	}
	b.WriteString("; else echo found; fi")
	stdout, _, code := ifRun(t, b.String())
	assert.Equal(t, "found\n", stdout)
	assert.Equal(t, 0, code)
}

// --- Context cancellation ---

func TestIfPentestContextCancellation(t *testing.T) {
	// If with a long-running condition should respect context timeout.
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	// This script has many if statements in sequence; with cancellation it should stop.
	var b strings.Builder
	for i := 0; i < 1000; i++ {
		b.WriteString("if true; then true; fi\n")
	}
	_, _, _ = ifRunCtx(ctx, t, b.String())
	// Just verifying it doesn't hang — no assertion on output needed.
}

// --- Break/continue edge cases ---

func TestIfPentestBreakInCondition(t *testing.T) {
	// break in a for loop body where the if condition uses break.
	// "break" in condition position should propagate correctly.
	stdout, _, code := ifRun(t, `
for i in a b c; do
  if true; then
    if [ "$i" = "b" ]; then break; fi
    echo $i
  fi
done`)
	assert.Equal(t, "a\n", stdout)
	assert.Equal(t, 0, code)
}

func TestIfPentestContinueInElif(t *testing.T) {
	stdout, _, code := ifRun(t, `
for i in a b c; do
  if false; then
    echo never
  elif [ "$i" = "b" ]; then
    continue
  else
    echo $i
  fi
done`)
	assert.Equal(t, "a\nc\n", stdout)
	assert.Equal(t, 0, code)
}

// --- Exit in various positions ---

func TestIfPentestExitInCondition(t *testing.T) {
	// exit in condition should terminate the shell.
	stdout, _, code := ifRun(t, `if exit 7; then echo nope; fi; echo nope2`)
	assert.Equal(t, "", stdout)
	assert.Equal(t, 7, code)
}

func TestIfPentestExitInElse(t *testing.T) {
	_, _, code := ifRun(t, `if false; then echo nope; else exit 13; fi; echo nope2`)
	assert.Equal(t, 13, code)
}

func TestIfPentestExitInElif(t *testing.T) {
	_, _, code := ifRun(t, `if false; then echo a; elif true; then exit 5; fi; echo nope`)
	assert.Equal(t, 5, code)
}

// --- Blocked features inside if constructs ---

func TestIfPentestSubshellInCondition(t *testing.T) {
	stdout, _, code := ifRun(t, `if (true); then echo yes; fi`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "yes\n", stdout)
}

func TestIfPentestFuncDeclInBody(t *testing.T) {
	_, stderr, code := ifRun(t, `if true; then foo() { echo bar; }; fi`)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "function declarations are not supported")
}

func TestIfPentestBackgroundInBody(t *testing.T) {
	_, stderr, code := ifRun(t, "if true; then echo bg &\nfi")
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "background execution (&) is not supported")
}

// --- Multiple if statements in sequence ---

func TestIfPentestSequentialIfs(t *testing.T) {
	stdout, _, code := ifRun(t, `
if true; then echo first; fi
if false; then echo skip; fi
if true; then echo third; fi`)
	assert.Equal(t, "first\nthird\n", stdout)
	assert.Equal(t, 0, code)
}

// --- If with pipe ---

func TestIfPentestIfInPipeline(t *testing.T) {
	stdout, _, code := ifRun(t, `echo hello | if true; then cat; fi`)
	assert.Equal(t, "hello\n", stdout)
	assert.Equal(t, 0, code)
}
