package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
	"go.opentelemetry.io/otel/trace"
	"golang.org/x/xerrors"

	"github.com/coder/coder/cli/cliflag"
	"github.com/coder/coder/coderd/tracing"
	"github.com/coder/coder/codersdk"
	"github.com/coder/coder/loadtest/harness"
)

const loadtestTracerName = "coder_loadtest"

func loadtest() *cobra.Command {
	var (
		configPath  string
		outputSpecs []string

		traceEnable          bool
		traceCoder           bool
		traceHoneycombAPIKey string
		tracePropagate       bool
	)
	cmd := &cobra.Command{
		Use:   "loadtest --config <path> [--output json[:path]] [--output text[:path]]]",
		Short: "Load test the Coder API",
		// TODO: documentation and a JSON schema file
		Long: "Perform load tests against the Coder server. The load tests are configurable via a JSON file.",
		Example: formatExamples(
			example{
				Description: "Run a loadtest with the given configuration file",
				Command:     "coder loadtest --config path/to/config.json",
			},
			example{
				Description: "Run a loadtest, reading the configuration from stdin",
				Command:     "cat path/to/config.json | coder loadtest --config -",
			},
			example{
				Description: "Run a loadtest outputting JSON results instead",
				Command:     "coder loadtest --config path/to/config.json --output json",
			},
			example{
				Description: "Run a loadtest outputting JSON results to a file",
				Command:     "coder loadtest --config path/to/config.json --output json:path/to/results.json",
			},
			example{
				Description: "Run a loadtest outputting text results to stdout and JSON results to a file",
				Command:     "coder loadtest --config path/to/config.json --output text --output json:path/to/results.json",
			},
		),
		Hidden: true,
		Args:   cobra.ExactArgs(0),
		RunE: func(cmd *cobra.Command, args []string) error {
			ctx := tracing.SetTracerName(cmd.Context(), loadtestTracerName)

			config, err := loadLoadTestConfigFile(configPath, cmd.InOrStdin())
			if err != nil {
				return err
			}
			outputs, err := parseLoadTestOutputs(outputSpecs)
			if err != nil {
				return err
			}

			client, err := CreateClient(cmd)
			if err != nil {
				return err
			}

			me, err := client.User(ctx, codersdk.Me)
			if err != nil {
				return xerrors.Errorf("fetch current user: %w", err)
			}

			// Only owners can do loadtests. This isn't a very strong check but
			// there's not much else we can do. Ratelimits are enforced for
			// non-owners so hopefully that limits the damage if someone
			// disables this check and runs it against a non-owner account.
			ok := false
			for _, role := range me.Roles {
				if role.Name == "owner" {
					ok = true
					break
				}
			}
			if !ok {
				return xerrors.Errorf("Not logged in as a site owner. Load testing is only available to site owners.")
			}

			// Setup tracing and start a span.
			var (
				shouldTrace                           = traceEnable || traceCoder || traceHoneycombAPIKey != ""
				tracerProvider   trace.TracerProvider = trace.NewNoopTracerProvider()
				closeTracingOnce sync.Once
				closeTracing     = func(_ context.Context) error {
					return nil
				}
			)
			if shouldTrace {
				tracerProvider, closeTracing, err = tracing.TracerProvider(ctx, loadtestTracerName, tracing.TracerOpts{
					Default:   traceEnable,
					Coder:     traceCoder,
					Honeycomb: traceHoneycombAPIKey,
				})
				if err != nil {
					return xerrors.Errorf("initialize tracing: %w", err)
				}
				defer func() {
					closeTracingOnce.Do(func() {
						// Allow time for traces to flush even if command
						// context is canceled.
						ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
						defer cancel()
						_ = closeTracing(ctx)
					})
				}()
			}
			tracer := tracerProvider.Tracer(loadtestTracerName)

			// Disable ratelimits and propagate tracing spans for future
			// requests. Individual tests will setup their own loggers.
			client.BypassRatelimits = true
			client.PropagateTracing = tracePropagate

			// Prepare the test.
			strategy := config.Strategy.ExecutionStrategy()
			th := harness.NewTestHarness(strategy)

			for i, t := range config.Tests {
				name := fmt.Sprintf("%s-%d", t.Type, i)

				for j := 0; j < t.Count; j++ {
					id := strconv.Itoa(j)
					runner, err := t.NewRunner(client.Clone())
					if err != nil {
						return xerrors.Errorf("create %q runner for %s/%s: %w", t.Type, name, id, err)
					}

					th.AddRun(name, id, &runnableTraceWrapper{
						tracer:   tracer,
						spanName: fmt.Sprintf("%s/%s", name, id),
						runner:   runner,
					})
				}
			}

			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "Running load test...")

			testCtx := ctx
			if config.Timeout > 0 {
				var cancel func()
				testCtx, cancel = context.WithTimeout(testCtx, time.Duration(config.Timeout))
				defer cancel()
			}

			// TODO: live progress output
			err = th.Run(testCtx)
			if err != nil {
				return xerrors.Errorf("run test harness (harness failure, not a test failure): %w", err)
			}

			// Print the results.
			res := th.Results()
			for _, output := range outputs {
				var (
					w = cmd.OutOrStdout()
					c io.Closer
				)
				if output.path != "-" {
					f, err := os.Create(output.path)
					if err != nil {
						return xerrors.Errorf("create output file: %w", err)
					}
					w, c = f, f
				}

				switch output.format {
				case loadTestOutputFormatText:
					res.PrintText(w)
				case loadTestOutputFormatJSON:
					err = json.NewEncoder(w).Encode(res)
					if err != nil {
						return xerrors.Errorf("encode JSON: %w", err)
					}
				}

				if c != nil {
					err = c.Close()
					if err != nil {
						return xerrors.Errorf("close output file: %w", err)
					}
				}
			}

			// Cleanup.
			_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\nCleaning up...")
			err = th.Cleanup(ctx)
			if err != nil {
				return xerrors.Errorf("cleanup tests: %w", err)
			}

			// Upload traces.
			if shouldTrace {
				_, _ = fmt.Fprintln(cmd.ErrOrStderr(), "\nUploading traces...")
				closeTracingOnce.Do(func() {
					ctx, cancel := context.WithTimeout(ctx, 1*time.Minute)
					defer cancel()
					err := closeTracing(ctx)
					if err != nil {
						_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "\nError uploading traces: %+v\n", err)
					}
				})
			}

			if res.TotalFail > 0 {
				return xerrors.New("load test failed, see above for more details")
			}

			return nil
		},
	}

	cliflag.StringVarP(cmd.Flags(), &configPath, "config", "", "CODER_LOADTEST_CONFIG_PATH", "", "Path to the load test configuration file, or - to read from stdin.")
	cliflag.StringArrayVarP(cmd.Flags(), &outputSpecs, "output", "", "CODER_LOADTEST_OUTPUTS", []string{"text"}, "Output formats, see usage for more information.")

	cliflag.BoolVarP(cmd.Flags(), &traceEnable, "trace", "", "CODER_LOADTEST_TRACE", false, "Whether application tracing data is collected. It exports to a backend configured by environment variables. See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/protocol/exporter.md")
	cliflag.BoolVarP(cmd.Flags(), &traceCoder, "trace-coder", "", "CODER_LOADTEST_TRACE_CODER", false, "Whether opentelemetry traces are sent to Coder. We recommend keeping this disabled unless we advise you to enable it.")
	cliflag.StringVarP(cmd.Flags(), &traceHoneycombAPIKey, "trace-honeycomb-api-key", "", "CODER_LOADTEST_TRACE_HONEYCOMB_API_KEY", "", "Enables trace exporting to Honeycomb.io using the provided API key.")
	cliflag.BoolVarP(cmd.Flags(), &tracePropagate, "trace-propagate", "", "CODER_LOADTEST_TRACE_PROPAGATE", false, "Enables trace propagation to the Coder backend, which will be used to correlate server-side spans with client-side spans. Only enable this if the server is configured with the exact same tracing configuration as the client.")

	return cmd
}

func loadLoadTestConfigFile(configPath string, stdin io.Reader) (LoadTestConfig, error) {
	if configPath == "" {
		return LoadTestConfig{}, xerrors.New("config is required")
	}

	var (
		configReader io.ReadCloser
	)
	if configPath == "-" {
		configReader = io.NopCloser(stdin)
	} else {
		f, err := os.Open(configPath)
		if err != nil {
			return LoadTestConfig{}, xerrors.Errorf("open config file %q: %w", configPath, err)
		}
		configReader = f
	}

	var config LoadTestConfig
	err := json.NewDecoder(configReader).Decode(&config)
	_ = configReader.Close()
	if err != nil {
		return LoadTestConfig{}, xerrors.Errorf("read config file %q: %w", configPath, err)
	}

	err = config.Validate()
	if err != nil {
		return LoadTestConfig{}, xerrors.Errorf("validate config: %w", err)
	}

	return config, nil
}

type loadTestOutputFormat string

const (
	loadTestOutputFormatText loadTestOutputFormat = "text"
	loadTestOutputFormatJSON loadTestOutputFormat = "json"
	// TODO: html format
)

type loadTestOutput struct {
	format loadTestOutputFormat
	// Up to one path (the first path) will have the value "-" which signifies
	// stdout.
	path string
}

func parseLoadTestOutputs(outputs []string) ([]loadTestOutput, error) {
	var stdoutFormat loadTestOutputFormat

	validFormats := map[loadTestOutputFormat]struct{}{
		loadTestOutputFormatText: {},
		loadTestOutputFormatJSON: {},
	}

	var out []loadTestOutput
	for i, o := range outputs {
		parts := strings.SplitN(o, ":", 2)
		format := loadTestOutputFormat(parts[0])
		if _, ok := validFormats[format]; !ok {
			return nil, xerrors.Errorf("invalid output format %q in output flag %d", parts[0], i)
		}

		if len(parts) == 1 {
			if stdoutFormat != "" {
				return nil, xerrors.Errorf("multiple output flags specified for stdout")
			}
			stdoutFormat = format
			continue
		}
		if len(parts) != 2 {
			return nil, xerrors.Errorf("invalid output flag %d: %q", i, o)
		}

		out = append(out, loadTestOutput{
			format: format,
			path:   parts[1],
		})
	}

	// Default to --output text
	if stdoutFormat == "" && len(out) == 0 {
		stdoutFormat = loadTestOutputFormatText
	}

	if stdoutFormat != "" {
		out = append([]loadTestOutput{{
			format: stdoutFormat,
			path:   "-",
		}}, out...)
	}

	return out, nil
}

type runnableTraceWrapper struct {
	tracer   trace.Tracer
	spanName string
	runner   harness.Runnable

	span trace.Span
}

var _ harness.Runnable = &runnableTraceWrapper{}
var _ harness.Cleanable = &runnableTraceWrapper{}

func (r *runnableTraceWrapper) Run(ctx context.Context, id string, logs io.Writer) error {
	ctx, span := r.tracer.Start(ctx, r.spanName, trace.WithNewRoot())
	defer span.End()
	r.span = span

	traceID := "unknown trace ID"
	spanID := "unknown span ID"
	if span.SpanContext().HasTraceID() {
		traceID = span.SpanContext().TraceID().String()
	}
	if span.SpanContext().HasSpanID() {
		spanID = span.SpanContext().SpanID().String()
	}
	_, _ = fmt.Fprintf(logs, "Trace ID: %s\n", traceID)
	_, _ = fmt.Fprintf(logs, "Span ID: %s\n\n", spanID)

	// Make a separate span for the run itself so the sub-spans are grouped
	// neatly. The cleanup span is also a child of the above span so this is
	// important for readability.
	ctx2, span2 := r.tracer.Start(ctx, r.spanName+" run")
	defer span2.End()
	return r.runner.Run(ctx2, id, logs)
}

func (r *runnableTraceWrapper) Cleanup(ctx context.Context, id string) error {
	c, ok := r.runner.(harness.Cleanable)
	if !ok {
		return nil
	}

	if r.span != nil {
		ctx = trace.ContextWithSpanContext(ctx, r.span.SpanContext())
	}
	ctx, span := r.tracer.Start(ctx, r.spanName+" cleanup")
	defer span.End()

	return c.Cleanup(ctx, id)
}