package aws

import (
	"context"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	ec2api "github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	stsapi "github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/tidbcloud/observability/pkg/common/metrics"
)

const (
	privateEndpointClientLabel = "private-endpoint-client"
	errDescribeSubnetsFailed   = "describe subnets failed"
	errCreateVpcEndpointFailed = "create vpc endpoint failed"
	errDeleteVpcEndpointFailed = "delete vpc endpoint failed"
)

type PrivateEndpointClient struct {
	cli       privateEndpointClient
	cfg       aws.Config
	rolecache map[string]*aws.CredentialsCache
}

func NewPrivateEndpoint(cfg aws.Config) PrivateEndpointClient {
	c := ec2api.NewFromConfig(cfg)
	return PrivateEndpointClient{
		cli:       c,
		cfg:       cfg,
		rolecache: make(map[string]*aws.CredentialsCache),
	}
}

func (p PrivateEndpointClient) AssumeRole(ctx context.Context, roleArn string) *PrivateEndpointClient {
	//check if the role is already in the cache
	if cred, ok := p.rolecache[roleArn]; ok {
		return &PrivateEndpointClient{
			cli: ec2api.NewFromConfig(p.cfg, func(o *ec2api.Options) {
				o.Credentials = cred
			}),
			cfg: p.cfg,
		}
	}
	newCred := aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(stsapi.NewFromConfig(p.cfg), roleArn, func(p *stscreds.AssumeRoleOptions) {
		p.RoleSessionName = sessionName
	}))
	p.rolecache[roleArn] = newCred
	return &PrivateEndpointClient{
		cli: ec2api.NewFromConfig(p.cfg, func(o *ec2api.Options) {
			o.Credentials = newCred
		}),
		cfg: p.cfg,
	}
}

func (p PrivateEndpointClient) DescribeSubnets(ctx context.Context, subnetIDs []string) (*ec2api.DescribeSubnetsOutput, error) {
	subnets, err := p.cli.DescribeSubnets(ctx, &ec2api.DescribeSubnetsInput{
		SubnetIds: subnetIDs,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(p.cfg.Region, privateEndpointClientLabel, "DescribeSubnets", strings.Join(subnetIDs, ","), errDescribeSubnetsFailed)
		return nil, err
	}
	return subnets, nil
}

func (p PrivateEndpointClient) DescribeVpcEndpointServices(ctx context.Context, serviceNames []string) (*ec2api.DescribeVpcEndpointServicesOutput, error) {
	services, err := p.cli.DescribeVpcEndpointServices(ctx, &ec2api.DescribeVpcEndpointServicesInput{
		ServiceNames: serviceNames,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(p.cfg.Region, privateEndpointClientLabel, "DescribeVpcEndpointServices", strings.Join(serviceNames, ","), errDescribeSubnetsFailed)
		return nil, err
	}
	return services, nil
}

func (p PrivateEndpointClient) DescribeVpcEndpoints(ctx context.Context, endpointIDs []string) (*ec2api.DescribeVpcEndpointsOutput, error) {
	endpoints, err := p.cli.DescribeVpcEndpoints(ctx, &ec2api.DescribeVpcEndpointsInput{
		VpcEndpointIds: endpointIDs,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(p.cfg.Region, privateEndpointClientLabel, "DescribeVpcEndpoints", strings.Join(endpointIDs, ","), errDescribeSubnetsFailed)
		return nil, err
	}
	return endpoints, nil
}

func (p PrivateEndpointClient) DescribeVpcEndpointsByName(ctx context.Context, name string) (*ec2api.DescribeVpcEndpointsOutput, error) {
	endpoints, err := p.cli.DescribeVpcEndpoints(ctx, &ec2api.DescribeVpcEndpointsInput{
		Filters: []ec2types.Filter{
			{
				Name:   aws.String("tag:Name"),
				Values: []string{name},
			},
		},
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(p.cfg.Region, privateEndpointClientLabel, "DescribeVpcEndpointsByName", name, errDescribeSubnetsFailed)
		return nil, err
	}
	return endpoints, nil
}

func (p PrivateEndpointClient) CreateVpcEndpoint(ctx context.Context, name string, vpcID string, serviceName string, subnetIDs []string, endpointType string, secretGroupIDs []string, tags map[string]string) (*ec2api.CreateVpcEndpointOutput, error) {
	tagspec := []ec2types.TagSpecification{
		{
			ResourceType: ec2types.ResourceTypeVpcEndpoint,
			Tags: []ec2types.Tag{
				{
					Key:   aws.String("Name"),
					Value: aws.String(name),
				},
			},
		},
	}
	for k, v := range tags {
		tagspec[0].Tags = append(tagspec[0].Tags, ec2types.Tag{
			Key:   aws.String(k),
			Value: aws.String(v),
		})
	}
	endpoint, err := p.cli.CreateVpcEndpoint(ctx, &ec2api.CreateVpcEndpointInput{
		VpcId:             aws.String(vpcID),
		ServiceName:       aws.String(serviceName),
		SubnetIds:         subnetIDs,
		VpcEndpointType:   ec2types.VpcEndpointType(endpointType),
		SecurityGroupIds:  secretGroupIDs,
		TagSpecifications: tagspec,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(p.cfg.Region, privateEndpointClientLabel, "CreateVpcEndpoint", name, errCreateVpcEndpointFailed)
		return nil, err
	}
	return endpoint, nil
}

func (p PrivateEndpointClient) DeleteVpcEndpoints(ctx context.Context, endpointIDs []string) (*ec2api.DeleteVpcEndpointsOutput, error) {
	endpoints, err := p.cli.DeleteVpcEndpoints(ctx, &ec2api.DeleteVpcEndpointsInput{
		VpcEndpointIds: endpointIDs,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(p.cfg.Region, privateEndpointClientLabel, "DeleteVpcEndpoints", strings.Join(endpointIDs, ","), errDeleteVpcEndpointFailed)
		return nil, err
	}
	return endpoints, nil
}