package billing

import (
	"fmt"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/athena"
	"go.uber.org/zap"

	. "github.com/tidbcloud/devtier-store-service/utils"
)

func GetAwsServerlessCostFromAthena(logger *zap.Logger) (date, cost string) {
	logger.Info("GetAwsServerlessCostFromAthena start...")
	defer logger.Info("GetAwsServerlessCostFromAthena end...")

	sess := session.Must(session.NewSession())
	creds := stscreds.NewCredentials(sess, Role)
	svc := athena.New(sess, &aws.Config{Region: aws.String("us-west-2"), Credentials: creds})

	var s athena.StartQueryExecutionInput
	s.SetQueryString("SELECT to_iso8601(current_date - interval '2' day) as date," +
		"round(sum(line_item_unblended_cost), 2) as cost " +
		"FROM athenacurcfn_" + Env + "_billing." + Env + "_billing " +
		"WHERE (resource_tags_user_servicetype = 'devtier' or resource_tags_user_servicetype = 'serverless') " +
		"and to_iso8601(line_item_usage_start_date) <= to_iso8601(current_date - interval '1' day) " +
		"and to_iso8601(line_item_usage_start_date) >= to_iso8601(current_date - interval '2' day)")
	var q athena.QueryExecutionContext
	q.SetDatabase("athenacurcfn_" + Env + "_billing")
	s.SetQueryExecutionContext(&q)

	var r athena.ResultConfiguration
	r.SetOutputLocation("s3://billing-yys/dest")
	s.SetResultConfiguration(&r)

	result, err := svc.StartQueryExecution(&s)
	if err != nil {
		logger.Error("svc.StartQueryExecution failed, env=%s,role=%s,err=%v", zap.String(Env, Role), zap.Error(err))
		return
	}
	logger.Info("svc.StartQueryExecution success", zap.Any("result", result.GoString()))

	var qri athena.GetQueryExecutionInput
	qri.SetQueryExecutionId(*result.QueryExecutionId)

	var qrop *athena.GetQueryExecutionOutput
	duration := time.Duration(2) * time.Second // Pause for 2 seconds

	for {
		qrop, err = svc.GetQueryExecution(&qri)
		if err != nil {
			logger.Error("svc.GetQueryExecution failed", zap.Error(err))
			return
		}
		if *qrop.QueryExecution.Status.State == "FAILED" ||
			*qrop.QueryExecution.Status.State == "CANCELLED" {
			logger.Error("svc.GetQueryExecution failed, =%v", zap.Any("result", *qrop.QueryExecution))
			return
		}
		if *qrop.QueryExecution.Status.State != "RUNNING" &&
			*qrop.QueryExecution.Status.State != "QUEUED" {
			logger.Info("svc.GetQueryExecution end", zap.Any("state", *qrop.QueryExecution.Status.State))
			break
		}
		logger.Info("waiting for ...", zap.Any("duration", duration))
		time.Sleep(duration)

	}
	logger.Info("svc.GetQueryExecution success", zap.Any("state", *qrop.QueryExecution.Status.State))
	if *qrop.QueryExecution.Status.State != "SUCCEEDED" {
		logger.Error("svc.GetQueryExecution failed", zap.Any("result", *qrop.QueryExecution))
		return
	}

	var ip athena.GetQueryResultsInput
	ip.SetQueryExecutionId(*result.QueryExecutionId)

	op, err := svc.GetQueryResults(&ip)
	if err != nil {
		logger.Error("svc.GetQueryResults failed", zap.Error(err))
		return
	}
	if (*op.ResultSet.Rows[1].Data[1] == athena.Datum{}) {
		logger.Error("svc.GetQueryResults failed", zap.Any("error", *op))
		return
	}
	return *op.ResultSet.Rows[1].Data[0].VarCharValue, *op.ResultSet.Rows[1].Data[1].VarCharValue
}

func GetAwsServerlessCostByRegion(logger *zap.Logger, offset uint32) (regionCost []ServerlessTierBilling) {
	//logger := otelzap.L().Ctx(ctx)
	logger.Info("GetAwsServerlessCostByRegion start...")
	defer logger.Info("GetAwsServerlessCostByRegion end...")
	//offset := 2

	sess := session.Must(session.NewSession())
	creds := stscreds.NewCredentials(sess, Role)
	svc := athena.New(sess, &aws.Config{Region: aws.String("us-west-2"), Credentials: creds})

	var s athena.StartQueryExecutionInput
	s.SetQueryString("SELECT to_iso8601(current_date - interval '" + fmt.Sprint(offset) + "' day) as date," +
		"product_region as region, " +
		"round(sum(line_item_unblended_cost), 2) as cost " +
		"FROM athenacurcfn_" + Env + "_billing." + Env + "_billing " +
		"WHERE (resource_tags_user_servicetype = 'devtier' or resource_tags_user_servicetype = 'serverless') " +
		"and to_iso8601(line_item_usage_start_date) <= to_iso8601(current_date - interval '" + fmt.Sprint(offset-1) + "' day) " +
		"and to_iso8601(line_item_usage_start_date) >= to_iso8601(current_date - interval '" + fmt.Sprint(offset) + "' day) " +
		"group by product_region")
	var q athena.QueryExecutionContext
	q.SetDatabase("athenacurcfn_" + Env + "_billing")
	s.SetQueryExecutionContext(&q)

	var r athena.ResultConfiguration
	r.SetOutputLocation("s3://billing-yys/dest")
	s.SetResultConfiguration(&r)

	result, err := svc.StartQueryExecution(&s)
	if err != nil {
		logger.Error("svc.StartQueryExecution failed, env,role,err", zap.Any(Env, Role), zap.Error(err))
		return
	}
	logger.Info("svc.StartQueryExecution success", zap.Any("result", result.GoString()))

	var qri athena.GetQueryExecutionInput
	qri.SetQueryExecutionId(*result.QueryExecutionId)

	var qrop *athena.GetQueryExecutionOutput
	duration := time.Duration(2) * time.Second // Pause for 2 seconds

	for {
		qrop, err = svc.GetQueryExecution(&qri)
		if err != nil {
			logger.Error("svc.GetQueryExecution failed", zap.Error(err))
			return
		}
		if *qrop.QueryExecution.Status.State == "FAILED" ||
			*qrop.QueryExecution.Status.State == "CANCELLED" {
			logger.Error("svc.GetQueryExecution failed", zap.Any("result", *qrop.QueryExecution))
			return
		}
		if *qrop.QueryExecution.Status.State != "RUNNING" &&
			*qrop.QueryExecution.Status.State != "QUEUED" {
			logger.Info("svc.GetQueryExecution end", zap.Any("state", *qrop.QueryExecution.Status.State))
			break
		}
		logger.Info("waiting for ...", zap.Any("duration", duration))
		time.Sleep(duration)

	}
	logger.Info("svc.GetQueryExecution success", zap.Any("state", *qrop.QueryExecution.Status.State))
	if *qrop.QueryExecution.Status.State != "SUCCEEDED" {
		logger.Error("svc.GetQueryExecution failed", zap.Any("result", *qrop.QueryExecution))
		return
	}

	var ip athena.GetQueryResultsInput
	ip.SetQueryExecutionId(*result.QueryExecutionId)

	op, err := svc.GetQueryResults(&ip)
	if err != nil {
		logger.Error("svc.GetQueryResults failed", zap.Error(err))
		return
	}
	for id, v := range op.ResultSet.Rows {
		if id == 0 || len(v.Data) != 3 {
			continue
		}
		var data ServerlessTierBilling
		data.Day = *v.Data[0].VarCharValue
		data.Provider = "aws"
		data.Region = *v.Data[1].VarCharValue
		data.Cost, _ = strconv.ParseFloat(*v.Data[2].VarCharValue, 64)
		regionCost = append(regionCost, data)
	}
	return regionCost
}