package profs3

import (
	"bytes"
	"context"
	crand "crypto/rand"
	"fmt"
	"io/ioutil"
	"math"
	"math/big"
	"math/rand"
	"net/url"
	"time"

	"code.justin.tv/video/autoprof"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

const (
	// defaultProfileInterval adjusts the upper limit on the interval between
	// profiles. The package applies jitter to the time between each profile,
	// which will result in slightly more frequent samples on average.
	defaultProfileInterval = 2 * time.Minute
)

// A Collector periodically builds a profile bundle for the process and
// uploads it to S3.
type Collector struct {
	// S3 is the S3 client used for profile bundle uploads.
	S3 s3iface.S3API
	// S3Bucket is the name of the S3 bucket to use for uploads.
	S3Bucket string

	// OnError receives any non-fatal error encountered by the Collector's Run
	// method, such as a problem connecting to S3. Fatal errors, such as the
	// provided Context expiring, cannot be ignored and so are not passed to
	// this function. The function may return nil to ignore the error.
	OnError func(err error) error
}

// checkError calls the user-provided OnError function to check if an error
// should be treated as fatal. If the user-provided function returns nil, the
// Collector will ignore the error.
func (c *Collector) checkError(err error) error {
	if err != nil && c.OnError != nil {
		err = c.OnError(err)
	}
	return err
}

// Run periodically builds a profile bundle for the processes and uploads it
// to S3.
func (c *Collector) Run(ctx context.Context) error {
	seed, err := crand.Int(crand.Reader, big.NewInt(math.MaxInt64))
	if err != nil {
		// fatal error, return immediately
		return err
	}

	r := &runner{
		c:   c,
		rng: rand.New(rand.NewSource(seed.Int64())),
	}

	for i := 0; ; i++ {
		r.delay(ctx, i)
		err = ctx.Err()
		if err != nil {
			// fatal error, return immediately
			return err
		}

		opts := r.options(i)
		err := r.upload(ctx, opts)
		if err := c.checkError(err); err != nil {
			return err
		}
	}
}

type runner struct {
	c *Collector

	rng           *rand.Rand
	nextExecTrace int
}

func (r *runner) options(i int) *autoprof.ArchiveOptions {
	var opts autoprof.ArchiveOptions

	opts.CPUProfileDuration = 5 * time.Second
	opts.CPUProfileByteTarget = 1e6

	// Execution traces are often quite large, and can't really be analyzed in
	// aggregate in the same ways that pprof-formatted profiles can. Store
	// fewer of them.
	if r.nextExecTrace == i {
		const execTraceMaxPeriod = 100
		r.nextExecTrace = i + 1 + r.rng.Intn(execTraceMaxPeriod)
		opts.ExecutionTraceDuration = 1 * time.Second
	}
	opts.ExecutionTraceByteTarget = 1e7

	// Don't include variable-duration profiles on the first run; we'd like a
	// good chance of getting at least a little data from short-lived
	// processes.
	if i == 0 {
		opts.CPUProfileDuration = 0
		opts.ExecutionTraceDuration = 0
	}

	return &opts
}

func (r *runner) delay(ctx context.Context, i int) {
	max := int64(defaultProfileInterval)

	// shorten the delay by up to 100% on the first run, and by up to 20% on
	// subsequent runs
	maxTrim := max
	if i > 0 {
		maxTrim = max / 5
	}

	trim := r.rng.Int63n(maxTrim)
	t := time.NewTimer(time.Duration(max - trim))
	defer t.Stop()

	select {
	case <-ctx.Done():
	case <-t.C:
	}
}

func (r *runner) upload(ctx context.Context, opts *autoprof.ArchiveOptions) error {
	meta := autoprof.CurrentArchiveMeta()

	// Some profile types are sensitive to latency when writing out their data.
	// The execution tracer is one such profile type, which results in
	// systematic failure to record traces of garbage collection in progress in
	// some applications. Use a buffer type that does not introduce large
	// latency spikes.
	buf := &linkedListBuffer{}
	err := autoprof.NewZipCollector(buf, meta, opts).Run(ctx)
	if err != nil {
		return err
	}

	// Now that the latency-sensitive portion is complete, convert the buffer
	// into a format convenient for the S3 SDK.
	body, err := ioutil.ReadAll(buf)
	if err != nil {
		return err
	}

	s3Timeout := 10 * time.Second
	ctx, cancel := context.WithTimeout(ctx, s3Timeout)
	defer cancel()

	_, err = r.c.S3.PutObjectWithContext(ctx, &s3.PutObjectInput{
		Bucket: aws.String(r.c.S3Bucket),
		Key:    aws.String(s3Key(meta)),
		Body:   bytes.NewReader(body),
	})
	if err != nil {
		return err
	}

	return nil
}

func s3Key(m *autoprof.ArchiveMeta) string {
	return fmt.Sprintf("pprof/%s/%s/%s/%s",
		url.PathEscape(m.Main),
		url.PathEscape(m.Hostname),
		url.PathEscape(m.ProcID),
		url.PathEscape(m.CaptureTime))
}
