// Package aws provides the AWS SSO (IAM Identity Center) provider implementation.
// It implements the provider.Provider interface for managing AWS SSO permissions.
//
// Requirements: 3.2, 3.3, 3.4, 3.6
package aws

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/identitystore"
	"github.com/aws/aws-sdk-go/service/ssoadmin"

	"AutoSSO/internal/config"
	"AutoSSO/internal/pkg/logger"
	"AutoSSO/internal/provider"
)

var waitBeforeOperationDelay = 10 * time.Second

// AWSProvider implements the provider.Provider interface for AWS SSO.
// It manages permissions through AWS IAM Identity Center (formerly AWS SSO).
//
// Requirements: 3.4, 3.6
type AWSProvider struct {
	// config holds the AWS configuration
	config *config.AWSConfig

	// session is the AWS SDK session
	session *session.Session

	// ssoAdminClient is the SSO Admin API client
	ssoAdminClient *ssoadmin.SSOAdmin

	// identityStoreClient is the Identity Store API client
	identityStoreClient *identitystore.IdentityStore

	// permissionSetCache caches permission set name to ARN mappings
	permissionSetCache map[string]string

	// cacheMu protects the permission set cache
	cacheMu sync.RWMutex

	// closed indicates if the provider has been closed
	closed bool

	// closeMu protects the closed flag
	closeMu sync.RWMutex

	// logger is the logger instance
	logger *logger.Logger
}

// NewAWSProvider creates a new AWS provider instance from the configuration.
// It initializes the AWS SDK session and creates the necessary API clients.
//
// Parameters:
//   - cfg: The AWS configuration containing Identity Center ARN, Identity Store ID, and region
//
// Returns:
//   - *AWSProvider: The initialized AWS provider
//   - error: Any error that occurred during initialization
//
// Requirements: 3.4, 3.6
func NewAWSProvider(cfg *config.AWSConfig) (*AWSProvider, error) {
	if cfg == nil {
		return nil, fmt.Errorf("AWS config cannot be nil")
	}

	if cfg.IdentityCenterARN == "" {
		return nil, fmt.Errorf("AWS Identity Center ARN is required")
	}

	if cfg.IdentityStoreID == "" {
		return nil, fmt.Errorf("AWS Identity Store ID is required")
	}

	if cfg.Region == "" {
		return nil, fmt.Errorf("AWS Region is required")
	}

	// Create AWS session with shared config
	sess, err := session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
		Config: aws.Config{
			Region: aws.String(cfg.Region),
		},
	})
	if err != nil {
		return nil, fmt.Errorf("failed to create AWS session: %w", err)
	}

	// Create SSO Admin client
	ssoAdminClient := ssoadmin.New(sess, &aws.Config{
		Region: aws.String(cfg.Region),
	})

	// Create Identity Store client
	// The endpoint is constructed from the region
	identityStoreEndpoint := fmt.Sprintf("https://identitystore.%s.amazonaws.com", cfg.Region)
	identityStoreClient := identitystore.New(sess, &aws.Config{
		Region:   aws.String(cfg.Region),
		Endpoint: aws.String(identityStoreEndpoint),
	})

	return &AWSProvider{
		config:              cfg,
		session:             sess,
		ssoAdminClient:      ssoAdminClient,
		identityStoreClient: identityStoreClient,
		permissionSetCache:  make(map[string]string),
		logger:              logger.NewLogger(logger.DefaultConfig()).WithModule("aws-provider"),
	}, nil
}

// WithLogger sets the logger for the provider.
func (p *AWSProvider) WithLogger(log *logger.Logger) *AWSProvider {
	p.logger = log.WithModule("aws-provider")
	return p
}

// Name returns the provider name.
//
// Returns:
//   - string: "AWS"
func (p *AWSProvider) Name() string {
	return string(provider.ProviderAWS)
}

// ValidateCredentials validates that the AWS credentials are valid
// by attempting to list permission sets.
//
// Parameters:
//   - ctx: Context for cancellation and timeout
//
// Returns:
//   - error: Any error if credentials are invalid or authentication fails
func (p *AWSProvider) ValidateCredentials(ctx context.Context) error {
	p.closeMu.RLock()
	if p.closed {
		p.closeMu.RUnlock()
		return fmt.Errorf("provider is closed")
	}
	p.closeMu.RUnlock()

	// Try to list permission sets to validate credentials
	input := &ssoadmin.ListPermissionSetsInput{
		InstanceArn: aws.String(p.config.IdentityCenterARN),
		MaxResults:  aws.Int64(1),
	}

	_, err := p.ssoAdminClient.ListPermissionSetsWithContext(ctx, input)
	if err != nil {
		return fmt.Errorf("failed to validate AWS credentials: %w", err)
	}

	return nil
}

// Close releases any resources held by the provider.
//
// Returns:
//   - error: Any error that occurred during cleanup
func (p *AWSProvider) Close() error {
	p.closeMu.Lock()
	defer p.closeMu.Unlock()

	if p.closed {
		return nil
	}

	p.closed = true

	// Clear the cache
	p.cacheMu.Lock()
	p.permissionSetCache = make(map[string]string)
	p.cacheMu.Unlock()

	return nil
}

// getPermissionSetARN retrieves the ARN for a permission set by name.
// It uses a cache to avoid repeated API calls.
//
// Parameters:
//   - ctx: Context for cancellation and timeout
//   - permissionSetName: The name of the permission set
//
// Returns:
//   - string: The ARN of the permission set
//   - error: Any error that occurred during the lookup
func (p *AWSProvider) getPermissionSetARN(ctx context.Context, permissionSetName string) (string, error) {
	// Check cache first
	p.cacheMu.RLock()
	if arn, exists := p.permissionSetCache[permissionSetName]; exists {
		p.cacheMu.RUnlock()
		return arn, nil
	}
	p.cacheMu.RUnlock()

	// List all permission sets and find the matching one
	input := &ssoadmin.ListPermissionSetsInput{
		InstanceArn: aws.String(p.config.IdentityCenterARN),
		MaxResults:  aws.Int64(100),
	}

	for {
		resp, err := p.ssoAdminClient.ListPermissionSetsWithContext(ctx, input)
		if err != nil {
			return "", fmt.Errorf("failed to list permission sets: %w", err)
		}

		for _, arn := range resp.PermissionSets {
			name, err := p.getPermissionSetName(ctx, *arn)
			if err != nil {
				continue
			}

			// Cache the mapping
			p.cacheMu.Lock()
			p.permissionSetCache[name] = *arn
			p.cacheMu.Unlock()

			if name == permissionSetName {
				return *arn, nil
			}
		}

		// Check for more results
		if resp.NextToken != nil {
			input.NextToken = resp.NextToken
		} else {
			break
		}
	}

	return "", fmt.Errorf("permission set not found: %s", permissionSetName)
}

// getPermissionSetName retrieves the name of a permission set by ARN.
//
// Parameters:
//   - ctx: Context for cancellation and timeout
//   - permissionSetARN: The ARN of the permission set
//
// Returns:
//   - string: The name of the permission set
//   - error: Any error that occurred during the lookup
func (p *AWSProvider) getPermissionSetName(ctx context.Context, permissionSetARN string) (string, error) {
	input := &ssoadmin.DescribePermissionSetInput{
		InstanceArn:      aws.String(p.config.IdentityCenterARN),
		PermissionSetArn: aws.String(permissionSetARN),
	}

	resp, err := p.ssoAdminClient.DescribePermissionSetWithContext(ctx, input)
	if err != nil {
		return "", fmt.Errorf("failed to describe permission set: %w", err)
	}

	return *resp.PermissionSet.Name, nil
}

// waitBeforeOperation waits for a short period before performing an operation.
// This is to avoid conflicts when adding and removing permissions in quick succession.
// AWS SDK may throw errors if operations are performed within 10 seconds of each other.
func (p *AWSProvider) waitBeforeOperation(ctx context.Context) error {
	timer := time.NewTimer(waitBeforeOperationDelay)
	defer timer.Stop()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}