package profiler

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"regexp"
	"runtime"
	"runtime/pprof"
	"sync"
	"time"

	gcemd "cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/internal/version"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"github.com/google/pprof/profile"
	gax "github.com/googleapis/gax-go/v2"
	"google.golang.org/api/option"
	gtransport "google.golang.org/api/transport/grpc"
	pb "google.golang.org/genproto/googleapis/devtools/cloudprofiler/v2"
	edpb "google.golang.org/genproto/googleapis/rpc/errdetails"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	grpcmd "google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

var (
	config       Config
	startOnce    allowUntilSuccess

	getProjectID     = gcemd.ProjectID
	getInstanceName  = gcemd.InstanceName
	getZone          = gcemd.Zone
	startCPUProfile  = pprof.StartCPUProfile
	stopCPUProfile   = pprof.StopCPUProfile
	writeHeapProfile = pprof.WriteHeapProfile
	sleep            = gax.Sleep
	dialGRPC         = gtransport.DialPool
	onGCE            = gcemd.OnGCE
	serviceRegexp    = regexp.MustCompile(`^[a-z]([-a-z0-9_.]{0,253}[a-z0-9])?$`)

	profilingDone chan bool
)

const (
	apiAddress       = "cloudprofiler.googleapis.com:443"
	xGoogAPIMetadata = "x-goog-api-client"
	zoneNameLabel    = "zone"
	versionLabel     = "version"
	languageLabel    = "language"
	instanceLabel    = "instance"
	scope            = "https://www.googleapis.com/auth/monitoring.write"

	maxBackoff        = time.Hour
	backoffMultiplier = 1.3 // Backoff envelope increases by this factor on each retry.
	retryInfoMetadata = "google.rpc.retryinfo-bin"
)

	Service string

	ServiceVersion string

	DebugLogging bool

	MutexProfiling bool

	NoCPUProfiling bool

	NoAllocProfiling bool

	AllocForceGC bool

	NoHeapProfiling bool

	NoGoroutineProfiling bool

	EnableOCTelemetry bool

	ProjectID string

	APIAddr string

	Instance string

	Zone string

	numProfiles int
}

type allowUntilSuccess struct {
	m    sync.Mutex
	done uint32
}

func (o *allowUntilSuccess) do(f func() error) (err error) {
	o.m.Lock()
	defer o.m.Unlock()
	if o.done == 0 {
		if err = f(); err == nil {
			o.done = 1
		}
	} else {
		debugLog("profiler.Start() called again after it was previously called")
		err = nil
	}
	return err
}

func Start(cfg Config, options ...option.ClientOption) error {
	startError := startOnce.do(func() error {
		return start(cfg, options...)
	})
	return startError
}

func start(cfg Config, options ...option.ClientOption) error {
	if err := initializeConfig(cfg); err != nil {
		debugLog("failed to initialize config: %v", err)
		return err
	}
	if config.MutexProfiling {
		if mutexEnabled = enableMutexProfiling(); !mutexEnabled {
			return fmt.Errorf("mutex profiling is not supported by %s, requires Go 1.8 or later", runtime.Version())
		}
	}

	ctx := context.Background()

	opts := []option.ClientOption{
		option.WithEndpoint(config.APIAddr),
		option.WithScopes(scope),
		option.WithUserAgent(fmt.Sprintf("gcloud-go-profiler/%s", version.Repo)),
	}
	if !config.EnableOCTelemetry {
		opts = append(opts, option.WithTelemetryDisabled())
	}
	opts = append(opts, options...)

	connPool, err := dialGRPC(ctx, opts...)
	if err != nil {
		debugLog("failed to dial GRPC: %v", err)
		return err
	}

	a, err := initializeAgent(pb.NewProfilerServiceClient(connPool))
	if err != nil {
		debugLog("failed to start the profiling agent: %v", err)
		return err
	}
	go pollProfilerService(withXGoogHeader(ctx), a)
	return nil
}

func debugLog(format string, e ...interface{}) {
	if config.DebugLogging {
		log.Printf(format, e...)
	}
}

type agent struct {
	client        pb.ProfilerServiceClient
	deployment    *pb.Deployment
	profileLabels map[string]string
	profileTypes  []pb.ProfileType
}

func abortedBackoffDuration(md grpcmd.MD) (time.Duration, error) {
	elem := md[retryInfoMetadata]
	if len(elem) <= 0 {
		return 0, errors.New("no retry info")
	}

	var retryInfo edpb.RetryInfo
	if err := proto.Unmarshal([]byte(elem[0]), &retryInfo); err != nil {
		return 0, err
	} else if time, err := ptypes.Duration(retryInfo.RetryDelay); err != nil {
		return 0, err
	} else {
		if time < 0 {
			return 0, errors.New("negative retry duration")
		}
		return time, nil
	}
}

type retryer struct {
	backoff gax.Backoff
	md      grpcmd.MD
}

func (r *retryer) Retry(err error) (time.Duration, bool) {
	st, _ := status.FromError(err)
	if st != nil && st.Code() == codes.Aborted {
		dur, err := abortedBackoffDuration(r.md)
		if err == nil {
			return dur, true
		}
		debugLog("failed to get backoff duration: %v", err)
	}
	return r.backoff.Pause(), true
}

func (a *agent) createProfile(ctx context.Context) *pb.Profile {
	req := pb.CreateProfileRequest{
		Parent:      "projects/" + a.deployment.ProjectId,
		Deployment:  a.deployment,
		ProfileType: a.profileTypes,
	}

	var p *pb.Profile
	md := grpcmd.New(map[string]string{})

	gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {
		debugLog("creating a new profile via profiler service")
		var err error
		p, err = a.client.CreateProfile(ctx, &req, grpc.Trailer(&md))
		if err != nil {
			debugLog("failed to create profile, will retry: %v", err)
		}
		return err
	}, gax.WithRetry(func() gax.Retryer {
		return &retryer{
			backoff: gax.Backoff{
				Initial:    initialBackoff,
				Max:        maxBackoff,
				Multiplier: backoffMultiplier,
			},
			md: md,
		}
	}))

	debugLog("successfully created profile %v", p.GetProfileType())
	return p
}

func (a *agent) profileAndUpload(ctx context.Context, p *pb.Profile) {
	var prof bytes.Buffer
	pt := p.GetProfileType()

	ptEnabled := false
	for _, enabled := range a.profileTypes {
		if enabled == pt {
			ptEnabled = true
			break
		}
	}

	if !ptEnabled {
		debugLog("skipping collection of disabled profile type: %v", pt)
		return
	}

	switch pt {
	case pb.ProfileType_CPU:
		duration, err := ptypes.Duration(p.Duration)
		if err != nil {
			debugLog("failed to get profile duration for CPU profile: %v", err)
			return
		}
		if err := startCPUProfile(&prof); err != nil {
			debugLog("failed to start CPU profile: %v", err)
			return
		}
		sleep(ctx, duration)
		stopCPUProfile()
	case pb.ProfileType_HEAP:
		if err := heapProfile(&prof); err != nil {
			debugLog("failed to write heap profile: %v", err)
			return
		}
	case pb.ProfileType_HEAP_ALLOC:
		duration, err := ptypes.Duration(p.Duration)
		if err != nil {
			debugLog("failed to get profile duration for allocation profile: %v", err)
			return
		}
		if err := deltaAllocProfile(ctx, duration, config.AllocForceGC, &prof); err != nil {
			debugLog("failed to collect allocation profile: %v", err)
			return
		}
	case pb.ProfileType_THREADS:
		if err := pprof.Lookup("goroutine").WriteTo(&prof, 0); err != nil {
			debugLog("failed to collect goroutine profile: %v", err)
			return
		}
	case pb.ProfileType_CONTENTION:
		duration, err := ptypes.Duration(p.Duration)
		if err != nil {
			debugLog("failed to get profile duration: %v", err)
			return
		}
		if err := deltaMutexProfile(ctx, duration, &prof); err != nil {
			debugLog("failed to collect mutex profile: %v", err)
			return
		}
	default:
		debugLog("unexpected profile type: %v", pt)
		return
	}

	p.ProfileBytes = prof.Bytes()
	p.Labels = a.profileLabels
	req := pb.UpdateProfileRequest{Profile: p}

	debugLog("start uploading profile")
	if _, err := a.client.UpdateProfile(ctx, &req); err != nil {
		debugLog("failed to upload profile: %v", err)
	}
}

func deltaMutexProfile(ctx context.Context, duration time.Duration, prof *bytes.Buffer) error {
	if !mutexEnabled {
		return errors.New("mutex profiling is not enabled")
	}
	p0, err := mutexProfile()
	if err != nil {
		return err
	}
	sleep(ctx, duration)
	p, err := mutexProfile()
	if err != nil {
		return err
	}

	p0.Scale(-1)
	p, err = profile.Merge([]*profile.Profile{p0, p})
	if err != nil {
		return err
	}

	return p.Write(prof)
}

func mutexProfile() (*profile.Profile, error) {
	p := pprof.Lookup("mutex")
	if p == nil {
		return nil, errors.New("mutex profiling is not supported")
	}
	var buf bytes.Buffer
	if err := p.WriteTo(&buf, 0); err != nil {
		return nil, err
	}
	return profile.Parse(&buf)
}

func withXGoogHeader(ctx context.Context, keyval ...string) context.Context {
	kv := append([]string{"gl-go", version.Go(), "gccl", version.Repo}, keyval...)
	kv = append(kv, "gax", gax.Version, "grpc", grpc.Version)

	md, _ := grpcmd.FromOutgoingContext(ctx)
	md = md.Copy()
	md[xGoogAPIMetadata] = []string{gax.XGoogHeader(kv...)}
	return grpcmd.NewOutgoingContext(ctx, md)
}

func initializeAgent(c pb.ProfilerServiceClient) (*agent, error) {
	labels := map[string]string{languageLabel: "go"}
	if config.Zone != "" {
		labels[zoneNameLabel] = config.Zone
	}
	if config.ServiceVersion != "" {
		labels[versionLabel] = config.ServiceVersion
	}
	d := &pb.Deployment{
		ProjectId: config.ProjectID,
		Target:    config.Service,
		Labels:    labels,
	}

	profileLabels := map[string]string{}

	if config.Instance != "" {
		profileLabels[instanceLabel] = config.Instance
	}

	var profileTypes []pb.ProfileType
	if !config.NoCPUProfiling {
		profileTypes = append(profileTypes, pb.ProfileType_CPU)
	}
	if !config.NoHeapProfiling {
		profileTypes = append(profileTypes, pb.ProfileType_HEAP)
	}
	if !config.NoGoroutineProfiling {
		profileTypes = append(profileTypes, pb.ProfileType_THREADS)
	}
	if !config.NoAllocProfiling {
		profileTypes = append(profileTypes, pb.ProfileType_HEAP_ALLOC)
	}
	if mutexEnabled {
		profileTypes = append(profileTypes, pb.ProfileType_CONTENTION)
	}

	if len(profileTypes) == 0 {
		return nil, fmt.Errorf("collection is not enabled for any profile types")
	}

	return &agent{
		client:        c,
		deployment:    d,
		profileLabels: profileLabels,
		profileTypes:  profileTypes,
	}, nil
}

func initializeConfig(cfg Config) error {
	config = cfg

	if config.Service == "" {
		for _, ev := range []string{"GAE_SERVICE", "K_SERVICE"} {
			if val := os.Getenv(ev); val != "" {
				config.Service = val
				break
			}
		}
	}
	if config.Service == "" {
		return errors.New("service name must be configured")
	}
	if !serviceRegexp.MatchString(config.Service) {
		return fmt.Errorf("service name %q does not match regular expression %v", config.Service, serviceRegexp)
	}

	if config.ServiceVersion == "" {
		for _, ev := range []string{"GAE_VERSION", "K_REVISION"} {
			if val := os.Getenv(ev); val != "" {
				config.ServiceVersion = val
				break
			}
		}
	}

		config.ProjectID = projectID
	}
	if onGCE() {
		var err error
		if config.ProjectID == "" {
			if config.ProjectID, err = getProjectID(); err != nil {
				return fmt.Errorf("failed to get the project ID from Compute Engine metadata: %v", err)
			}
		}

		if config.Zone == "" {
			if config.Zone, err = getZone(); err != nil {
				return fmt.Errorf("failed to get zone from Compute Engine metadata: %v", err)
			}
		}

		if config.Instance == "" {
			if config.Instance, err = getInstanceName(); err != nil {
				if _, ok := err.(gcemd.NotDefinedError); !ok {
					return fmt.Errorf("failed to get instance name from Compute Engine metadata: %v", err)
				}
				debugLog("failed to get instance name from Compute Engine metadata, will use empty name: %v", err)
			}
		}
	} else {
		if config.ProjectID == "" {
			return fmt.Errorf("project ID must be specified in the configuration if running outside of GCP")
		}
	}

	if config.APIAddr == "" {
		config.APIAddr = apiAddress
	}
	return nil
}

func pollProfilerService(ctx context.Context, a *agent) {
	debugLog("Cloud Profiler Go Agent version: %s", version.Repo)
	debugLog("profiler has started")
	for i := 0; config.numProfiles == 0 || i < config.numProfiles; i++ {
		p := a.createProfile(ctx)
		a.profileAndUpload(ctx, p)
	}

	if profilingDone != nil {
		profilingDone <- true
	}