package cmd

import (
	"context"
	"fmt"
	"os"
	"strings"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/innerr/ticat/pkg/core/model"

	cpaws "github.com/tidbcloud/controlplane9/internal/aws"
	"github.com/tidbcloud/controlplane9/internal/config"
	cp9errors "github.com/tidbcloud/controlplane9/internal/errors"
	"github.com/tidbcloud/controlplane9/internal/host"
	"github.com/tidbcloud/controlplane9/internal/host/remoteimage"
)

var (
	makeScopedReaper       = newScopedReaper
	saveReapState          = cpaws.SaveReapState
	loadReapState          = cpaws.LoadReapState
	removeReapState        = cpaws.RemoveReapState
	findReapStateBySession = cpaws.FindReapStateBySession
)

func setupCmd() model.NormalCmd {
	return func(args model.ArgVals, cc *model.Cli, env *model.Env, _ []model.ParsedCmd) error {
		syncVerbose(env)

		if !strings.EqualFold(strings.TrimSpace(args.GetRaw("checkout")), "false") {
			if err := continueSetupWithLatestRelease(context.Background()); err != nil {
				return err
			}
		}

		cfgPath := args.GetRaw("config")
		version, err := config.DetectAPIVersion(cfgPath)
		if err != nil {
			return err
		}
		if version == config.APIVersionV1Alpha2 {
			return runV1Alpha2Setup(cfgPath, cc, env)
		}
		return runLegacySetup(cfgPath, cc, env)
	}
}

func runLegacySetup(cfgPath string, cc *model.Cli, env *model.Env) error {
	cfg, _, err := config.Load(cfgPath, "")
	if err != nil {
		return err
	}

	ctx := context.Background()
	awsCfg, err := cmdLoadAWSConfig(ctx, cfg.Clouds.AWS.Profile, cfg.Clouds.AWS.Region)
	if err != nil {
		return cp9errors.Newf(cp9errors.CategoryResource, "setup: %v", err)
	}

	stsClient := cmdNewSTSClient(&awsCfg)
	result, err := cpaws.ValidateAccess(ctx, stsClient, cfg.Clouds.AWS.Profile, cfg.Clouds.AWS.Region)
	if err != nil {
		return cp9errors.Newf(cp9errors.CategoryResource, "setup: %v", err)
	}
	owner, err := cmdResolveOwnerID(result, cfg.CI.OwnerAllowlist)
	if err != nil {
		return cp9errors.Newf(cp9errors.CategoryConfig, "setup: %v", err)
	}
	result.Owner = owner

	return model.Output(cc, env, result)
}

func runV1Alpha2Setup(cfgPath string, cc *model.Cli, env *model.Env) error {
	session, err := config.LoadSession(cfgPath)
	if err != nil {
		return err
	}
	session, err = config.ResolveSession(session, cfgPath)
	if err != nil {
		return err
	}

	ctx := context.Background()
	result, err := runV1Alpha2ReadOnlyPreflight(ctx, session)
	if err != nil {
		if result != nil {
			if outputErr := model.Output(cc, env, result); outputErr != nil {
				return outputErr
			}
		}
		category := cp9errors.CategoryResource
		if cpaws.IsContractError(err) {
			category = cp9errors.CategoryConfig
		}
		return cp9errors.Newf(category, "setup: %v", err)
	}
	if config.UsesCloudSubstrate(session) {
		sourceProfile := strings.TrimSpace(session.Target.AWS.IAM.SourceProfile)
		if sourceProfile == "" {
			sourceProfile = strings.TrimSpace(session.Target.AWS.IAM.Profile)
		}
		localProfile, err := cpaws.WriteLocalProfile(&cpaws.LocalProfileBinding{
			Name:          session.Target.AWS.IAM.Profile,
			SourceProfile: sourceProfile,
			RoleARN:       session.Target.AWS.IAM.RoleARN,
			SessionName:   session.Name,
			Owner:         result.SessionTags[cpaws.TagOwner],
			Region:        session.Target.Region,
		})
		if err != nil {
			return cp9errors.Newf(cp9errors.CategoryResource, "setup: persist local AWS profile: %v", err)
		}
		result.LocalProfile = localProfile
		result.NoMutation = false
	}
	return model.Output(cc, env, result)
}

func runV1Alpha2ReadOnlyPreflight(ctx context.Context, session *config.SessionConfig) (*cpaws.SetupResult, error) {
	if session == nil {
		return nil, fmt.Errorf("session config is required")
	}
	if !config.UsesCloudSubstrate(session) {
		owner, tags, err := cpaws.ResolveLocalSessionIdentity(session.Name)
		if err != nil {
			return nil, err
		}
		result := &cpaws.SetupResult{
			Valid:            true,
			NoMutation:       true,
			CleanupReady:     true,
			CleanupCommand:   "cp9 destroy",
			Account:          "local",
			ARN:              "local",
			UserID:           "local",
			Owner:            owner,
			Profile:          "local",
			Region:           "local",
			CredentialSource: "local",
			SessionTags:      tags,
		}
		if session.Host.TTL.Duration > 0 {
			result.TTL = session.Host.TTL.String()
		}
		collision, err := runV1Alpha2SessionCollisionPreflight(ctx, session, owner, false)
		result.SessionCollision = collision
		if err != nil {
			result.Valid = false
		}
		return result, err
	}
	profile := session.Target.AWS.IAM.SourceProfile
	if strings.TrimSpace(profile) == "" {
		profile = session.Target.AWS.IAM.Profile
	}
	awsCfg, err := cmdLoadAWSConfig(ctx, profile, session.Target.Region)
	if err != nil {
		return nil, err
	}
	result, awsCfg, err := cmdRunV1Alpha2SetupIAMPreflight(
		ctx,
		&awsCfg,
		profile,
		session.Target.Region,
		session.Target.AWS.IAM.RoleARN,
		session.Name,
		session.CI.OwnerAllowlist,
	)
	if err != nil {
		return result, err
	}
	if err := cmdValidateIAMCleanup(ctx, &awsCfg, result, session.Name); err != nil {
		return result, err
	}
	effective := cpaws.EffectiveAWSConfig(&awsCfg, result)
	if session.Host.Mode == config.HostModeRemoteEC2K3d {
		validation, err := validateV1Alpha2UserManagedEC2Key(ctx, session, &effective, "", result.Account)
		result.UserManagedEC2Key = validation
		if err != nil {
			result.Valid = false
			return result, err
		}
	}
	collision, err := cmdRunV1Alpha2SessionCollisionPreflight(ctx, session, result.Owner, true)
	result.SessionCollision = collision
	if err != nil {
		result.Valid = false
		return result, err
	}
	if err := cmdResolveV1Alpha2RemoteImagePlan(ctx, &awsCfg, result, session); err != nil {
		return result, err
	}
	if session.Host.TTL.Duration > 0 {
		result.TTL = session.Host.TTL.String()
	}
	result.CleanupReady = true
	result.CleanupCommand = "cp9 destroy"
	return result, nil
}

func validateV1Alpha2UserManagedEC2Key(
	ctx context.Context,
	session *config.SessionConfig,
	awsCfg *awssdk.Config,
	stateDir string,
	account string,
) (validation *cpaws.UserManagedEC2KeyValidation, err error) {
	if session == nil {
		return nil, fmt.Errorf("session config is required for remote EC2 key validation")
	}
	if session.Host.Mode != config.HostModeRemoteEC2K3d {
		return nil, fmt.Errorf("remote EC2 key validation requires host.mode remote-ec2-k3d")
	}
	if strings.TrimSpace(stateDir) == "" {
		stateDir = host.ResolveSessionConfig(session).StateDir
	}
	instanceID, err := persistedRemoteInstanceID(stateDir)
	if err != nil {
		return nil, err
	}
	if awsCfg == nil && strings.TrimSpace(instanceID) == "" {
		return nil, fmt.Errorf("AWS config is required for EC2 key validation before the first remote host mutation")
	}
	validation, err = cmdValidateUserManagedEC2Key(
		ctx,
		awsCfg,
		account,
		session.Target.Region,
		session.Host.EC2.KeyPair,
		session.Host.EC2.SSHPrivateKey,
		instanceID,
	)
	if err != nil {
		diagnosis := cpaws.DiagnoseUserManagedEC2Key(err)
		return validation, fmt.Errorf(
			"validate host.ec2 key pair: %w (failure_class=%s field=%s remediation=%s)",
			err,
			diagnosis.FailureClass,
			diagnosis.Field,
			diagnosis.Remediation,
		)
	}
	return validation, nil
}

func persistedRemoteInstanceID(stateDir string) (string, error) {
	selected, err := host.Load(stateDir)
	if os.IsNotExist(err) {
		return "", nil
	}
	if err != nil {
		return "", fmt.Errorf("load persisted host for EC2 key validation: %w", err)
	}
	if selected == nil || selected.Remote == nil {
		return "", nil
	}
	return strings.TrimSpace(selected.Remote.InstanceID), nil
}

func resolveV1Alpha2RemoteImagePlan(
	ctx context.Context,
	awsCfg *awssdk.Config,
	result *cpaws.SetupResult,
	session *config.SessionConfig,
) error {
	if session == nil || session.Host.Mode != config.HostModeRemoteEC2K3d {
		return nil
	}
	if awsCfg == nil {
		return fmt.Errorf("AWS config is required to derive the remote image plan")
	}
	effective := cpaws.EffectiveAWSConfig(awsCfg, result)
	resolver := &remoteimage.Resolver{
		EC2:    cmdNewRemoteImageEC2Client(effective),
		SSM:    cmdNewRemoteImageSSMClient(effective),
		Region: session.Target.Region,
	}
	plan, err := resolver.Resolve(ctx, session.Host.EC2.InstanceType)
	if err != nil {
		return fmt.Errorf("derive remote image plan: %w", err)
	}
	result.RemoteImage = plan
	return nil
}

func runV1Alpha2SetupIAMPreflight(
	ctx context.Context,
	awsCfg *awssdk.Config,
	profile, region, roleARN, sessionName string,
	ownerAllowlist map[string][]string,
) (*cpaws.SetupResult, awssdk.Config, error) {
	result, err := cmdValidateIAM(
		ctx,
		cmdNewSTSClient(awsCfg),
		profile,
		region,
		roleARN,
		sessionName,
		ownerAllowlist,
	)
	if err == nil || !shouldRetrySetupWithAWSCLIExport(err) {
		return result, *awsCfg, err
	}

	exportedCfg, exportErr := cmdLoadAWSConfigFromExportedCredential(ctx, profile, region, nil)
	if exportErr != nil {
		return result, *awsCfg, fmt.Errorf("%w; AWS CLI credential export fallback failed: %w", err, exportErr)
	}
	result, retryErr := cmdValidateIAM(
		ctx,
		cmdNewSTSClient(&exportedCfg),
		profile,
		region,
		roleARN,
		sessionName,
		ownerAllowlist,
	)
	if retryErr != nil {
		return result, exportedCfg, retryErr
	}
	cpaws.PrefixCredentialSource(result, "aws-cli-export")
	return result, exportedCfg, nil
}

func shouldRetrySetupWithAWSCLIExport(err error) bool {
	return cpaws.IsSSOCredentialRefreshError(err)
}

func resolveInvokerOwner(ctx context.Context, cfg *config.Config) (string, error) {
	awsCfg, err := cpaws.LoadAWSConfig(ctx, cfg.Clouds.AWS.Profile, cfg.Clouds.AWS.Region)
	if err != nil {
		return "", err
	}
	result, err := cpaws.ValidateAccess(ctx, newSTSFromConfig(&awsCfg), cfg.Clouds.AWS.Profile, cfg.Clouds.AWS.Region)
	if err != nil {
		return "", err
	}
	return cpaws.ResolveOwnerIdentity(result, cfg.CI.OwnerAllowlist)
}

func newScopedReaper(profile string) (*cpaws.Reaper, error) {
	return newScopedAssumeRoleReaper(profile, "", true)
}

func newScopedAssumeRoleReaper(profile, assumeRoleARN string, assumeRoleWithTags bool) (*cpaws.Reaper, error) {
	return newScopedAssumeRoleReaperWithIAMInventory(profile, assumeRoleARN, assumeRoleWithTags, nil)
}

// newScopedAssumeRoleReaperWithIAMInventory is newScopedAssumeRoleReaper with an
// optional shared IAM snapshot. Broad cleanup passes one so tagged-IAM discovery
// costs a single account scan for the whole run rather than one per session; a
// nil inventory keeps the original per-session scan.
func newScopedAssumeRoleReaperWithIAMInventory(
	profile, assumeRoleARN string,
	assumeRoleWithTags bool,
	iamInventory *cpaws.IAMInventory,
) (*cpaws.Reaper, error) {
	scanner := &cpaws.RegionalScanner{
		Factory: func(region string) (cpaws.Tracker, error) {
			awsCfg, err := cpaws.LoadAWSConfig(context.Background(), profile, region)
			if err != nil {
				return nil, err
			}
			awsCfg = withDestroyAWSHTTPClient(&awsCfg)
			if assumeRoleARN != "" {
				awsCfg = reaperAssumeRoleConfig(&awsCfg, assumeRoleARN, assumeRoleWithTags)
			}
			if err := cpaws.FreezeConfigCredentials(context.Background(), &awsCfg); err != nil {
				return nil, fmt.Errorf("freeze reaper credentials: %w", err)
			}
			tracker, _, _ := cpaws.NewClients(&awsCfg)
			return tracker.WithIAMInventory(iamInventory), nil
		},
	}
	cleaner := &cpaws.RegionalCleaner{
		Factory: func(region string) (cpaws.Cleaner, error) {
			awsCfg, err := cpaws.LoadAWSConfig(context.Background(), profile, region)
			if err != nil {
				return nil, err
			}
			awsCfg = withDestroyAWSHTTPClient(&awsCfg)
			if assumeRoleARN != "" {
				awsCfg = reaperAssumeRoleConfig(&awsCfg, assumeRoleARN, assumeRoleWithTags)
			}
			if err := cpaws.FreezeConfigCredentials(context.Background(), &awsCfg); err != nil {
				return nil, fmt.Errorf("freeze reaper credentials: %w", err)
			}
			_, regionalCleaner, _ := cpaws.NewClients(&awsCfg)
			return regionalCleaner, nil
		},
	}
	return cpaws.NewReaper(scanner, cleaner), nil
}

func reaperAssumeRoleConfig(cfg *awssdk.Config, roleARN string, withTags bool) awssdk.Config {
	if withTags {
		return cpaws.ConfigWithAssumeRole(cfg, roleARN, os.Getenv("CP9_OWNER"), map[string]string{
			cpaws.TagOwner:   os.Getenv("CP9_OWNER"),
			cpaws.TagSession: os.Getenv("CP9_SESSION_NAME"),
		})
	}
	assumed := cfg.Copy()
	assumed.Credentials = awssdk.NewCredentialsCache(stscreds.NewAssumeRoleProvider(
		sts.NewFromConfig(*cfg),
		strings.TrimSpace(roleARN),
		func(o *stscreds.AssumeRoleOptions) {
			o.RoleSessionName = firstNonEmptyString(os.Getenv("CP9_SESSION_NAME"), "cp9-reaper")
		},
	))
	return assumed
}