package aws

import (
	"context"
	"fmt"
	"strings"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"

	"github.com/tidbcloud/controlplane9/internal/cluster"
)

var (
	iamReadyLoadAWSConfig                       = LoadAWSConfig
	iamReadyLoadAWSConfigFromExportedCredential = LoadAWSConfigFromExportedProfileCredentials
	iamReadyNewSTSClient                        = newSTSClient
	iamReadyValidateIAMCleanupFromConfig        = ValidateIAMCleanupReadinessFromConfig
)

// OwnerIdentityReadyStep resolves the active cp9-owner before any substrate step runs.
type OwnerIdentityReadyStep struct {
	Profile        string
	SourceProfile  string
	Region         string
	OwnerAllowlist map[string][]string
	STSClient      STSAPI
}

func (s *OwnerIdentityReadyStep) Name() cluster.StepName { return cluster.StepOwnerIdentityReady }

func (s *OwnerIdentityReadyStep) Run(_ context.Context) error {
	return fmt.Errorf("owner_identity_ready requires lifecycle state")
}

func (s *OwnerIdentityReadyStep) RunWithState(ctx context.Context, state *cluster.State) error {
	result, err := s.validateAccess(ctx)
	if err != nil {
		return fmt.Errorf("owner_identity_ready: %w", err)
	}
	owner, err := ResolveOwnerIdentity(result, s.OwnerAllowlist)
	if err != nil {
		return fmt.Errorf("owner_identity_ready: %w", err)
	}
	state.Owner = owner
	return nil
}

func (s *OwnerIdentityReadyStep) validateAccess(ctx context.Context) (*SetupResult, error) {
	stsClient := s.STSClient
	profile := effectiveSourceProfile(s.Profile, s.SourceProfile)
	var cfg *awssdk.Config
	if stsClient == nil {
		loaded, err := iamReadyLoadAWSConfig(ctx, profile, s.Region)
		if err != nil {
			return nil, fmt.Errorf("load AWS config: %w", err)
		}
		cfg = &loaded
		stsClient = iamReadyNewSTSClient(cfg)
	}
	result, err := ValidateAccess(ctx, stsClient, profile, s.Region)
	if err == nil || s.STSClient != nil || !IsSSOCredentialRefreshError(err) {
		return result, err
	}
	exportedCfg, exportErr := iamReadyLoadAWSConfigFromExportedCredential(ctx, profile, s.Region, nil)
	if exportErr != nil {
		return result, fmt.Errorf("%w; AWS CLI credential export fallback failed: %w", err, exportErr)
	}
	return ValidateAccess(ctx, iamReadyNewSTSClient(&exportedCfg), profile, s.Region)
}

// IAMReadyStep proves the local AWS/IAM contract and persists the session metadata.
type IAMReadyStep struct {
	SessionID        string
	StateDir         string
	Profile          string
	SourceProfile    string
	Region           string
	RoleARN          string
	OwnerAllowlist   map[string][]string
	STSClient        STSAPI
	IAMCleanupClient IAMCleanupReadinessAPI
}

func (s *IAMReadyStep) Name() cluster.StepName { return cluster.StepIAMReady }

func (s *IAMReadyStep) Run(_ context.Context) error {
	return fmt.Errorf("iam_ready requires lifecycle state")
}

func (s *IAMReadyStep) RunWithState(ctx context.Context, state *cluster.State) error {
	result, err := s.validateIAM(ctx)
	if err != nil {
		return fmt.Errorf("iam_ready: %w", err)
	}
	if state.Owner != "" && state.Owner != result.Owner {
		return fmt.Errorf("iam_ready: resolved owner %q does not match lifecycle owner %q", result.Owner, state.Owner)
	}
	state.Owner = result.Owner

	_, err = WriteLocalProfile(&LocalProfileBinding{
		Name:          s.Profile,
		SourceProfile: effectiveSourceProfile(s.Profile, s.SourceProfile),
		RoleARN:       s.RoleARN,
		SessionName:   s.SessionID,
		Owner:         result.SessionTags[TagOwner],
		Region:        s.Region,
	})
	if err != nil {
		return fmt.Errorf("iam_ready: persist local AWS profile: %w", err)
	}

	sessionState := &SessionState{
		SessionID:        s.SessionID,
		Owner:            result.Owner,
		Profile:          s.Profile,
		SourceProfile:    effectiveSourceProfile(s.Profile, s.SourceProfile),
		Region:           s.Region,
		Account:          result.Account,
		RoleARN:          result.RoleARN,
		CredentialSource: result.CredentialSource,
		SessionTags:      result.SessionTags,
	}
	if err := SaveSessionState(s.StateDir, sessionState); err != nil {
		return fmt.Errorf("iam_ready: save state: %w", err)
	}
	return nil
}

func (s *IAMReadyStep) validateIAM(ctx context.Context) (*SetupResult, error) {
	stsClient := s.STSClient
	profile := effectiveSourceProfile(s.Profile, s.SourceProfile)
	var cfg *awssdk.Config
	if stsClient == nil || s.IAMCleanupClient == nil {
		loaded, err := iamReadyLoadAWSConfig(ctx, profile, s.Region)
		if err != nil {
			return nil, fmt.Errorf("load AWS config: %w", err)
		}
		cfg = &loaded
	}
	if stsClient == nil {
		stsClient = iamReadyNewSTSClient(cfg)
	}
	result, err := ValidateIAMReadiness(
		ctx,
		stsClient,
		profile,
		s.Region,
		s.RoleARN,
		s.SessionID,
		s.OwnerAllowlist,
	)
	if err != nil {
		if cfg == nil || s.STSClient != nil || !IsSSOCredentialRefreshError(err) {
			return result, err
		}
		exportedCfg, exportErr := iamReadyLoadAWSConfigFromExportedCredential(ctx, profile, s.Region, nil)
		if exportErr != nil {
			return result, fmt.Errorf("%w; AWS CLI credential export fallback failed: %w", err, exportErr)
		}
		cfg = &exportedCfg
		result, err = ValidateIAMReadiness(
			ctx,
			iamReadyNewSTSClient(cfg),
			profile,
			s.Region,
			s.RoleARN,
			s.SessionID,
			s.OwnerAllowlist,
		)
		if err != nil {
			return result, err
		}
		PrefixCredentialSource(result, "aws-cli-export")
	}
	switch {
	case s.IAMCleanupClient != nil:
		err = ValidateIAMCleanupReadiness(ctx, s.IAMCleanupClient, s.SessionID)
	case cfg != nil:
		err = iamReadyValidateIAMCleanupFromConfig(ctx, cfg, result, s.SessionID)
	}
	if err != nil {
		return result, err
	}
	return result, nil
}

func effectiveSourceProfile(profile, sourceProfile string) string {
	sourceProfile = strings.TrimSpace(sourceProfile)
	if sourceProfile != "" {
		return sourceProfile
	}
	return strings.TrimSpace(profile)
}