package ddb

import (
	"context"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"net/http"
	"strconv"
	"sync"
	"time"

	comcfg "dataservices/commons/config"

	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"

	"github.com/aws/aws-sdk-go-v2/aws"
	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/expression"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
	"github.com/aws/aws-sdk-go-v2/service/dynamodbstreams"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

const (
	batchSize     = 25  // DynamoDB batch write allows a maximum batch size of 25 items.
	batchReadSize = 100 // DynamoDB batch read allows a maximum batch size of 100 items.
)

var dynamodbInitOnce sync.Once
var dynamodbClient *dynamodb.Client
var streamsClient *dynamodbstreams.Client

func Init(ctx context.Context, cfg *comcfg.CentralConfig) error {
	dynamodbInitOnce.Do(func() {
		var awsConf *aws.Config
		if cfg.IsCentralMode() {
			awsConf = LoadAwsConfig(ctx, cfg.AwsConf.Region, "", "")
		} else {
			awsConf = LoadAwsConfig(ctx, cfg.AwsConf.Region, cfg.AwsConf.RoleArn, cfg.AwsConf.Session)
		}
		dynamodbClient = dynamodb.NewFromConfig(*awsConf)
		streamsClient = dynamodbstreams.NewFromConfig(*awsConf)
		tilog.Info(ctx, "successfully init dynamodb client",
			zap.Any("awsConfig", cfg.AwsConf),
			zap.String("deployMode", cfg.ServerConf.DeployMode))
	})

	return nil
}

func Client() *dynamodb.Client {
	return dynamodbClient
}

func StreamClient() *dynamodbstreams.Client {
	return streamsClient
}

func NewDynamodbClient(ctx context.Context, region string) (*dynamodb.Client, error) {
	cfg, err := config.LoadDefaultConfig(
		context.Background(),
		config.WithRegion(region),
	)
	if err != nil {
		return nil, err
	}
	return dynamodb.NewFromConfig(cfg), nil
}

func LoadAwsConfig(ctx context.Context, region, roleArn, roleSessionName string) *aws.Config {
	buildableClient := awshttp.NewBuildableClient().WithTransportOptions(func(t *http.Transport) {
		t.MaxIdleConnsPerHost = 150
		t.MaxIdleConns = 300
		// t.TLSClientConfig.InsecureSkipVerify = false
	})
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithHTTPClient(buildableClient))
	if err != nil {
		tilog.Fatal(ctx, "failed to load aws default config", zap.String("region", region), zap.Error(err))
	}

	if roleArn == "" || roleSessionName == "" {
		tilog.Info(ctx, "successfully loading aws config without assume role",
			zap.String("region", region),
			zap.String("role-arn", roleArn),
			zap.String("role-session-name", roleSessionName))
		return &cfg
	}

	stsClient := sts.NewFromConfig(cfg)
	durationSeconds := int32(60 * 60)
	expiryWindow := 30 * time.Second
	_, err = stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
		RoleArn:         &roleArn,
		RoleSessionName: &roleSessionName,
		DurationSeconds: &durationSeconds,
	})
	if err != nil {
		tilog.Fatal(ctx, "failed to assume aws role",
			zap.String("region", region),
			zap.String("role-arn", roleArn),
			zap.String("role-session-name", roleSessionName),
			zap.Int32("duration-second", durationSeconds),
			zap.Duration("expire-window", expiryWindow),
			zap.Error(err))
	}
	cfg.Credentials = aws.NewCredentialsCache(
		stscreds.NewAssumeRoleProvider(stsClient, roleArn),
		func(options *aws.CredentialsCacheOptions) {
			options.ExpiryWindow = expiryWindow
		},
	)
	tilog.Info(ctx, "successfully loading aws config with assume role",
		zap.String("region", region),
		zap.String("role-arn", roleArn),
		zap.String("role-session-name", roleSessionName),
		zap.Int32("duration-second", durationSeconds),
		zap.Duration("expire-window", expiryWindow))
	return &cfg
}

func NewAssumeDynamoDBClient(ctx context.Context, region, roleArn, roleSessionName string) (*dynamodb.Client, error) {
	cfg := LoadAwsConfig(ctx, region, roleArn, roleSessionName)
	ddbClient := dynamodb.NewFromConfig(*cfg)
	return ddbClient, nil
}

func CreateTable(ctx context.Context, cli *dynamodb.Client, createTableInput *dynamodb.CreateTableInput) (*types.TableDescription, error) {
	table, err := cli.CreateTable(ctx, createTableInput)
	if err != nil {
		tilog.Error(ctx, "Create table failed", zap.String("table", *createTableInput.TableName), zap.Error(err))
		return nil, err
	}
	tilog.Info(ctx, "Create table done, wait table exist", zap.String("table", *createTableInput.TableName))
	waiter := dynamodb.NewTableExistsWaiter(cli)
	err = waiter.Wait(ctx, &dynamodb.DescribeTableInput{
		TableName: createTableInput.TableName}, 5*time.Minute)
	if err != nil {
		tilog.Error(ctx, "Wait table exist failed", zap.String("table", *createTableInput.TableName), zap.Error(err))
		return nil, err
	}
	tilog.Info(ctx, "Create table finish, table is exist now", zap.String("table", *createTableInput.TableName))
	return table.TableDescription, nil
}

func IsTableExist(ctx context.Context, cli *dynamodb.Client, tableName string) (bool, error) {
	// Maybe use cli.DescribeTable api is better when there are to much tables.
	resp, err := cli.ListTables(ctx, &dynamodb.ListTablesInput{})
	if err != nil {
		return false, err
	}
	for _, t := range resp.TableNames {
		if t == tableName {
			return true, nil
		}
	}
	return false, nil
}

func DeleteTable(ctx context.Context, cli *dynamodb.Client, tableName string) error {
	_, err := cli.DeleteTable(ctx, &dynamodb.DeleteTableInput{TableName: aws.String(tableName)})
	if err != nil {
		tilog.Error(ctx, "Delete table failed", zap.String("table", tableName), zap.Error(err))
		return err
	}
	tilog.Info(ctx, "Delete table success", zap.String("table", tableName))
	return nil
}

func PutItem(ctx context.Context, cli *dynamodb.Client, tableName string, item interface{}) error {
	var err error
	defer func() {
		metrics.DDBCounter.WithLabelValues(strconv.FormatBool(err == nil)).Inc()
	}()
	data, err := attributevalue.MarshalMap(item)
	if err != nil {
		tilog.Error(ctx, "PutItem marshal map failed, should never happen", zap.String("table", tableName), zap.Any("item", item), zap.Error(err))
		return errs.NewInternalServerError("put item failed", err)
	}
	_, err = cli.PutItem(ctx, &dynamodb.PutItemInput{
		TableName: aws.String(tableName),
		Item:      data,
	})
	if err != nil {
		tilog.Error(ctx, "PutItem failed", zap.String("table", tableName), zap.Any("item", item), zap.Error(err))
		return errs.NewInternalServerError("put item failed", err)
	}
	return nil
}

func GetItem(ctx context.Context, cli *dynamodb.Client, tableName string, key map[string]types.AttributeValue, target any) error {
	res, err := cli.GetItem(ctx, &dynamodb.GetItemInput{
		TableName: aws.String(tableName),
		Key:       key,
	})
	if err != nil {
		tilog.Error(ctx, "GetItem failed", zap.String("table", tableName), zap.Any("key", key), zap.Error(err))
		return err
	} else if len(res.Item) == 0 {
		return errs.NewNotFound("item not exists", nil)
	}
	return attributevalue.UnmarshalMap(res.Item, target)
}

func BatchPutItem(ctx context.Context, cli *dynamodb.Client, tableName string, datas ...any) (int, error) {
	var (
		err     error
		item    map[string]types.AttributeValue
		written = 0
		start   = 0
		end     = start + batchSize
	)
	for start < len(datas) {
		var writeReqs []types.WriteRequest
		if end > len(datas) {
			end = len(datas)
		}
		for _, data := range datas[start:end] {
			item, err = attributevalue.MarshalMap(data)
			if err != nil {
				tilog.Error(ctx, "BatchPutItem failed to marshal data for batch writing", zap.String("table", tableName), zap.Any("data", data), zap.Error(err))
				continue
			}
			writeReqs = append(
				writeReqs,
				types.WriteRequest{PutRequest: &types.PutRequest{Item: item}},
			)
		}
		res, err := cli.BatchWriteItem(ctx, &dynamodb.BatchWriteItemInput{
			RequestItems: map[string][]types.WriteRequest{tableName: writeReqs}})
		if err != nil {
			tilog.Error(ctx, "BatchPutItem failed to add a batch of data to ddb table", zap.String("table", tableName), zap.Error(err))
		} else {
			written += len(writeReqs) - len(res.UnprocessedItems[tableName])
		}
		start = end
		end += batchSize
	}

	return written, err
}

func DeleteItem(ctx context.Context, cli *dynamodb.Client, tableName string, key map[string]types.AttributeValue) error {
	_, err := cli.DeleteItem(ctx, &dynamodb.DeleteItemInput{
		TableName: aws.String(tableName), Key: key,
	})
	if err != nil {
		tilog.Error(ctx, "DeleteItem failed", zap.String("table", tableName), zap.Error(err))
	}
	return err
}

func BatchDeleteItem(ctx context.Context, cli *dynamodb.Client, tableName string, keys ...map[string]types.AttributeValue) (int, error) {
	var (
		err     error
		written = 0
		start   = 0
		end     = start + batchSize
	)
	for start < len(keys) {
		var writeReqs []types.WriteRequest
		if end > len(keys) {
			end = len(keys)
		}
		for _, key := range keys[start:end] {
			writeReqs = append(
				writeReqs,
				types.WriteRequest{DeleteRequest: &types.DeleteRequest{Key: key}},
			)
		}
		res, err := cli.BatchWriteItem(ctx, &dynamodb.BatchWriteItemInput{
			RequestItems: map[string][]types.WriteRequest{tableName: writeReqs}})
		if err != nil {
			tilog.Error(ctx, "BatchDeleteItem failed to delete a batch of data from ddb table", zap.String("table", tableName), zap.Error(err))
		} else {
			written += len(writeReqs) - len(res.UnprocessedItems[tableName])
		}
		start = end
		end += batchSize
	}

	return written, err
}

func BatchGetItem(ctx context.Context, cli *dynamodb.Client, tableName string, keys ...map[string]types.AttributeValue) ([]map[string]types.AttributeValue, error) {
	var (
		err    error
		start  = 0
		end    = start + batchReadSize
		result []map[string]types.AttributeValue
	)
	for start < len(keys) {
		if end > len(keys) {
			end = len(keys)
		}
		keyAndAttrs := types.KeysAndAttributes{Keys: keys[start:end]}
		res, err := cli.BatchGetItem(ctx, &dynamodb.BatchGetItemInput{
			RequestItems: map[string]types.KeysAndAttributes{tableName: keyAndAttrs}})
		if err != nil {
			logger.Error(ctx, "BatchGetItem failed to get a batch of data from ddb table", zap.String("table", tableName), zap.Error(err))
		} else {
			result = append(result, res.Responses[tableName]...)
		}
		start = end
		end += batchReadSize
	}
	return result, err
}

func UpdateItem(ctx context.Context, cli *dynamodb.Client, tableName string, key map[string]types.AttributeValue, updateDdbMap map[string]interface{}) error {
	if len(updateDdbMap) == 0 {
		return nil
	}
	var update expression.UpdateBuilder
	i := 0
	for k, v := range updateDdbMap {
		if i == 0 {
			update = expression.Set(expression.Name(k), expression.Value(v))
		} else {
			update.Set(expression.Name(k), expression.Value(v))
		}
		i++
	}
	expr, err := expression.NewBuilder().WithUpdate(update).Build()
	if err != nil {
		return err
	}
	_, err = cli.UpdateItem(ctx, &dynamodb.UpdateItemInput{
		TableName:                 aws.String(tableName),
		Key:                       key,
		ExpressionAttributeNames:  expr.Names(),
		ExpressionAttributeValues: expr.Values(),
		UpdateExpression:          expr.Update(),
		ReturnValues:              types.ReturnValueUpdatedNew,
	})
	return err
}