// 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 2024 Datadog, Inc.

package gotesting

import (
	"bytes"
	"context"
	"errors"
	"io"
	"reflect"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	"unsafe"
)

// getFieldPointerFrom gets an unsafe.Pointer (gc-safe type of pointer) to a struct field
// useful to get or set values to private field
func getFieldPointerFrom(value any, fieldName string) (unsafe.Pointer, error) {
	return getFieldPointerFromValue(reflect.Indirect(reflect.ValueOf(value)), fieldName)
}

// getFieldPointerFromValue gets an unsafe.Pointer (gc-safe type of pointer) to a struct field
// useful to get or set values to private field
func getFieldPointerFromValue(value reflect.Value, fieldName string) (unsafe.Pointer, error) {
	member := value.FieldByName(fieldName)
	if member.IsValid() {
		return unsafe.Pointer(member.UnsafeAddr()), nil
	}

	return unsafe.Pointer(nil), errors.New("member is invalid")
}

// copyFieldUsingPointers copies a private field value from one struct to another of the same type
func copyFieldUsingPointers[V any](source any, target any, fieldName string) error {
	return copyFieldUsingPointersWithConversion[V](source, target, fieldName, func(sourceValue V) V { return sourceValue })
}

// copyFieldUsingPointersWithConversion copies a private field value from one struct to another of the same type
func copyFieldUsingPointersWithConversion[V any](source any, target any, fieldName string, convertFunc func(sourceValue V) V) error {
	sourcePtr, err := getFieldPointerFrom(source, fieldName)
	if err != nil {
		return err
	}
	targetPtr, err := getFieldPointerFrom(target, fieldName)
	if err != nil {
		return err
	}

	if targetPtr == nil {
		return errors.New("target pointer is nil")
	}

	if sourcePtr == nil {
		return errors.New("source pointer is nil")
	}

	if (*V)(targetPtr) == nil {
		return errors.New("target pointer value is nil")
	}

	if (*V)(sourcePtr) == nil {
		return errors.New("source pointer value is nil")
	}

	*(*V)(targetPtr) = convertFunc(*(*V)(sourcePtr))
	return nil
}

// ****************
// COMMON
// ****************

// commonPrivateFields is collection of required private fields from testing.common
type commonPrivateFields struct {
	mu      *sync.RWMutex
	output  *[]byte // Output generated by test or benchmark.
	level   *int
	name    *string         // Name of test or benchmark.
	failed  *bool           // Test or benchmark has failed.
	skipped *bool           // Test or benchmark has been skipped.
	parent  *unsafe.Pointer // Parent common
	barrier *chan bool      // Barrier for parallel tests
	signal  *chan bool      // Signal channel for test completion
	sub     *[]*testing.T   // Queue of subtests to be run in parallel.
}

// AddLevel increase or decrease the testing.common.level field value, used by
// testing.B to create the name of the benchmark test
func (c *commonPrivateFields) AddLevel(delta int) int {
	if c.mu == nil {
		return 0
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.level == nil {
		return 0
	}
	*c.level = *c.level + delta
	return *c.level
}

// SetFailed set the boolean value in testing.common.failed field value
func (c *commonPrivateFields) SetFailed(value bool) {
	if c.mu == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.failed == nil {
		return
	}
	*c.failed = value
}

// SetSkipped set the boolean value in testing.common.skipped field value
func (c *commonPrivateFields) SetSkipped(value bool) {
	if c.mu == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.skipped == nil {
		return
	}
	*c.skipped = value
}

// GetOutput returns the output generated by the test or benchmark
func (c *commonPrivateFields) GetOutput() []byte {
	if c.mu == nil {
		return nil
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.output == nil {
		return nil
	}
	return *c.output
}

// AppendOutput appends the output generated by the test or benchmark
func (c *commonPrivateFields) AppendOutput(output []byte) {
	if c.mu == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.output == nil {
		return
	}
	*c.output = append(*c.output, output...)
}

// ****************
// TESTING
// ****************

// getInternalTestArray gets the pointer to the testing.InternalTest array inside a
// testing.M instance containing all the "root" tests
func getInternalTestArray(m *testing.M) *[]testing.InternalTest {
	if ptr, err := getFieldPointerFrom(m, "tests"); err == nil && ptr != nil {
		return (*[]testing.InternalTest)(ptr)
	}
	return nil
}

// getTestPrivateFields is a method to retrieve all required privates field from
// testing.T, returning a commonPrivateFields instance
func getTestPrivateFields(t *testing.T) *commonPrivateFields {
	testFields := &commonPrivateFields{}

	// testing.common
	if ptr, err := getFieldPointerFrom(t, "mu"); err == nil && ptr != nil {
		testFields.mu = (*sync.RWMutex)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "output"); err == nil && ptr != nil {
		testFields.output = (*[]byte)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "level"); err == nil && ptr != nil {
		testFields.level = (*int)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "name"); err == nil && ptr != nil {
		testFields.name = (*string)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "failed"); err == nil && ptr != nil {
		testFields.failed = (*bool)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "skipped"); err == nil && ptr != nil {
		testFields.skipped = (*bool)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "parent"); err == nil && ptr != nil {
		testFields.parent = (*unsafe.Pointer)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "barrier"); err == nil {
		testFields.barrier = (*chan bool)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "signal"); err == nil {
		testFields.signal = (*chan bool)(ptr)
	}
	if ptr, err := getFieldPointerFrom(t, "sub"); err == nil {
		testFields.sub = (*[]*testing.T)(ptr)
	}

	return testFields
}

// getTestParentPrivateFields is a method to retrieve all required parent privates field from
// testing.T.parent, returning a commonPrivateFields instance
func getTestParentPrivateFields(t *testing.T) *commonPrivateFields {
	indirectValue := reflect.Indirect(reflect.ValueOf(t))
	member := indirectValue.FieldByName("parent")
	if member.IsValid() {
		value := member.Elem()
		testFields := &commonPrivateFields{}

		// testing.common
		if ptr, err := getFieldPointerFromValue(value, "mu"); err == nil && ptr != nil {
			testFields.mu = (*sync.RWMutex)(ptr)
		}
		if ptr, err := getFieldPointerFromValue(value, "output"); err == nil && ptr != nil {
			testFields.output = (*[]byte)(ptr)
		}
		if ptr, err := getFieldPointerFromValue(value, "level"); err == nil && ptr != nil {
			testFields.level = (*int)(ptr)
		}
		if ptr, err := getFieldPointerFromValue(value, "name"); err == nil && ptr != nil {
			testFields.name = (*string)(ptr)
		}
		if ptr, err := getFieldPointerFromValue(value, "failed"); err == nil && ptr != nil {
			testFields.failed = (*bool)(ptr)
		}
		if ptr, err := getFieldPointerFromValue(value, "skipped"); err == nil && ptr != nil {
			testFields.skipped = (*bool)(ptr)
		}
		if ptr, err := getFieldPointerFromValue(value, "barrier"); err == nil {
			testFields.barrier = (*chan bool)(ptr)
		}

		return testFields
	}
	return nil
}

// contextMatcher is collection of required private fields from testing.context.match
type contextMatcher struct {
	mu       *sync.Mutex
	subNames *map[string]int32
}

// ClearSubNames clears the subname map used for creating unique names for subtests
func (c *contextMatcher) ClearSubNames() {
	if c.mu == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.subNames == nil || *c.subNames == nil {
		return
	}
	clear(*c.subNames)
}

// getTestContextMatcherPrivateFields is a method to retrieve all required privates field from
// testing.T.context.match, returning a contextMatcher instance
func getTestContextMatcherPrivateFields(t *testing.T) *contextMatcher {
	indirectValue := reflect.Indirect(reflect.ValueOf(t))
	contextMember := indirectValue.FieldByName("context")
	if !contextMember.IsValid() {
		// In 1.24 they changed the name of the field to tstate
		contextMember = indirectValue.FieldByName("tstate")
		if !contextMember.IsValid() {
			return nil
		}
	}
	contextMember = contextMember.Elem()
	matchMember := contextMember.FieldByName("match")
	if !matchMember.IsValid() {
		return nil
	}
	matchMember = matchMember.Elem()

	fields := &contextMatcher{}
	if ptr, err := getFieldPointerFromValue(matchMember, "mu"); err == nil && ptr != nil {
		fields.mu = (*sync.Mutex)(ptr)
	}
	if ptr, err := getFieldPointerFromValue(matchMember, "subNames"); err == nil && ptr != nil {
		fields.subNames = (*map[string]int32)(ptr)
	}

	return fields
}

var copyMutex sync.Mutex // Mutex to protect concurrent read and writes to testing.T

// copyTestWithoutParent tries to copy all private fields except the t.parent from a *testing.T to another
func copyTestWithoutParent(source *testing.T, target *testing.T) {
	copyMutex.Lock()
	defer copyMutex.Unlock()
	// Copy important field values
	_ = copyFieldUsingPointers[[]byte](source, target, "output") // Output generated by test or benchmark.
	_ = copyFieldUsingPointersWithConversion[io.Writer](source, target, "w", func(sourceValue io.Writer) io.Writer {
		return getThreadSafeWriter(sourceValue) // Wrap the original writer in a thread-safe writer.
	}) // For flushToParent.
	_ = copyFieldUsingPointers[bool](source, target, "ran")                        // Test or benchmark (or one of its subtests) was executed.
	_ = copyFieldUsingPointers[bool](source, target, "failed")                     // Test or benchmark has failed.
	_ = copyFieldUsingPointers[bool](source, target, "skipped")                    // Test or benchmark has been skipped.
	_ = copyFieldUsingPointers[bool](source, target, "done")                       // Test is finished and all subtests have completed.
	_ = copyFieldUsingPointers[map[uintptr]struct{}](source, target, "helperPCs")  // functions to be skipped when writing file/line info
	_ = copyFieldUsingPointers[map[string]struct{}](source, target, "helperNames") // helperPCs converted to function names
	_ = copyFieldUsingPointers[[]func()](source, target, "cleanups")               // optional functions to be called at the end of the test
	_ = copyFieldUsingPointers[string](source, target, "cleanupName")              // Name of the cleanup function.
	_ = copyFieldUsingPointers[[]uintptr](source, target, "cleanupPc")             // The stack trace at the point where Cleanup was called.
	_ = copyFieldUsingPointers[bool](source, target, "finished")                   // Test function has completed.
	_ = copyFieldUsingPointers[bool](source, target, "inFuzzFn")                   // Whether the fuzz target, if this is one, is running.

	_ = copyFieldUsingPointers[unsafe.Pointer](source, target, "chatty")      // A copy of chattyPrinter, if the chatty flag is set.
	_ = copyFieldUsingPointers[bool](source, target, "bench")                 // Whether the current test is a benchmark.
	_ = copyFieldUsingPointers[atomic.Bool](source, target, "hasSub")         // whether there are sub-benchmarks.
	_ = copyFieldUsingPointers[atomic.Bool](source, target, "cleanupStarted") // Registered cleanup callbacks have started to execute
	_ = copyFieldUsingPointers[string](source, target, "runner")              // Function name of tRunner running the test.
	_ = copyFieldUsingPointers[bool](source, target, "isParallel")            // Whether the test is parallel.

	_ = copyFieldUsingPointers[int](source, target, "level")            // Nesting depth of test or benchmark.
	_ = copyFieldUsingPointers[[]uintptr](source, target, "creator")    // If level > 0, the stack trace at the point where the parent called t.Run.
	_ = copyFieldUsingPointers[string](source, target, "name")          // Name of test or benchmark.
	_ = copyFieldUsingPointers[unsafe.Pointer](source, target, "start") // Time test or benchmark started
	_ = copyFieldUsingPointers[time.Duration](source, target, "duration")
	_ = copyFieldUsingPointers[[]*testing.T](source, target, "sub")            // Queue of subtests to be run in parallel.
	_ = copyFieldUsingPointers[atomic.Int64](source, target, "lastRaceErrors") // Max value of race.Errors seen during the test or its subtests.
	_ = copyFieldUsingPointers[atomic.Bool](source, target, "raceErrorLogged")
	_ = copyFieldUsingPointers[string](source, target, "tempDir")
	_ = copyFieldUsingPointers[error](source, target, "tempDirErr")
	_ = copyFieldUsingPointers[int32](source, target, "tempDirSeq")

	_ = copyFieldUsingPointers[bool](source, target, "isEnvSet")
	_ = copyFieldUsingPointers[unsafe.Pointer](source, target, "context") // For running tests and subtests.

	// New 1.24 fields
	_ = copyFieldUsingPointers[context.Context](source, target, "ctx")
	_ = copyFieldUsingPointers[context.CancelFunc](source, target, "cancelCtx")

	_ = copyFieldUsingPointers[bool](source, target, "denyParallel")
	_ = copyFieldUsingPointers[unsafe.Pointer](source, target, "tstate") // For running tests and subtests.
}

// createNewTest creates a new testing.T instance
func createNewTest() *testing.T {
	nT := &testing.T{}
	var ctxPtr *context.Context
	if ptr, err := getFieldPointerFrom(nT, "barrier"); err == nil && ptr != nil {
		value := (*chan bool)(ptr)
		*value = make(chan bool) // Initialize the barrier channel
	}
	if ptr, err := getFieldPointerFrom(nT, "signal"); err == nil && ptr != nil {
		value := (*chan bool)(ptr)
		*value = make(chan bool, 1) // Initialize the signal channel
	}
	if ptr, err := getFieldPointerFrom(nT, "ctx"); err == nil && ptr != nil {
		ctxPtr = (*context.Context)(ptr)
		*ctxPtr = context.Background() // Initialize the context
	}
	if ptr, err := getFieldPointerFrom(nT, "cancelCtx"); err == nil && ptr != nil && ctxPtr != nil {
		ctx, cancelCtx := context.WithCancel(*ctxPtr)
		value := (*context.CancelFunc)(ptr)
		*ctxPtr = ctx      // Update the context with the new one
		*value = cancelCtx // Initialize the cancel function
	}
	return nT
}

// chattyPrinter is a struct that holds the private fields of testing.T.chatty
type chattyPrinter struct {
	w        *io.Writer // Writer to which output is sent.
	lastName *string    // Last name used for output.
}

// getTestChattyPrinter retrieves the chattyPrinter from a testing.T instance.
func getTestChattyPrinter(t *testing.T) *chattyPrinter {
	indirectValue := reflect.Indirect(reflect.ValueOf(t))
	contextMember := indirectValue.FieldByName("chatty")
	if !contextMember.IsValid() {
		return nil
	}
	contextMember = contextMember.Elem()
	if !contextMember.IsValid() {
		return nil
	}

	fields := &chattyPrinter{}
	if ptr, err := getFieldPointerFromValue(contextMember, "w"); err == nil && ptr != nil {
		fields.w = (*io.Writer)(ptr)
	}
	if ptr, err := getFieldPointerFromValue(contextMember, "lastName"); err == nil && ptr != nil {
		fields.lastName = (*string)(ptr)
	} else {
		fields.lastName = new(string) // Initialize to avoid nil pointer dereference
	}

	return fields
}

type customWriter struct {
	mu      sync.Mutex               // Mutex to protect concurrent writes.
	chatty  *chattyPrinter           // Pointer to the chattyPrinter for output.
	writer  io.Writer                // The original writer to which output is sent.
	outputs map[string]*bytes.Buffer // Map to store outputs by test name.
}

// internalWrite writes the output to the internal buffer of customWriter.
func (cw *customWriter) internalWrite(lastName string, p []byte) {
	// get the lock
	cw.mu.Lock()
	defer cw.mu.Unlock()
	if cw.outputs == nil {
		cw.outputs = make(map[string]*bytes.Buffer)
	}
	var buf *bytes.Buffer
	if v, ok := cw.outputs[lastName]; !ok {
		buf = new(bytes.Buffer)
		cw.outputs[lastName] = buf
	} else {
		buf = v
	}
	buf.Write(p)
}

// Write implements the io.Writer interface for customWriter.
func (cw *customWriter) Write(p []byte) (n int, err error) {
	if cw.chatty == nil || cw.chatty.w == nil {
		return 0, errors.New("chattyPrinter is not initialized")
	}

	// Write to the internal buffer
	cw.internalWrite(*cw.chatty.lastName, p)

	// Use the chattyPrinter's writer to write the output
	if cw.writer != nil {
		return cw.writer.Write(p)
	}
	return len(p), nil
}

// GetOutput retrieves the output for a specific test name from the customWriter.
func (cw *customWriter) GetOutput(name string) string {
	cw.mu.Lock()
	defer cw.mu.Unlock()
	if buf, ok := cw.outputs[name]; ok {
		return buf.String()
	}
	return ""
}

type threadSafeWriter struct {
	mu     sync.Mutex // Mutex to protect concurrent writes.
	writer io.Writer  // The original writer to which output is sent.
}

// Write implements the io.Writer interface for threadSafeWriter.
func (tsw *threadSafeWriter) Write(p []byte) (n int, err error) {
	tsw.mu.Lock()
	defer tsw.mu.Unlock()
	copyMutex.Lock()
	defer copyMutex.Unlock()
	return tsw.writer.Write(p)
}

// threadSafeWriters is a map to hold thread-safe writers for each original writer
var threadSafeWriters = make(map[io.Writer]io.Writer)

// threadSafeWritersMu is a mutex to protect access to the threadSafeWriters map
var threadSafeWritersMu sync.Mutex

// getThreadSafeWriter returns a thread-safe writer for the given writer.
func getThreadSafeWriter(writer io.Writer) io.Writer {
	threadSafeWritersMu.Lock()
	defer threadSafeWritersMu.Unlock()
	if writer == nil {
		return nil
	}
	if tsw, ok := threadSafeWriters[writer]; ok {
		return tsw
	}
	tsw := &threadSafeWriter{writer: writer}
	threadSafeWriters[writer] = tsw
	return tsw
}

// ****************
// BENCHMARKS
// ****************

// get the pointer to the internal benchmark array
// getInternalBenchmarkArray gets the pointer to the testing.InternalBenchmark array inside
// a testing.M instance containing all the "root" benchmarks
func getInternalBenchmarkArray(m *testing.M) *[]testing.InternalBenchmark {
	if ptr, err := getFieldPointerFrom(m, "benchmarks"); err == nil && ptr != nil {
		return (*[]testing.InternalBenchmark)(ptr)
	}
	return nil
}

// benchmarkPrivateFields is a collection of required private fields from testing.B
// also contains a pointer to the original testing.B for easy access
type benchmarkPrivateFields struct {
	commonPrivateFields
	B         *testing.B
	benchFunc *func(b *testing.B)
	result    *testing.BenchmarkResult
}

// getBenchmarkPrivateFields is a method to retrieve all required privates field from
// testing.B, returning a benchmarkPrivateFields instance
func getBenchmarkPrivateFields(b *testing.B) *benchmarkPrivateFields {
	benchFields := &benchmarkPrivateFields{
		B: b,
	}

	// testing.common
	if ptr, err := getFieldPointerFrom(b, "mu"); err == nil && ptr != nil {
		benchFields.mu = (*sync.RWMutex)(ptr)
	}
	if ptr, err := getFieldPointerFrom(b, "level"); err == nil && ptr != nil {
		benchFields.level = (*int)(ptr)
	}
	if ptr, err := getFieldPointerFrom(b, "name"); err == nil && ptr != nil {
		benchFields.name = (*string)(ptr)
	}
	if ptr, err := getFieldPointerFrom(b, "failed"); err == nil && ptr != nil {
		benchFields.failed = (*bool)(ptr)
	}
	if ptr, err := getFieldPointerFrom(b, "skipped"); err == nil && ptr != nil {
		benchFields.skipped = (*bool)(ptr)
	}
	if ptr, err := getFieldPointerFrom(b, "parent"); err == nil && ptr != nil {
		benchFields.parent = (*unsafe.Pointer)(ptr)
	}

	// testing.B
	if ptr, err := getFieldPointerFrom(b, "benchFunc"); err == nil && ptr != nil {
		benchFields.benchFunc = (*func(b *testing.B))(ptr)
	}
	if ptr, err := getFieldPointerFrom(b, "result"); err == nil && ptr != nil {
		benchFields.result = (*testing.BenchmarkResult)(ptr)
	}

	return benchFields
}
