package iamgen

import (
	"encoding/json"
	"fmt"

	ackv1alpha1 "github.com/aws-controllers-k8s/iam-controller/apis/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// GenerateActorRole creates an ACK Role for a workload
func GenerateActorRole(cfg ActorConfig) (*ackv1alpha1.Role, error) {
	var trustPolicy IAMPolicy
	var err error

	switch cfg.Type {
	case ActorTypePodIdentity:
		trustPolicy, err = generatePodIdentityTrust(cfg.PodIdentity)
	case ActorTypeIRSA:
		trustPolicy, err = generateIRSATrust(cfg.AccountID, cfg.IRSA)
	case ActorTypeSPIFFE:
		trustPolicy, err = generateSPIFFETrust(cfg.AccountID, cfg.SPIFFE)
	default:
		return nil, fmt.Errorf("unknown actor type: %s", cfg.Type)
	}

	if err != nil {
		return nil, err
	}

	return buildACKRole(
		cfg.Name, cfg.Namespace, cfg.RoleName,
		trustPolicy,
		"sts:AssumeRole", cfg.TargetRoleARNs,
		"CrossAccountHops",
	), nil
}

// GenerateTargetRole creates an ACK Role for a resource
func GenerateTargetRole(cfg TargetConfig) (*ackv1alpha1.Role, error) {
	// 1. Build Trust Policy (trusting the actors)
	trustPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustActors",
			Effect:    "Allow",
			Principal: map[string][]string{"AWS": cfg.TrustedRoleARNs},
			Action:    "sts:AssumeRole",
		}},
	}

	// 2. Determine Permissions
	var action, resource string
	switch cfg.Type {
	case TargetTypeDatabase:
		if cfg.Database == nil {
			return nil, fmt.Errorf("database config missing")
		}
		action = "rds-db:connect"
		resource = fmt.Sprintf("arn:aws:rds-db:%s:%s:dbuser:%s/%s",
			cfg.Database.Region, cfg.AccountID, cfg.Database.DBClusterID, cfg.Database.DBUser)
	case TargetTypeGeneric:
		if cfg.Generic == nil {
			return nil, fmt.Errorf("generic config missing")
		}
		action = cfg.Generic.Action
		resource = cfg.Generic.ResourceARN
	}

	return buildACKRole(
		cfg.Name, cfg.Namespace, cfg.RoleName,
		trustPolicy,
		action, []string{resource},
		"ResourcePermissions",
	), nil
}

// --- INTERNAL HELPERS ---

func buildACKRole(k8sName, ns, awsRoleName string, trust IAMPolicy, action string, resources []string, policyName string) *ackv1alpha1.Role {
	trustJSON, _ := json.Marshal(trust)

	permPolicy := IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:      "AllowAction",
			Effect:   "Allow",
			Action:   action,
			Resource: resources,
		}},
	}
	permJSON, _ := json.Marshal(permPolicy)

	return &ackv1alpha1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      k8sName,
			Namespace: ns,
		},
		Spec: ackv1alpha1.RoleSpec{
			Name:                     &awsRoleName,
			AssumeRolePolicyDocument: string(trustJSON),
			InlinePolicies:           map[string]*string{policyName: awsString(string(permJSON))},
		},
	}
}

func generatePodIdentityTrust(cfg *PodIdentityConfig) (IAMPolicy, error) {
	if cfg == nil {
		return IAMPolicy{}, fmt.Errorf("PodIdentity config missing")
	}
	return IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustPIA",
			Effect:    "Allow",
			Principal: map[string]string{"Service": cfg.TrustPrincipal},
			Action:    []string{"sts:AssumeRole", "sts:TagSession"},
		}},
	}, nil
}

func generateIRSATrust(accountID string, cfg *IRSAConfig) (IAMPolicy, error) {
	if cfg == nil {
		return IAMPolicy{}, fmt.Errorf("IRSA config missing")
	}
	aud := cfg.OIDCAudience
	if aud == "" {
		aud = "sts.amazonaws.com"
	}
	sub := fmt.Sprintf("system:serviceaccount:%s:%s", cfg.ServiceNamespace, cfg.ServiceAccount)
	return buildOIDCTrust(accountID, cfg.OIDCProvider, sub, aud), nil
}

func generateSPIFFETrust(accountID string, cfg *SPIFFEConfig) (IAMPolicy, error) {
	if cfg == nil {
		return IAMPolicy{}, fmt.Errorf("SPIFFE config missing")
	}
	aud := cfg.OIDCAudience
	if aud == "" {
		aud = "sts.amazonaws.com"
	}
	return buildOIDCTrust(accountID, cfg.OIDCProvider, cfg.SPIFFEID, aud), nil
}

func buildOIDCTrust(accountID, provider, sub, aud string) IAMPolicy {
	oidcARN := fmt.Sprintf("arn:aws:iam::%s:oidc-provider/%s", accountID, provider)
	return IAMPolicy{
		Version: "2012-10-17",
		Statement: []IAMStatement{{
			Sid:       "TrustOIDC",
			Effect:    "Allow",
			Principal: map[string]string{"Federated": oidcARN},
			Action:    "sts:AssumeRoleWithWebIdentity",
			Condition: map[string]interface{}{
				"StringEquals": map[string]string{
					fmt.Sprintf("%s:sub", provider): sub,
					fmt.Sprintf("%s:aud", provider): aud,
				},
			},
		}},
	}
}

type IAMPolicy struct {
	Version   string         `json:"Version"`
	Statement []IAMStatement `json:"Statement"`
}
type IAMStatement struct {
	Sid       string                 `json:"Sid,omitempty"`
	Effect    string                 `json:"Effect"`
	Principal interface{}            `json:"Principal,omitempty"`
	Action    interface{}            `json:"Action"`
	Resource  interface{}            `json:"Resource,omitempty"`
	Condition map[string]interface{} `json:"Condition,omitempty"`
}

func awsString(v string) *string { return &v }

// In your AccessSubscription controller's Reconcile loop...

func (r *AccessSubscriptionReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
    // 1. Fetch AccessSubscription
    var sub AccessSubscription
    if err := r.Get(ctx, req.NamespacedName, &sub); err != nil { return ctrl.Result{}, client.IgnoreNotFound(err) }

    // 2. CALCULATE CONFIG (Mocked here - this comes from your ApplicationInstance CRD)
    // Assume you found the app runs in 2 accounts.
    actorConfigs := []iamgen.ActorConfig{
        {
            Name: sub.Name + "-actor-acct1", Namespace: sub.Namespace,
            AccountID: "111111111111", RoleName: "my-app-role",
            Type: iamgen.ActorTypeIRSA,
            IRSA: &iamgen.IRSAConfig{OIDCProvider: "oidc...", ServiceAccount: "sa", ServiceNamespace: "ns"},
             // We know the target ARN ahead of time deterministically
            TargetRoleARNs: []string{fmt.Sprintf("arn:aws:iam::%s:role/%s", "333333333333", "target-db-role")},
        },
        {
             Name: sub.Name + "-actor-acct2", Namespace: sub.Namespace,
             AccountID: "222222222222", RoleName: "my-app-role",
             // ... similar config for second account ...
             TargetRoleARNs: []string{fmt.Sprintf("arn:aws:iam::%s:role/%s", "333333333333", "target-db-role")},
        },
    }

    targetConfig := iamgen.TargetConfig{
        Name: sub.Name + "-target", Namespace: sub.Namespace,
        AccountID: "333333333333", RoleName: "target-db-role",
        Type: iamgen.TargetTypeDatabase,
        Database: &iamgen.DatabaseConfig{Region: "us-west-2", DBClusterID: "mydb", DBUser: "appuser"},
        // Trust the known ARNs of the actors
        TrustedRoleARNs: []string{
             "arn:aws:iam::111111111111:role/my-app-role",
             "arn:aws:iam::222222222222:role/my-app-role",
        },
    }

    // 3. RECONCILE ACTORS
    allActorsSynced := true
    for _, cfg := range actorConfigs {
        desiredRole, _ := iamgen.GenerateActorRole(cfg)
        // Set owner ref so deleting Subscription deletes Roles
        ctrl.SetControllerReference(&sub, desiredRole, r.Scheme)
        
        synced, err := r.applyAndCheckStatus(ctx, desiredRole)
        if err != nil { return ctrl.Result{}, err }
        if !synced { allActorsSynced = false }
    }

    // 4. RECONCILE TARGET
    desiredTarget, _ := iamgen.GenerateTargetRole(targetConfig)
    ctrl.SetControllerReference(&sub, desiredTarget, r.Scheme)
    targetSynced, err := r.applyAndCheckStatus(ctx, desiredTarget)
    if err != nil { return ctrl.Result{}, err }

    // 5. AGGREGATE STATUS
    newStatus := AccessSubscriptionStatus{
        IAMReady: allActorsSynced && targetSynced,
        // ... other status fields ...
    }
    
    // Update status if changed
    if sub.Status.IAMReady != newStatus.IAMReady {
        sub.Status = newStatus
        r.Status().Update(ctx, &sub)
    }

    return ctrl.Result{}, nil
}

// Helper to Apply ACK Role and check its 'ResourceSynced' condition
func (r *AccessSubscriptionReconciler) applyAndCheckStatus(ctx context.Context, desired *ackv1alpha1.Role) (bool, error) {
    var current ackv1alpha1.Role
    err := r.Get(ctx, types.NamespacedName{Name: desired.Name, Namespace: desired.Namespace}, &current)
    
    if errors.IsNotFound(err) {
        if err := r.Create(ctx, desired); err != nil { return false, err }
        return false, nil // Created, not yet synced
    } else if err != nil { return false, err }

    // Update if spec changed
    // Note: In prod, do a deep comparison of Spec before updating to avoid noise
    current.Spec = desired.Spec
    if err := r.Update(ctx, &current); err != nil { return false, err }

    // CHECK ACK STATUS
    for _, cond := range current.Status.Conditions {
        // ACK standard condition for "AWS resource is created and matches spec"
        if cond.Type == ackv1alpha1.ConditionTypeResourceSynced && cond.Status == corev1.ConditionTrue {
            return true, nil
        }
    }
    return false, nil
}