// 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 wc_test

import (
	"bytes"
	"context"
	"errors"
	"os"
	"path/filepath"
	"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 wcRun(t *testing.T, script, dir string) (string, string, int) {
	t.Helper()
	return wcRunCtx(context.Background(), t, script, dir)
}

func wcRunCtx(ctx context.Context, t *testing.T, script, dir 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
	opts := []interp.RunnerOption{
		interp.StdIO(nil, &outBuf, &errBuf),
		interp.AllowedPaths([]string{dir}),
		interpoption.AllowAllCommands().(interp.RunnerOption),
	}

	runner, err := interp.New(opts...)
	require.NoError(t, err)
	defer runner.Close()

	if dir != "" {
		runner.Dir = dir
	}

	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
}

func wcWriteFile(t *testing.T, dir, name, content string) {
	t.Helper()
	require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(content), 0644))
}

// --- Flag and argument injection ---

func TestWcPentestUnknownFlags(t *testing.T) {
	dir := t.TempDir()
	for _, flag := range []string{"-f", "--follow", "--no-such-flag", "--files0-from=foo"} {
		_, stderr, code := wcRun(t, "wc "+flag, dir)
		assert.Equal(t, 1, code, "flag: %s", flag)
		assert.Contains(t, stderr, "wc:", "flag: %s", flag)
	}
}

func TestWcPentestDoubleDashFlagLikeFile(t *testing.T) {
	dir := t.TempDir()
	wcWriteFile(t, dir, "-v", "hello\n")
	stdout, _, code := wcRun(t, "wc -- -v", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "-v")
}

func TestWcPentestMultipleStdin(t *testing.T) {
	dir := t.TempDir()
	wcWriteFile(t, dir, "file.txt", "hello\n")
	stdout, _, code := wcRun(t, "cat file.txt | wc - -", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "total")
}

// --- Path edge cases ---

func TestWcPentestNonexistentFile(t *testing.T) {
	dir := t.TempDir()
	stdout, stderr, code := wcRun(t, "wc nonexistent.txt", dir)
	assert.Equal(t, 1, code)
	assert.Equal(t, "", stdout)
	assert.Contains(t, stderr, "wc:")
}

func TestWcPentestEmptyFilename(t *testing.T) {
	dir := t.TempDir()
	stdout, stderr, code := wcRun(t, "wc ''", dir)
	assert.Equal(t, 1, code)
	assert.Equal(t, "", stdout)
	assert.Contains(t, stderr, "wc:")
}

// --- Special files ---

func TestWcPentestDevNull(t *testing.T) {
	dir := t.TempDir()
	wcWriteFile(t, dir, "empty.txt", "")
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stdout, _, code := wcRunCtx(ctx, t, "wc empty.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "0")
}

// --- Context cancellation ---

func TestWcPentestContextCancelled(t *testing.T) {
	dir := t.TempDir()
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, _, _ = wcRunCtx(ctx, t, "wc", dir)
}

func TestWcPentestContextTimeout(t *testing.T) {
	dir := t.TempDir()
	wcWriteFile(t, dir, "file.txt", strings.Repeat("hello\n", 10000))
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stdout, _, code := wcRunCtx(ctx, t, "wc file.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "10000")
}

// --- Large input ---

func TestWcPentestLargeFile(t *testing.T) {
	dir := t.TempDir()
	content := strings.Repeat("word word word word word\n", 40000)
	wcWriteFile(t, dir, "large.txt", content)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	stdout, _, code := wcRunCtx(ctx, t, "wc -l large.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "40000")
}

// --- Many files (FD leak check) ---

func TestWcPentestManyFiles(t *testing.T) {
	dir := t.TempDir()
	var args []string
	for i := 0; i < 50; i++ {
		name := filepath.Join(dir, strings.ReplaceAll(filepath.Base(t.Name()), "/", "_")+"_"+string(rune('a'+i%26))+string(rune('0'+i/26))+".txt")
		require.NoError(t, os.WriteFile(name, []byte("x\n"), 0644))
		args = append(args, filepath.Base(name))
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	stdout, _, code := wcRunCtx(ctx, t, "wc "+strings.Join(args, " "), dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "total")
}

// --- Edge case: file with only newlines ---

func TestWcPentestOnlyNewlines(t *testing.T) {
	dir := t.TempDir()
	wcWriteFile(t, dir, "file.txt", strings.Repeat("\n", 100))
	stdout, _, code := wcRun(t, "wc file.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "100")
	assert.Contains(t, stdout, "  0")
}

// --- Edge case: long line ---

func TestWcPentestLongLine(t *testing.T) {
	dir := t.TempDir()
	longLine := strings.Repeat("x", 1024*1024) + "\n"
	wcWriteFile(t, dir, "file.txt", longLine)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stdout, _, code := wcRunCtx(ctx, t, "wc -L file.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "1048576")
}

// --- Invalid UTF-8 at chunk boundary ---

func TestWcPentestInvalidUTF8AtChunkBoundary(t *testing.T) {
	dir := t.TempDir()
	// Build content so that invalid UTF-8 bytes (0xC0 0x80) land at the
	// exact 32 KiB read boundary. This exercises the carry buffer edge
	// case where invalid bytes must be handled in-place (not carried).
	const chunkSize = 32 * 1024
	padding := strings.Repeat("A", chunkSize-1) // fills up to byte 32767
	// Place 0xC0 at offset 32767 (last byte of first chunk) and 0x80 at
	// offset 32768 (first byte of second chunk).
	content := []byte(padding)
	content = append(content, 0xC0, 0x80)
	content = append(content, '\n')

	require.NoError(t, os.WriteFile(filepath.Join(dir, "invalid_utf8.txt"), content, 0644))

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// -c should report exact byte count
	stdout, _, code := wcRunCtx(ctx, t, "wc -c invalid_utf8.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "32770") // chunkSize - 1 + 2 invalid bytes + newline

	// -l should count the newline
	stdout, _, code = wcRunCtx(ctx, t, "wc -l invalid_utf8.txt", dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "1")
}

// --- Flag expansion in loop ---

func TestWcPentestFlagExpansion(t *testing.T) {
	dir := t.TempDir()
	wcWriteFile(t, dir, "file.txt", "hello\n")
	_, stderr, code := wcRun(t, "for flag in --follow; do wc $flag file.txt; done", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "wc:")
}
