package tasks

import (
	"bytes"
	"context"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/shopspring/decimal"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protojson"
	"gopkg.in/yaml.v3"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/selection"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"stathat.com/c/consistent"

	"github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	observabilitypingcapcomv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/pkg/diagcontroller/config"
	"github.com/tidbcloud/observability/pkg/diagcontroller/helper"
	diagstorage "github.com/tidbcloud/observability/pkg/diagcontroller/storage"
)

type DiagTask struct {
	conf    *config.DiagConfig
	kubeCli client.Client
	store   diagstorage.Storage
}

func NewDiagTask(
	kubeCli client.Client,
	conf *config.DiagConfig,
	store diagstorage.Storage,
) Tasker {
	return &DiagTask{
		conf:    conf,
		kubeCli: kubeCli,
		store:   store,
	}
}

func (d *DiagTask) Interval() time.Duration {
	return d.conf.Task.Interval
}

func (d *DiagTask) DispatchTask() error {
	clusters := make([]ClusterInfo, 0)
	if d.conf.Task.EnableDedicated {
		dedicated, err := d.listClusters([]string{"tidbcloud"}, DedicatedCluster)
		if err != nil {
			return err
		}
		clusters = append(clusters, dedicated...)
	}
	if d.conf.Task.EnableServerless {
		serverlessVIP, err := d.listServerlessVipClusters()
		if err != nil {
			log.Error("list serverless vip clusters failed",
				zap.Error(err),
			)
		} else {
			clusters = append(clusters, serverlessVIP...)
		}
	}

	diagEngines, err := helper.ListAvailableEnginePods(d.kubeCli, d.conf.DiagEngine)
	if err != nil {
		return err
	}

	if diagEngines == nil || len(diagEngines.Items) == 0 {
		log.Info("no available diag-engine pod, skip")
		return nil
	}

	tasks, err := d.syncTasks(clusters)
	if err != nil {
		return err
	}
	catalog := generateCatalog(tasks, diagEngines)
	if err := d.syncCatalogFile(catalog); err != nil {
		return err
	}

	return nil
}

func (d *DiagTask) syncTasks(clusters []ClusterInfo) ([]*Task, error) {
	tasks := make([]*Task, 0, len(clusters))
	// TODO: batch sync cluster task
	for _, cluster := range clusters {
		task, err := d.syncClusterTaskFile(cluster)
		if err != nil {
			log.Error("sync cluster task file to s3 failed",
				zap.Error(err),
				zap.String("Cluster", cluster.ID),
			)
			continue
		}
		if task != nil {
			tasks = append(tasks, task)
		}
	}
	return tasks, nil
}

func (d *DiagTask) listClusters(bizType []string, clusterType string) ([]ClusterInfo, error) {
	clusters := &observabilitypingcapcomv1.ClusterList{}
	specify, err := labels.NewRequirement("observability.pingcap.com/bizType", selection.In, bizType)
	if err != nil {
		return nil, err
	}
	ls := labels.NewSelector().Add(*specify)
	opts := &client.ListOptions{
		LabelSelector: ls,
	}
	if err := d.kubeCli.List(context.Background(), clusters, opts); err != nil {
		return nil, err
	}

	clusterInfos := make([]ClusterInfo, 0, len(clusters.Items))
	for _, cluster := range clusters.Items {
		if id, ok := cluster.GetLabels()["observability.pingcap.com/applicationID"]; ok {
			projectID, ok := cluster.GetLabels()["observability.pingcap.com/projectID"]
			if !ok {
				log.Warn("cluster has no project id", zap.String("Cluster", id))
			}
			if _, err := strconv.Atoi(id); err == nil {
				clusterInfos = append(clusterInfos, ClusterInfo{
					ProjectID: projectID,
					Type:      clusterType,
					ID:        id,
				})
			}
		}
	}

	return clusterInfos, nil
}

func (d *DiagTask) listServerlessVipClusters() ([]ClusterInfo, error) {
	data, err := d.store.LoadObject(d.conf.Task.Bucket, d.conf.Task.ServerlessVIPListPath)
	if err != nil {
		return nil, err
	}
	records, err := csv.NewReader(bytes.NewBuffer(data)).ReadAll()
	if err != nil {
		return nil, err
	}

	spendLimits, err := d.listServerlessSpendLimits()
	if err != nil {
		log.Error("list serverless spend limits failed", zap.Error(err))
	}

	ruPrices, err := d.listServerlessRuPrices()
	if err != nil {
		log.Error("list serverless ru prices failed", zap.Error(err))
	}

	clusterInfos := make([]ClusterInfo, 0, len(records))
	for _, raw := range records {
		if len(raw) < 2 {
			continue
		}
		// skip header
		if raw[0] == "cluster_id" {
			continue
		}

		limit, waterMark, freeCredits := decimal.Zero, decimal.Zero, decimal.Zero
		if spendLimits != nil {
			if setting, ok := spendLimits[raw[0]]; ok {
				waterMark = setting.WaterMark
				var bizData ServerlessBizData
				if err := json.Unmarshal([]byte(setting.BizData), &bizData); err == nil {
					limit = bizData.InvoiceLimit
					freeCredits = bizData.FreeCredits
				} else {
					log.Error("unmarshal serverless biz data failed", zap.String("cluster", raw[0]), zap.Error(err))
				}
			} else {
				log.Error("cluster has no spend limit setting", zap.String("cluster", raw[0]))
			}
		} else {
			log.Error("spend limit setting is nil")
		}

		info := ClusterInfo{
			Type:        ServerlessCluster,
			ID:          raw[0],
			KeyspaceID:  raw[1],
			SpendLimit:  limit,
			WaterMark:   waterMark,
			FreeCredits: freeCredits,
		}

		if ruPrices != nil {
			if cost, ok := ruPrices[ROW_STORAGE_TYPE]; ok {
				info.RowStoragePrice = cost.CostPrice
			}
			if cost, ok := ruPrices[COLUMN_STORAGE_TYPE]; ok {
				info.ColumnStoragePrice = cost.CostPrice
			}
			if cost, ok := ruPrices[REQUEST_TYPE]; ok {
				info.RequestUnitPrice = cost.CostPrice
			}
		}

		clusterInfos = append(clusterInfos, info)
	}

	return clusterInfos, nil
}

func (d *DiagTask) listServerlessSpendLimits() (map[string]*SpendLimitSetting, error) {
	data, err := d.store.LoadObject(d.conf.Task.Bucket, d.conf.Task.ServerlessSpendLimitsPath)
	if err != nil {
		return nil, err
	}

	settings := make([]*SpendLimitSetting, 0, 10)
	if err := json.Unmarshal(data, &settings); err != nil {
		return nil, err
	}

	output := make(map[string]*SpendLimitSetting, len(settings))
	for _, setting := range settings {
		output[strconv.FormatUint(setting.RangeID, 10)] = setting
	}

	return output, nil
}

func (d *DiagTask) listServerlessRuPrices() (map[string]*UnitCost, error) {
	data, err := d.store.LoadObject(d.conf.Task.Bucket, d.conf.Task.ServerlessSkuPricePath)
	if err != nil {
		return nil, err
	}

	prices := make([]*UnitCost, 0, 3)
	if err := json.Unmarshal(data, &prices); err != nil {
		return nil, err
	}

	output := make(map[string]*UnitCost, len(prices))
	for _, price := range prices {
		output[getSkuType(price.Unit)] = price
	}
	return output, nil
}

const (
	ROW_STORAGE_TYPE    = "serverless-row-storage"
	COLUMN_STORAGE_TYPE = "serverless-column-storage"
	REQUEST_TYPE        = "serverless-request-unit"
	UNKNOWN_TYPE        = "unknown"
)

func getSkuType(unit string) string {
	if strings.HasPrefix(unit, "serverless-row-storage") {
		return ROW_STORAGE_TYPE
	}
	if strings.HasPrefix(unit, "serverless-column-storage") {
		return COLUMN_STORAGE_TYPE
	}
	if strings.HasPrefix(unit, "serverless-request-unit") {
		return REQUEST_TYPE
	}
	return UNKNOWN_TYPE
}

func (d *DiagTask) syncClusterTaskFile(cluster ClusterInfo) (*Task, error) {
	meta, err := d.getClusterMeta(cluster.ID)
	if err != nil {
		return nil, err
	}

	if !isNormalCluster(meta) {
		log.Info("cluster is not normal, skip sync diag task file", zap.String("cluster", cluster.ID))
		return nil, nil
	}

	task := generateTask(cluster, d.conf.Task, meta, d.conf.Provider)

	objectBytes, err := yaml.Marshal(task)
	if err != nil {
		return nil, err
	}
	if err := d.store.UploadObject(d.conf.Task.Bucket, task.FilePath, objectBytes); err != nil {
		return nil, errors.Errorf("failed to sync task file %s to s3: %s", task.FilePath, err.Error())
	}
	log.Info("sync task file", zap.String("cluster", cluster.ID), zap.String("file", task.FilePath))
	return task, nil
}

func generateTask(base ClusterInfo, conf config.TaskConfig, meta *core.Cluster, provider string) *Task {
	ts := &Task{
		Interval: conf.Interval,
		Validity: ValidityPeriod{
			Type: LongTermTask,
		},
		Metas: Meta{
			ClusterID:          base.ID,
			ClusterType:        base.Type,
			KeyspaceID:         base.KeyspaceID,
			SpendLimit:         base.SpendLimit,
			WaterMark:          base.WaterMark,
			FreeCredits:        base.FreeCredits,
			Provider:           provider,
			RowStoragePrice:    base.RowStoragePrice,
			ColumnStoragePrice: base.ColumnStoragePrice,
			RequestUnitPrice:   base.RequestUnitPrice,
		},
		FilePath: filepath.Join(conf.TaskFilePrefix, base.ID+".yaml"),
	}
	if meta != nil {
		ts.Metas.DatabaseVersion = meta.Version
		ts.Metas.CreateTime = time.Unix(meta.CreatedAt, 0).Format("2006-01-02 15:04:05")
		for _, comp := range meta.Components {
			ts.Metas.Topology = append(ts.Metas.Topology, Component{
				Name:     componentName[comp.ComponentType],
				Replicas: comp.Replicas,
				CPU:      comp.CpuNum,
				Memory:   fmt.Sprint(comp.MemoryGi, "GiB"),
				Storage:  fmt.Sprint(comp.StorageSizeGi, "GiB"),
			})
		}
	}
	return ts
}

func (d *DiagTask) getClusterMeta(id string) (*core.Cluster, error) {
	key := filepath.Join(d.conf.Task.ClusterMetaFilePrefix, id+".json")
	data, err := d.store.LoadObject(d.conf.Task.Bucket, key)
	if err != nil {
		return nil, err
	}

	cluster := &core.Cluster{}
	err = protojson.UnmarshalOptions{DiscardUnknown: true}.Unmarshal(data, cluster)
	if err != nil {
		return nil, errors.Annotate(err, "decode cluster meta failed")
	}

	return cluster, nil
}

func (d *DiagTask) syncCatalogFile(catalog *Catalog) error {
	catalogBytes, err := yaml.Marshal(catalog)
	if err != nil {
		return err
	}

	if err := d.store.UploadObject(d.conf.Task.Bucket, d.conf.Task.CatalogPath, catalogBytes); err != nil {
		return err
	}

	log.Info("sync catalog file to s3 successfully", zap.String("catalog filepath", d.conf.Task.CatalogPath), zap.String("catalog", string(catalogBytes)))

	return nil
}

func isNormalCluster(meta *core.Cluster) bool {
	if meta.Status != core.ClusterStatus_CLUSTER_STATUS_NORMAL && meta.Status != core.ClusterStatus_CLUSTER_STATUS_SCALING {
		return false
	}

	return true
}

func generateCatalog(tasks []*Task, engines *corev1.PodList) *Catalog {
	catalog := &Catalog{
		Tasks: make(map[string][]string, len(engines.Items)),
	}

	if len(engines.Items) > 0 {
		ch := consistent.New()
		for _, engine := range engines.Items {
			catalog.Tasks[engine.Name] = make([]string, 0)
			ch.Add(engine.Name)
		}
		for _, task := range tasks {
			engine, err := ch.Get(task.Metas.ClusterID)
			if err != nil {
				log.Error("get engine name failed", zap.Error(err))
				continue
			}
			catalog.Tasks[engine] = append(catalog.Tasks[engine], task.FilePath)
		}
	}

	return catalog
}

// Task is the task to be executed by diag-engine
type Task struct {
	Interval   time.Duration  `json:"interval" yaml:"interval"`
	LastReport string         `json:"lastReport" yaml:"lastReport"`
	Validity   ValidityPeriod `json:"validityPeriod" yaml:"validityPeriod"`
	Metas      Meta           `json:"metas" yaml:"metas"`
	FilePath   string         `json:"filePath" yaml:"filePath"`
}

type Meta struct {
	ClusterID          string          `json:"clusterID" yaml:"clusterID"`
	DatabaseVersion    string          `json:"databaseVersion" yaml:"databaseVersion"`
	CreateTime         string          `json:"createTime" yaml:"createTime"`
	Topology           []Component     `json:"topology" yaml:"topology"`
	ClusterType        string          `json:"clusterType" yaml:"clusterType"`
	KeyspaceID         string          `json:"keyspaceID" yaml:"keyspaceID"`
	SpendLimit         decimal.Decimal `json:"spendLimit" yaml:"spendLimit"`
	WaterMark          decimal.Decimal `json:"waterMark" yaml:"waterMark"`
	FreeCredits        decimal.Decimal `json:"freeCredits" yaml:"freeCredits"`
	Provider           string          `json:"provider" yaml:"provider"`
	RowStoragePrice    decimal.Decimal `json:"rowStoragePrice" yaml:"rowStoragePrice"`
	ColumnStoragePrice decimal.Decimal `json:"columnStoragePrice" yaml:"columnStoragePrice"`
	RequestUnitPrice   decimal.Decimal `json:"requestUnitPrice" yaml:"requestUnitPrice"`
}

type Component struct {
	Name     string `json:"name" yaml:"name"`
	CPU      int32  `json:"cpu" yaml:"cpu"`
	Memory   string `json:"memory" yaml:"memory"`
	Storage  string `json:"storage" yaml:"storage"`
	Replicas int32  `json:"replicas" yaml:"replicas"`
}

var componentName = map[core.ComponentType]string{
	core.ComponentType_COMPONENT_TYPE_PD:      "pd",
	core.ComponentType_COMPONENT_TYPE_TIKV:    "tikv",
	core.ComponentType_COMPONENT_TYPE_TIDB:    "tidb",
	core.ComponentType_COMPONENT_TYPE_TIFLASH: "tiflash",
}

const LongTermTask = "long-term"

type ValidityPeriod struct {
	Type string `json:"type" yaml:"type"`
}

const (
	ServerlessCluster = "serverless"
	DedicatedCluster  = "dedicated"
	O11YCluster       = "o11y"
)

// ClusterInfo Record the base information of the diagnosed cluster
type ClusterInfo struct {
	ProjectID             string          // TiDB Cloud project id
	Type                  string          // Cluster type : dedicated, serverless, o11y
	ID                    string          // Cluster id
	DiagnosticRuleVersion string          // diagnostic rule version
	InspectionRuleVersion string          // inspection rule version
	KeyspaceID            string          // Keyspace id
	SpendLimit            decimal.Decimal // Spend limit
	WaterMark             decimal.Decimal // Water mark
	FreeCredits           decimal.Decimal // Free quota
	RowStoragePrice       decimal.Decimal // Row unit price
	ColumnStoragePrice    decimal.Decimal // Column unit price
	RequestUnitPrice      decimal.Decimal // Request unit price
}

// Catalog is the catalog of diag-engine, which contains all the diag tasks
type Catalog struct {
	// Tasks is the map of diag tasks,
	// key is the diag-engine name and value is the tasks to be executed by the diag-engine
	Tasks map[string][]string `json:"tasks" yaml:"tasks"`
}

func (c *Catalog) TasksCount() int {
	count := 0
	for _, tasks := range c.Tasks {
		count += len(tasks)
	}
	return count
}

type SpendLimitSetting struct {
	SpendLimitSettingID uint64          `json:"spendlimit_setting_id"`
	TenantID            uint64          `json:"tenant_id"`
	Interval            string          `json:"interval"`
	CostType            string          `json:"cost_type"`
	RangeType           string          `json:"range_type"`
	RangeID             uint64          `json:"range_id"`
	Limit               decimal.Decimal `json:"limit"`
	Status              string          `json:"status"`
	BizType             string          `json:"biz_type"`
	BizData             string          `json:"biz_data"`
	WaterMark           decimal.Decimal `json:"water_mark"`
}

type ServerlessBizData struct {
	FreeCredits          decimal.Decimal `json:"free_credits"`
	InvoiceLimit         decimal.Decimal `json:"invoice_limit"`
	ActiveTotalCredits   decimal.Decimal `json:"active_total_credits"`
	RemainingCredits     decimal.Decimal `json:"remaining_credits"`
	RechargePlanID       uint            `json:"recharge_plan_id"`
	TenantRechargePlanID uint            `json:"tenant_recharge_plan_id"`
	TieredBillingID      uint            `json:"tiered_billing_id"`
}

type UnitCost struct {
	CostID          string            `json:"cost_id"`
	CostTimeStamp   int64             `json:"cost_time_stamp"`
	Unit            string            `json:"unit"`
	Cost            int               `json:"cost"`
	Labels          map[string]string `json:"labels"`
	CostPrice       decimal.Decimal   `json:"price"`
	RawPrice        decimal.Decimal   `json:"raw_price"`
	PackageID       uint              `json:"package_id"`
	PackageVersion  uint              `json:"package_version"`
	BindedPackageID uint              `json:"binded_package_id"`
	ConsumedSource  string            `json:"consumed_source"`
	InvoiceID       int               `json:"invoice_id"`
	Discount        decimal.Decimal   `json:"discount"`
	Internal        bool              `json:"internal"`
}