package awsauth

import (
	"context"
	"crypto/sha256"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/patrickmn/go-cache"
	"github.com/pingcap/log"
	"github.com/tidbcloud/observability/pkg/common/constants"
	"github.com/tidbcloud/observability/pkg/common/errors"
	"go.uber.org/zap"
)

const (
	// apiGatewayService is the service name for AWS API Gateway
	apiGatewayService = "execute-api"
)

var (
	// Set default expiration to 1 hour, cleanup every 2 hours
	sigv4CredCache     = cache.New(1*time.Hour, 2*time.Hour)
	sigv4CredCacheLock sync.Mutex

	// signer is the shared SigV4 signer instance
	signer = v4.NewSigner()
)

// SigV4AuthInfo contains the information needed for SigV4 signing
type SigV4AuthInfo struct {
	RoleArn    string
	ExternalID string
	Region     string
}

// GetSigV4Headers generates SigV4 signed headers for AWS API Gateway access
// @Description: Generate SigV4 authentication headers for requests to AWS API Gateway
// @Parameter ctx - context for the request
// @Parameter authInfo - SigV4 authentication information (role ARN, external ID, region)
// @Parameter method - HTTP method
// @Parameter targetURL - full URL to sign
// @Parameter body - request body (can be empty)
// @return map[string]string - headers to add to the request
// @return error
func GetSigV4Headers(ctx context.Context, authInfo SigV4AuthInfo, method, targetURL string, body []byte) (map[string]string, error) {
	if authInfo.RoleArn == "" || authInfo.Region == "" {
		return nil, errors.NewErrorf(errors.ParameterInvalid, "authInfo missing for SigV4 signing")
	}

	// Get or create credentials
	creds, err := getOrAssumeRole(ctx, authInfo)
	if err != nil {
		return nil, err
	}

	// Create the HTTP request for signing
	req, err := http.NewRequestWithContext(ctx, method, targetURL, nil)
	if err != nil {
		return nil, errors.NewErrorf(errors.HttpClientResponseError, "failed to create request: %s", err.Error())
	}

	// Calculate the payload hash
	var payloadHash string
	if len(body) > 0 {
		hash := sha256.Sum256(body)
		payloadHash = fmt.Sprintf("%x", hash)
	} else {
		// Empty payload hash
		hash := sha256.Sum256([]byte(""))
		payloadHash = fmt.Sprintf("%x", hash)
	}

	// Sign the request, it's thread safe
	err = signer.SignHTTP(ctx, creds, req, payloadHash, apiGatewayService, authInfo.Region, time.Now().UTC())
	if err != nil {
		return nil, errors.NewErrorf(errors.AWSClientError, "failed to sign request: %s", err.Error())
	}

	// Extract the signed headers
	headers := make(map[string]string)
	for key, values := range req.Header {
		if len(values) > 0 {
			headers[key] = values[0]
		}
	}

	log.Debug("SigV4 signing completed",
		zap.String("method", method),
		zap.String("host", targetURL),
	)

	return headers, nil
}

// getOrAssumeRole gets cached credentials or assumes the role to get new ones
func getOrAssumeRole(ctx context.Context, authInfo SigV4AuthInfo) (aws.Credentials, error) {
	cacheKey := fmt.Sprintf("%s:%s", authInfo.RoleArn, authInfo.ExternalID)

	// Check cache first
	if cached, ok := sigv4CredCache.Get(cacheKey); ok {
		creds := cached.(aws.Credentials)
		// Verify credentials haven't expired
		if !creds.Expired() {
			log.Debug("Using cached SigV4 credentials", zap.String("roleArn", authInfo.RoleArn))
			return creds, nil
		}
	}

	sigv4CredCacheLock.Lock()
	defer sigv4CredCacheLock.Unlock()

	// Double-check cache after acquiring lock
	if cached, ok := sigv4CredCache.Get(cacheKey); ok {
		creds := cached.(aws.Credentials)
		if !creds.Expired() {
			return creds, nil
		}
	}

	// Assume the role
	creds, err := assumeRole(ctx, authInfo)
	if err != nil {
		return aws.Credentials{}, err
	}

	// Cache the credentials
	// Use the expiration time from the credentials to set the cache TTL, minus a buffer
	// to avoid using credentials that are about to expire
	ttl := time.Until(creds.Expires.Add(-1 * time.Minute))
	if ttl > 0 {
		sigv4CredCache.Set(cacheKey, creds, ttl)
	}

	log.Info("Assumed role for SigV4 signing", zap.String("roleArn", authInfo.RoleArn))
	return creds, nil
}

// assumeRole assumes an IAM role and returns temporary credentials
func assumeRole(ctx context.Context, authInfo SigV4AuthInfo) (aws.Credentials, error) {
	region := authInfo.Region

	// Load default AWS config
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return aws.Credentials{}, errors.NewErrorf(errors.AWSClientError, "failed to load AWS config: %s", err.Error())
	}

	// Create STS client
	stsClient := sts.NewFromConfig(cfg)

	// Build assume role input
	var externalIDPtr *string
	if authInfo.ExternalID != "" {
		externalIDPtr = aws.String(authInfo.ExternalID)
	}
	assumeRoleInput := &sts.AssumeRoleInput{
		RoleArn:         aws.String(authInfo.RoleArn),
		RoleSessionName: aws.String(roleSessionNameForSigV4()),
		ExternalId:      externalIDPtr,
	}

	// Assume the role
	assumeRoleOutput, err := stsClient.AssumeRole(ctx, assumeRoleInput)
	if err != nil {
		log.Error("Failed to assume role for SigV4",
			zap.String("roleArn", authInfo.RoleArn),
			zap.Error(err))
		return aws.Credentials{}, errors.NewErrorf(errors.GenStsCredentialError, "failed to assume role: %s", err.Error())
	}

	return aws.Credentials{
		AccessKeyID:     *assumeRoleOutput.Credentials.AccessKeyId,
		SecretAccessKey: *assumeRoleOutput.Credentials.SecretAccessKey,
		SessionToken:    *assumeRoleOutput.Credentials.SessionToken,
		CanExpire:       true,
		Expires:         *assumeRoleOutput.Credentials.Expiration,
	}, nil
}

// roleSessionNameForSigV4 generates a role session name for SigV4 signing
func roleSessionNameForSigV4() string {
	suffix, err := os.Hostname()
	if err != nil {
		suffix = fmt.Sprintf("%d", time.Now().Unix())
	}
	return fmt.Sprintf("%s_sigv4_%s", constants.ObservabilityGlobalServer, suffix)
}