package conprof

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"regexp"
	"sort"
	"strings"

	"github.com/pingcap/log"
	"go.uber.org/zap"

	"github.com/pingcap-inc/clinic/pkg/ngm/storage"
)

const maxTimeRangeSecs = 60 * 60 * 3 // 3h

var profileMetaRegexp = regexp.MustCompile(`^(\d+)-([a-z]+)-([a-z]+)-([a-zA-Z0-9]+)$`)

type ConProf interface {
	List(ctx context.Context, params *ListParams) ([]ProfileMeta, error)

	Download(ctx context.Context, params *DownloadParams) ([]byte, error)
}

type ClusterInfo struct {
	Provider   string
	Region     string
	TenantID   string
	ProjectID  string
	ClusterID  string
	DeployType string
}

type ListParams struct {
	ClusterInfo
	BeginTime int64
	EndTime   int64
}

type ProfileMeta struct {
	Ts          int64  // in secs
	Address     string // e.g. 127.0.0.1:10080
	Component   string // tidb/tikv/pd
	ProfileType string // cpu/heap/mutex/goroutine
}

type DownloadParams struct {
	ClusterInfo
	ProfileMeta
}

func NewConProf(env string, stores storage.StorageFactory) ConProf {
	return &conprof{
		env:    env,
		stores: stores,
	}
}

type conprof struct {
	env    string
	stores storage.StorageFactory
}

func (c *conprof) List(ctx context.Context, params *ListParams) ([]ProfileMeta, error) {
	if params.BeginTime <= 0 {
		return nil, errors.New("invalid begin time")
	}
	if params.EndTime <= 0 {
		return nil, errors.New("invalid end time")
	}
	store, err := c.stores.Storage(ctx, params.Provider, params.Region)
	if err != nil {
		return nil, err
	}
	begin := params.BeginTime - (params.BeginTime % 60)
	end := params.EndTime - (params.EndTime % 60)
	bucket := bucketName(c.env, params.Provider, params.Region)
	if end < begin {
		return nil, errors.New("invalid time range")
	}
	if begin+maxTimeRangeSecs < end {
		return nil, errors.New("time range too large")
	}
	if begin == end {
		prefix := commonKeyPrefixBetweenBeginEnd(params.TenantID, params.ProjectID, params.ClusterID, begin, end)
		objects, err := store.Objects(ctx, bucket, prefix)
		if err != nil {
			return nil, err
		}
		return extractProfileMetasFromObjectMetas(objects), nil
	}
	firstKey, err := tryFindFirstKey(ctx, store, bucket, params.TenantID, params.ProjectID, params.ClusterID, begin)
	if err != nil {
		return nil, err
	}
	var objects []storage.ObjectMeta
	if firstKey != "" && params.Provider != "azure" {
		dir := fmt.Sprintf("0/%s/%s/%s/profiles/", params.TenantID, params.ProjectID, params.ClusterID)
		objects, err = listObjectMetasFromFirstKey(ctx, store, bucket, dir, firstKey, end)
		if err != nil {
			return nil, err
		}
	} else {
		prefix := commonKeyPrefixBetweenBeginEnd(params.TenantID, params.ProjectID, params.ClusterID, begin, end)
		objects, err = listObjectMetasBetweenBeginEndCommonPrefix(ctx, store, bucket, prefix, begin, end)
		if err != nil {
			return nil, err
		}
	}
	metas := extractProfileMetasFromObjectMetas(objects)
	sort.Slice(metas, func(i, j int) bool {
		return metas[i].Ts > metas[j].Ts
	})
	return metas, nil
}

func (c *conprof) Download(ctx context.Context, params *DownloadParams) ([]byte, error) {
	store, err := c.stores.Storage(ctx, params.Provider, params.Region)
	if err != nil {
		return nil, err
	}
	bucket := bucketName(c.env, params.Provider, params.Region)
	log.Info("continuous_profiling download params", zap.Any("params", params))

	addressB64 := base64.RawURLEncoding.EncodeToString([]byte(params.Address))
	key := fmt.Sprintf("0/%s/%s/%s/profiles/%d-%s-%s-%s.log",
		params.TenantID, params.ProjectID, params.ClusterID,
		params.Ts, params.Component, params.ProfileType, addressB64)
	if strings.ToLower(params.Provider) == "aws" || strings.ToLower(params.Provider) == "azure" {
		key += ".gz"
	}
	object, err := store.Object(ctx, bucket, key)
	if err != nil {
		log.Error("continuous_profiling download failed", zap.String("file key", key), zap.Error(err))
		return nil, err
	}
	decoder := base64.NewDecoder(base64.StdEncoding, object.Reader)
	return io.ReadAll(decoder)
}

func tryFindFirstKey(
	ctx context.Context,
	store storage.Storage,
	bucket string,
	tenant, project, cluster string,
	begin int64,
) (string, error) {
	ts := begin
	for n := 0; n < 10; n++ {
		prefix := commonKeyPrefixBetweenBeginEnd(tenant, project, cluster, ts, ts)
		objects, err := store.Objects(ctx, bucket, prefix)
		if err != nil {
			return "", err
		}
		if len(objects) > 0 {
			return objects[0].Key, nil
		}
		ts += 60
	}
	return "", nil
}

func listObjectMetasFromFirstKey(
	ctx context.Context,
	store storage.Storage,
	bucket, dir, startKey string,
	end int64,
) ([]storage.ObjectMeta, error) {
	metas := []storage.ObjectMeta{}
	err := store.ObjectsCallbackWithStartKey(ctx, bucket, dir, startKey, func(meta storage.ObjectMeta) bool {
		filename := extractFilenameFromObjectKey(meta.Key)
		if filename == "" {
			log.Warn("invalid profile object key", zap.String("key", meta.Key))
			return false
		}
		ts := extractTsFromFilename(filename)
		if ts == 0 {
			log.Warn("invalid profile object key", zap.String("key", meta.Key))
			return false
		}
		if ts > end {
			return false // break
		}
		metas = append(metas, meta)
		return true // continue
	})
	if err != nil {
		return nil, err
	}
	return metas, nil
}

func listObjectMetasBetweenBeginEndCommonPrefix(
	ctx context.Context,
	store storage.Storage,
	bucket, prefix string,
	begin, end int64,
) ([]storage.ObjectMeta, error) {
	metas := []storage.ObjectMeta{}
	err := store.ObjectsCallback(ctx, bucket, prefix, func(meta storage.ObjectMeta) bool {
		filename := extractFilenameFromObjectKey(meta.Key)
		if filename == "" {
			log.Warn("invalid profile object key", zap.String("key", meta.Key))
			return false
		}
		ts := extractTsFromFilename(filename)
		if ts == 0 {
			log.Warn("invalid profile object key", zap.String("key", meta.Key))
			return false
		}
		if ts < begin {
			return true // continue
		}
		if ts > end {
			return false // break
		}
		metas = append(metas, meta)
		return true // continue
	})
	if err != nil {
		return nil, err
	}
	return metas, nil
}

func extractFilenameFromObjectKey(key string) string {
	begin := strings.LastIndexByte(key, '/')
	if begin != -1 {
		key = key[begin+1:]
	}
	end := strings.IndexByte(key, '.')
	if end != -1 {
		key = key[:end]
	}
	return key
}

func extractTsFromFilename(filename string) (ts int64) {
	_, _ = fmt.Sscanf(filename, "%d", &ts)
	return
}

func extractProfileMetaFromFilename(filename string) ProfileMeta {
	meta := ProfileMeta{}
	rs := profileMetaRegexp.FindAllStringSubmatch(filename, -1)
	if len(rs) != 1 || len(rs[0]) != 5 {
		return meta
	}
	_, _ = fmt.Sscanf(rs[0][1], "%d", &meta.Ts)
	meta.Component = rs[0][2]
	meta.ProfileType = rs[0][3]
	if rs[0][4] != "" {
		r, err := base64.RawURLEncoding.DecodeString(rs[0][4])
		if err == nil {
			meta.Address = string(r)
		}
	}
	return meta
}

func extractProfileMetasFromObjectMetas(metas []storage.ObjectMeta) []ProfileMeta {
	if len(metas) == 0 {
		return []ProfileMeta{}
	}
	rs := make([]ProfileMeta, 0, len(metas))
	for _, meta := range metas {
		filename := extractFilenameFromObjectKey(meta.Key)
		if filename == "" {
			continue
		}
		r := extractProfileMetaFromFilename(filename)
		if r.Ts == 0 || r.Component == "" || r.ProfileType == "" || r.Address == "" {
			continue
		}
		rs = append(rs, r)
	}
	return rs
}

func bucketName(env, provider, region string) string {
	switch env {
	case "prod":
		return fmt.Sprintf("o11y-prod-shared-%s", region)
	default:
		if provider == "aws" && region == "us-west-2" {
			return fmt.Sprintf("o11y-prod-shared-%s-staging", region)
		} else {
			return fmt.Sprintf("o11y-test-shared-%s", region)
		}
	}
}

func commonKeyPrefixBetweenBeginEnd(tenant, project, cluster string, begin, end int64) string {
	tsPrefix := longestCommonPrefix([]string{
		fmt.Sprintf("%d", begin),
		fmt.Sprintf("%d", end),
	})
	return fmt.Sprintf("0/%s/%s/%s/profiles/%s", tenant, project, cluster, tsPrefix)
}

func longestCommonPrefix(strs []string) string {
	lenStrs := len(strs)
	if lenStrs == 0 {
		return ""
	}
	firstString := strs[0]
	lenFirstString := len(firstString)
	commonPrefix := ""
	for i := 0; i < lenFirstString; i++ {
		firstStringChar := string(firstString[i])
		match := true
		for j := 1; j < lenStrs; j++ {
			if (len(strs[j]) - 1) < i {
				match = false
				break
			}
			if string(strs[j][i]) != firstStringChar {
				match = false
				break
			}
		}
		if match {
			commonPrefix += firstStringChar
		} else {
			break
		}
	}
	return commonPrefix
}