// bktec fetches and runs test plans generated by Buildkite
// Test Engine.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"syscall"
	"time"

	"github.com/buildkite/test-engine-client/internal/api"
	"github.com/buildkite/test-engine-client/internal/config"
	"github.com/buildkite/test-engine-client/internal/debug"
	"github.com/buildkite/test-engine-client/internal/env"
	"github.com/buildkite/test-engine-client/internal/plan"
	"github.com/buildkite/test-engine-client/internal/runner"
	"github.com/buildkite/test-engine-client/internal/version"
	"github.com/olekukonko/tablewriter"
)

const Logo = `
______ ______ _____
___  /____  /___  /____________
__  __ \_  //_/  __/  _ \  ___/
_  /_/ /  ,<  / /_ /  __/ /__
/_.___//_/|_| \__/ \___/\___/
`

func printStartUpMessage() {
	const green = "\033[32m"
	const reset = "\033[0m"
	fmt.Println("+++ Buildkite Test Engine Client: bktec " + version.Version + "\n")
	fmt.Println(green + Logo + reset)
}

type TestRunner interface {
	// Run takes testCases as input, executes the test against the test cases, and mutates the runner.RunResult with the test results.
	Run(result *runner.RunResult, testCases []plan.TestCase, retry bool) error
	// GetExamples discovers all tests within given files.
	// This function is only used for split by example use case. Currently only supported by RSpec.
	GetExamples(files []string) ([]plan.TestCase, error)
	// GetFiles discover all test files that the runner should execute.
	// This is sent to server-side when creating test plan.
	// This is also used to obtain a fallback non-intelligent test splitting mechanism.
	GetFiles() ([]string, error)
	Name() string
}

func main() {
	env := env.OS{}

	debug.SetDebug(env.Get("BUILDKITE_TEST_ENGINE_DEBUG_ENABLED") == "true")

	versionFlag := flag.Bool("version", false, "print version information")

	flag.Parse()

	if *versionFlag {
		fmt.Printf("bktec %s\n", version.Version)
		os.Exit(0)
	}

	printStartUpMessage()

	// get config
	cfg, err := config.New(env)
	if err != nil {
		logErrorAndExit(16, "Invalid configuration...\n%v", err)
	}

	testRunner, err := runner.DetectRunner(cfg)
	if err != nil {
		logErrorAndExit(16, "Unsupported value for BUILDKITE_TEST_ENGINE_TEST_RUNNER %q: %v", cfg.TestRunner, err)
	}

	files, err := testRunner.GetFiles()
	if err != nil {
		logErrorAndExit(16, "Couldn't get files: %v", err)
	}

	// get plan
	ctx := context.Background()
	apiClient := api.NewClient(api.ClientConfig{
		ServerBaseUrl:    cfg.ServerBaseUrl,
		AccessToken:      cfg.AccessToken,
		OrganizationSlug: cfg.OrganizationSlug,
	})

	testPlan, err := fetchOrCreateTestPlan(ctx, apiClient, cfg, files, testRunner)
	if err != nil {
		logErrorAndExit(16, "Couldn't fetch or create test plan: %v", err)
	}

	debug.Printf("My favourite ice cream is %s", testPlan.Experiment)

	// get plan for this node
	thisNodeTask := testPlan.Tasks[strconv.Itoa(cfg.NodeIndex)]

	// execute tests
	var timeline []api.Timeline
	runResult, err := runTestsWithRetry(testRunner, &thisNodeTask.Tests, cfg.MaxRetries, testPlan.MutedTests, &timeline, cfg.RetryForMutedTest)

	// Handle errors that prevent the runner from finishing.
	// By finishing, it means that the runner has completed with a readable result.
	if err != nil {
		// runner terminated by signal: exit with 128 + signal number
		if ProcessSignaledError := new(runner.ProcessSignaledError); errors.As(err, &ProcessSignaledError) {
			logSignalAndExit(testRunner.Name(), ProcessSignaledError.Signal)
		}

		// runner exited with error: exit with the exit code
		if exitError := new(exec.ExitError); errors.As(err, &exitError) {
			logErrorAndExit(exitError.ExitCode(), "%s exited with error: %v", testRunner.Name(), err)
		}

		// other errors: exit with 16
		logErrorAndExit(16, "Couldn't run tests: %v", err)
	}

	// At this point, the runner is expected to have completed

	if !testPlan.Fallback {
		sendMetadata(ctx, apiClient, cfg, timeline, runResult.Statistics())
	}

	printReport(runResult, testPlan.SkippedTests, testRunner.Name())

	if runResult.Status() == runner.RunStatusFailed || runResult.Status() == runner.RunStatusError {
		os.Exit(1)
	}
}

func printReport(runResult runner.RunResult, testsSkippedByTestEngine []plan.TestCase, runnerName string) {
	fmt.Println("+++ ========== Buildkite Test Engine Report  ==========")

	switch runResult.Status() {
	case runner.RunStatusPassed:
		fmt.Println("✅ All tests passed.")
	case runner.RunStatusFailed:
		fmt.Println("❌ Some tests failed.")
	case runner.RunStatusError:
		fmt.Printf("🚨 %s\n", runResult.Error())
	}
	fmt.Println("")

	// Print statistics
	statistics := runResult.Statistics()
	data := [][]string{
		{"Passed", "first run", strconv.Itoa(statistics.PassedOnFirstRun)},
		{"Passed", "on retry", strconv.Itoa(statistics.PassedOnRetry)},
		{"Muted", "passed", strconv.Itoa(statistics.MutedPassed)},
		{"Muted", "failed", strconv.Itoa(statistics.MutedFailed)},
		{"Failed", "", strconv.Itoa(statistics.Failed)},
		{"Skipped", "", strconv.Itoa(statistics.Skipped)},
	}
	table := tablewriter.NewWriter(os.Stdout)
	table.AppendBulk(data)
	table.SetFooter([]string{"", "Total", strconv.Itoa(statistics.Total)})
	table.SetFooterAlignment(tablewriter.ALIGN_RIGHT)
	table.SetAutoMergeCellsByColumnIndex([]int{0, 1})
	table.SetRowLine(true)
	table.Render()

	// Print muted and failed tests
	mutedTests := runResult.MutedTests()
	if len(mutedTests) > 0 {
		fmt.Println("")
		fmt.Println("+++ Muted Tests:")
		for _, mutedTest := range runResult.MutedTests() {
			fmt.Printf("- %s %s (%s)\n", mutedTest.Scope, mutedTest.Name, mutedTest.Status)
		}
	}

	failedTests := runResult.FailedTests()
	if len(failedTests) > 0 {
		fmt.Println("")
		fmt.Println("+++ Failed Tests:")
		for _, failedTests := range runResult.FailedTests() {
			fmt.Printf("- %s %s\n", failedTests.Scope, failedTests.Name)
		}
	}

	testsSkippedByRunner := runResult.SkippedTests()
	if len(testsSkippedByRunner) > 0 {
		fmt.Println("")
		fmt.Printf("+++ Skipped by %s:\n", runnerName)
		for _, skippedTest := range testsSkippedByRunner {
			fmt.Printf("- %s %s\n", skippedTest.Scope, skippedTest.Name)
		}
	}

	if len(testsSkippedByTestEngine) > 0 {
		fmt.Println("")
		fmt.Println("+++ Skipped by Test Engine:")
		for _, skippedTest := range testsSkippedByTestEngine {
			fmt.Printf("- %s %s\n", skippedTest.Scope, skippedTest.Name)
		}
	}

	fmt.Println("===================================================")
}

func createTimestamp() string {
	return time.Now().Format(time.RFC3339Nano)
}

func sendMetadata(ctx context.Context, apiClient *api.Client, cfg config.Config, timeline []api.Timeline, statistics runner.RunStatistics) {
	err := apiClient.PostTestPlanMetadata(ctx, cfg.SuiteSlug, cfg.Identifier, api.TestPlanMetadataParams{
		Timeline:   timeline,
		Env:        cfg.DumpEnv(),
		Version:    version.Version,
		Statistics: statistics,
	})

	// Error is suppressed because we don't want to fail the build if we can't send metadata.
	if err != nil {
		fmt.Printf("Failed to send metadata to Test Engine: %v\n", err)
	}
}

// runTestsWithRetry can be considered the core of bktec.
// It invoke testRunner, orchestrate retries.
// It returns RunResult.
//
// For next reader, there is a small caveat with current implementation:
// - testCases and timeline are both expected to be mutated.
// - testCases in this case serve both as input and output -> we should probably change it.
func runTestsWithRetry(testRunner TestRunner, testsCases *[]plan.TestCase, maxRetries int, mutedTests []plan.TestCase, timeline *[]api.Timeline, retryForMutedTest bool) (runner.RunResult, error) {
	attemptCount := 0

	// Create a new run result with muted tests to keep track of the results.
	runResult := runner.NewRunResult(mutedTests)

	for attemptCount <= maxRetries {
		if attemptCount == 0 {
			fmt.Printf("+++ Buildkite Test Engine Client: Running tests\n")
			*timeline = append(*timeline, api.Timeline{
				Event:     "test_start",
				Timestamp: createTimestamp(),
			})
		} else {
			fmt.Printf("+++ Buildkite Test Engine Client: ♻️ Attempt %d of %d to retry failing tests\n", attemptCount, maxRetries)
			*timeline = append(*timeline, api.Timeline{
				Event:     fmt.Sprintf("retry_%d_start", attemptCount),
				Timestamp: createTimestamp(),
			})
		}

		err := testRunner.Run(runResult, *testsCases, attemptCount > 0)

		if attemptCount == 0 {
			*timeline = append(*timeline, api.Timeline{
				Event:     "test_end",
				Timestamp: createTimestamp(),
			})
		} else {
			*timeline = append(*timeline, api.Timeline{
				Event:     fmt.Sprintf("retry_%d_end", attemptCount),
				Timestamp: createTimestamp(),
			})
		}

		// Don't retry if there is an error that is not a test failure.
		if err != nil {
			return *runResult, err
		}

		// Don't retry if we've reached max retries.
		if attemptCount == maxRetries {
			return *runResult, nil
		}

		failedTests := runResult.FailedTests()
		failedMutedTests := runResult.FailedMutedTests()

		shouldRetryForHardFailedTests := len(failedTests) > 0
		shouldRetryForMutedTests := retryForMutedTest && len(failedMutedTests) > 0
		shouldRetry := shouldRetryForHardFailedTests || shouldRetryForMutedTests

		if shouldRetry {
			*testsCases = failedTests

			if shouldRetryForMutedTests {
				*testsCases = append(*testsCases, failedMutedTests...)
			}

			attemptCount++
		} else {
			return *runResult, nil
		}
	}

	return *runResult, nil
}

func logSignalAndExit(name string, signal syscall.Signal) {
	fmt.Printf("Buildkite Test Engine Client: %s was terminated with signal: %v\n", name, signal)

	// Exit with 128 + signal number, the standard convention.
	exitCode := 128 + int(signal)
	os.Exit(exitCode)
}

// logErrorAndExit logs an error message and exits with the given exit code.
func logErrorAndExit(exitCode int, format string, v ...any) {
	fmt.Printf("Buildkite Test Engine Client: "+format+"\n", v...)
	os.Exit(exitCode)
}

// fetchOrCreateTestPlan fetches a test plan from the server, or creates a
// fallback plan if the server is unavailable or returns an error plan.
func fetchOrCreateTestPlan(ctx context.Context, apiClient *api.Client, cfg config.Config, files []string, testRunner TestRunner) (plan.TestPlan, error) {
	debug.Println("Fetching test plan")

	// Fetch the plan from the server's cache.
	cachedPlan, err := apiClient.FetchTestPlan(ctx, cfg.SuiteSlug, cfg.Identifier, cfg.JobRetryCount)

	handleError := func(err error) (plan.TestPlan, error) {
		if errors.Is(err, api.ErrRetryTimeout) {
			fmt.Println("⚠️ Could not fetch or create plan from server, falling back to non-intelligent splitting. Your build may take longer than usual.")
			p := plan.CreateFallbackPlan(files, cfg.Parallelism)
			return p, nil
		}

		if billingError := new(api.BillingError); errors.As(err, &billingError) {
			fmt.Println(billingError.Message)
			fmt.Println("⚠️ Falling back to non-intelligent splitting. Your build may take longer than usual.")
			p := plan.CreateFallbackPlan(files, cfg.Parallelism)
			return p, nil
		}

		return plan.TestPlan{}, err
	}

	if err != nil {
		return handleError(err)
	}

	if cachedPlan != nil {
		// The server can return an "error" plan indicated by an empty task list (i.e. `{"tasks": {}}`).
		// In this case, we should create a fallback plan.
		if len(cachedPlan.Tasks) == 0 {
			fmt.Println("⚠️ Error plan received, falling back to non-intelligent splitting. Your build may take longer than usual.")
			testPlan := plan.CreateFallbackPlan(files, cfg.Parallelism)
			return testPlan, nil
		}

		debug.Printf("Test plan found. Identifier: %q", cfg.Identifier)
		return *cachedPlan, nil
	}

	debug.Println("No test plan found, creating a new plan")
	// If the cache is empty, create a new plan.
	params, err := createRequestParam(ctx, cfg, files, *apiClient, testRunner)
	if err != nil {
		return handleError(err)
	}

	debug.Println("Creating test plan")
	testPlan, err := apiClient.CreateTestPlan(ctx, cfg.SuiteSlug, params)

	if err != nil {
		return handleError(err)
	}

	// The server can return an "error" plan indicated by an empty task list (i.e. `{"tasks": {}}`).
	// In this case, we should create a fallback plan.
	if len(testPlan.Tasks) == 0 {
		fmt.Println("⚠️ Error plan received, falling back to non-intelligent splitting. Your build may take longer than usual.")
		testPlan = plan.CreateFallbackPlan(files, cfg.Parallelism)
		return testPlan, nil
	}

	debug.Printf("Test plan created. Identifier: %q", cfg.Identifier)
	return testPlan, nil
}

// createRequestParam generates the parameters needed for a test plan request.
// For runners other than "rspec", it constructs the test plan parameters with all test files.
// For the "rspec" runner, it filters the test files through the Test Engine API and splits the filtered files into examples.
func createRequestParam(ctx context.Context, cfg config.Config, files []string, client api.Client, runner TestRunner) (api.TestPlanParams, error) {
	testFiles := []plan.TestCase{}
	for _, file := range files {
		testFiles = append(testFiles, plan.TestCase{
			Path: file,
		})
	}

	// Splitting files by example is only supported for rspec runner & cucumber
	if runner.Name() != "RSpec" && runner.Name() != "Cucumber" {
		params := api.TestPlanParams{
			Identifier:  cfg.Identifier,
			Parallelism: cfg.Parallelism,
			Branch:      cfg.Branch,
			Runner:      cfg.TestRunner,
			Tests: api.TestPlanParamsTest{
				Files: testFiles,
			},
		}

		// This is a workaround for the fact that the pytest-pants runner is not
		// supported by the Test Engine API. For now, we use the pytest runner. At
		// some point, there may be a difference between the two runners, but for
		// now the response from the Test Engine API is the same for both runners.
		if cfg.TestRunner == "pytest-pants" {
			params.Runner = "pytest"
		}

		return params, nil
	}

	if cfg.SplitByExample {
		debug.Println("Splitting by example")
	}

	// The SplitByExample flag indicates whether to filter slow files for splitting by example.
	// Regardless of the flag's state, the API will still filter other files that need to be split by example, such as those containing skipped tests.
	// Therefore, we must filter and split files even when SplitByExample is disabled.
	testParams, err := filterAndSplitFiles(ctx, cfg, client, testFiles, runner)
	if err != nil {
		return api.TestPlanParams{}, err
	}

	return api.TestPlanParams{
		Identifier:  cfg.Identifier,
		Parallelism: cfg.Parallelism,
		Branch:      cfg.Branch,
		Runner:      cfg.TestRunner,
		Tests:       testParams,
	}, nil
}

// filterAndSplitFiles filters the test files through the Test Engine API and splits the filtered files into examples.
// It returns the test plan parameters with the examples from the filtered files and the remaining files.
// An error is returned if there is a failure in any of the process.
func filterAndSplitFiles(ctx context.Context, cfg config.Config, client api.Client, files []plan.TestCase, runner TestRunner) (api.TestPlanParamsTest, error) {
	// Filter files that need to be split.
	debug.Printf("Filtering %d files", len(files))
	filteredFiles, err := client.FilterTests(ctx, cfg.SuiteSlug, api.FilterTestsParams{
		Files: files,
		Env:   cfg.DumpEnv(),
	})

	if err != nil {
		return api.TestPlanParamsTest{}, fmt.Errorf("filter tests: %w", err)
	}

	// If no files are filtered, return the all files.
	if len(filteredFiles) == 0 {
		debug.Println("No filtered files found")
		return api.TestPlanParamsTest{
			Files: files,
		}, nil
	}

	debug.Printf("Filtered %d files", len(filteredFiles))
	debug.Printf("Getting examples for %d filtered files", len(filteredFiles))

	filteredFilesMap := make(map[string]bool, len(filteredFiles))
	filteredFilesPath := make([]string, 0, len(filteredFiles))
	for _, file := range filteredFiles {
		filteredFilesMap[file.Path] = true
		filteredFilesPath = append(filteredFilesPath, file.Path)
	}

	examples, err := runner.GetExamples(filteredFilesPath)
	if err != nil {
		return api.TestPlanParamsTest{}, fmt.Errorf("get examples: %w", err)
	}

	debug.Printf("Got %d examples within the filtered files", len(examples))

	// Get the remaining files that are not filtered.
	remainingFiles := make([]plan.TestCase, 0, len(files)-len(filteredFiles))
	for _, file := range files {
		if _, ok := filteredFilesMap[file.Path]; !ok {
			remainingFiles = append(remainingFiles, file)
		}
	}

	return api.TestPlanParamsTest{
		Examples: examples,
		Files:    remainingFiles,
	}, nil
}
