package awsquota

import (
	"context"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/servicequotas"
	"google.golang.org/api/iterator"
	"k8s.io/utils/ptr"
)

const (
	ebsServiceName      = "EBS"
	ebsServiceCode      = "ebs"
	io2IOPSQuotaCode    = "L-8D977E7E"
	io2StorageQuotaCode = "L-09BD8365"
	gp3StorageQuotaCode = "L-7A658B76"
	query               = `SELECT 
    region,
    volume_type,
    SUM(size) AS total_size,
    SUM(iops) AS total_iops,
    COALESCE(SUM(throughput), 0) AS total_throughput
FROM 
    ` + "`dbaas-eng-bi.resource_asset_snapshot.aws_ec2_ebs_volumes`" + ` 
WHERE
    TIMESTAMP_TRUNC(_cq_sync_time, DAY) = TIMESTAMP("%s") and _cq_sync_time = (
        SELECT MAX(_cq_sync_time)
        FROM dbaas-eng-bi.resource_asset_snapshot.aws_ec2_ebs_volumes
        WHERE DATE(_cq_sync_time) = CURRENT_DATE()
        ) and account_id = "%s" and (volume_type = "gp3" or volume_type = "io2")
GROUP BY
    region,
    volume_type
LIMIT 1000;`
)

type ebsVolumeRecord struct {
	Region     string `bigquery:"region"`
	VolumeType string `bigquery:"volume_type"`
	SizeInGiB  int    `bigquery:"total_size"`
	IOPS       int    `bigquery:"total_iops"`
	Throughput int    `bigquery:"total_throughput"`
}

func (r ebsVolumeRecord) GetUsage() (storageInTiB, iops, throughput float64) {
	return float64(r.SizeInGiB) / 1024, float64(r.IOPS), float64(r.Throughput)
}

func (c *quotaClient) queryEBSQuotas(ctx context.Context) []Quota {
	date := time.Now().Format("2006-01-02")
	q := fmt.Sprintf(query, date, c.account)
	it, err := c.bigQueryCli.Query(q).Read(ctx)
	if err != nil {
		c.logger.Error(err, "failed to execute big query sql", "query", q)
		return nil
	}

	var quotas []Quota
	for {
		var row ebsVolumeRecord
		if e := it.Next(&row); e != nil {
			if e != iterator.Done {
				c.logger.Error(e, "failed to get next row")
			}
			break
		} else {
			storageInTiB, iops, _ := row.GetUsage()
			switch row.VolumeType {
			case "io2":
				if quota := c.createEBSQuota(ctx, row.Region, ebsServiceCode, io2StorageQuotaCode, "Storage for Provisioned IOPS SSD (io2) volumes, in TiB", storageInTiB); quota != nil {
					quotas = append(quotas, *quota)
				}
				if quota := c.createEBSQuota(ctx, row.Region, ebsServiceCode, io2IOPSQuotaCode, "IOPS for Provisioned IOPS SSD (io2) volumes", iops); quota != nil {
					quotas = append(quotas, *quota)
				}
			case "gp3":
				if quota := c.createEBSQuota(ctx, row.Region, ebsServiceCode, gp3StorageQuotaCode, "Storage for General Purpose SSD (gp3) volumes, in TiB", storageInTiB); quota != nil {
					quotas = append(quotas, *quota)
				}
			}
		}
	}
	return quotas
}

// getServiceQuotaLimitValue returns the limit value of the given quota code in the given region.
func (c *quotaClient) getServiceQuotaLimitValue(ctx context.Context, region, serviceCode, quotaCode string) (float64, error) {
	output, err := c.svcQuota.GetServiceQuota(ctx, &servicequotas.GetServiceQuotaInput{
		QuotaCode:   ptr.To(quotaCode),
		ServiceCode: ptr.To(serviceCode),
	}, func(opt *servicequotas.Options) { opt.Region = region })
	if err != nil || output.Quota == nil || output.Quota.Value == nil {
		return 0, fmt.Errorf("failed to get limit for service code %s, quota code %s in region %s: %w", serviceCode, quotaCode, region, err)
	}
	return *output.Quota.Value, nil
}

func (c *quotaClient) createEBSQuota(ctx context.Context, region, serviceCode, quotaCode, resource string, usage float64) *Quota {
	if usage <= 0 {
		return nil
	}
	limit, err := c.getServiceQuotaLimitValue(ctx, region, serviceCode, quotaCode)
	if err != nil {
		c.logger.Error(err, "failed to get quota limit", "region", region, "resource", resource)
		return nil
	}
	return &Quota{
		Region:   region,
		Service:  ebsServiceName,
		Resource: resource,
		Limit:    limit,
		Usage:    usage,
	}
}