// Package azure provides the Microsoft Azure provider implementation.
// It implements the provider.Provider interface for managing Azure RBAC permissions.
//
// Requirements: 3.2, 3.3, 3.4, 3.6
package azure

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"sync"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization/v2"

	"AutoSSO/internal/config"
	"AutoSSO/internal/provider"
)

// AzureProvider implements the provider.Provider interface for Azure.
// It manages permissions through Azure RBAC (Role-Based Access Control).
//
// Requirements: 3.4, 3.6
type AzureProvider struct {
	// config holds the Azure configuration
	config *config.AzureConfig

	// credential is the Azure AD client secret credential
	credential *azidentity.ClientSecretCredential

	// subscriptionClients caches role assignment clients per subscription
	subscriptionClients map[string]*armauthorization.RoleAssignmentsClient

	// roleDefinitionsClients caches role definitions clients per subscription
	roleDefinitionsClients map[string]*armauthorization.RoleDefinitionsClient

	// roleDefinitionCache caches role name to role definition ID mappings
	// Key format: "subscriptionID:roleName"
	roleDefinitionCache map[string]string

	// subscriptionMap maps subscription names to their IDs and role suffixes
	subscriptionMap map[string]*config.CloudAccount

	// cacheMu protects the caches
	cacheMu sync.RWMutex

	// clientMu protects the client maps
	clientMu sync.RWMutex

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

	// closeMu protects the closed flag
	closeMu sync.RWMutex
}

// NewAzureProvider creates a new Azure provider instance from the configuration.
// It initializes the Azure AD credential and prepares the subscription mappings.
//
// Parameters:
//   - cfg: The Azure configuration containing tenant ID, client ID, client secret, and subscriptions
//
// Returns:
//   - *AzureProvider: The initialized Azure provider
//   - error: Any error that occurred during initialization
//
// Requirements: 3.4, 3.6
func NewAzureProvider(cfg *config.AzureConfig) (*AzureProvider, error) {
	if cfg == nil {
		return nil, fmt.Errorf("Azure config cannot be nil")
	}

	if cfg.TenantID == "" {
		return nil, fmt.Errorf("Azure tenant ID is required")
	}

	if cfg.ClientID == "" {
		return nil, fmt.Errorf("Azure client ID is required")
	}

	if cfg.ClientSecret == "" {
		return nil, fmt.Errorf("Azure client secret is required")
	}

	// Create Azure AD client secret credential
	cred, err := azidentity.NewClientSecretCredential(
		cfg.TenantID,
		cfg.ClientID,
		cfg.ClientSecret,
		nil,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create Azure credential: %w", err)
	}

	// Build subscription map for quick lookup
	subscriptionMap := make(map[string]*config.CloudAccount)
	for i := range cfg.Subscriptions {
		sub := &cfg.Subscriptions[i]
		// Map by both ID and name for flexibility
		if sub.ID != "" {
			subscriptionMap[sub.ID] = sub
		}
		if sub.Name != "" {
			subscriptionMap[sub.Name] = sub
		}
	}

	return &AzureProvider{
		config:                 cfg,
		credential:             cred,
		subscriptionClients:    make(map[string]*armauthorization.RoleAssignmentsClient),
		roleDefinitionsClients: make(map[string]*armauthorization.RoleDefinitionsClient),
		roleDefinitionCache:    make(map[string]string),
		subscriptionMap:        subscriptionMap,
	}, nil
}

// Name returns the provider name.
//
// Returns:
//   - string: "Azure"
func (p *AzureProvider) Name() string {
	return string(provider.ProviderAzure)
}

// ValidateCredentials validates that the Azure credentials are valid
// by attempting to get a token from Azure AD.
//
// Parameters:
//   - ctx: Context for cancellation and timeout
//
// Returns:
//   - error: Any error if credentials are invalid or authentication fails
func (p *AzureProvider) 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 get a token to validate credentials
	_, err := p.credential.GetToken(ctx, policy.TokenRequestOptions{
		Scopes: []string{"https://management.azure.com/.default"},
	})
	if err != nil {
		return fmt.Errorf("failed to validate Azure credentials: %w", err)
	}

	return nil
}

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

	if p.closed {
		return nil
	}

	p.closed = true

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

	// Clear the client maps
	p.clientMu.Lock()
	p.subscriptionClients = make(map[string]*armauthorization.RoleAssignmentsClient)
	p.roleDefinitionsClients = make(map[string]*armauthorization.RoleDefinitionsClient)
	p.clientMu.Unlock()

	return nil
}

// GetUserID retrieves the Azure AD user Object ID for a given email address.
// It uses the Microsoft Graph API to look up the user.
//
// Parameters:
//   - ctx: Context for cancellation and timeout
//   - email: User's email address
//
// Returns:
//   - string: The Azure AD user Object ID
//   - error: Any error that occurred during the lookup
func (p *AzureProvider) GetUserID(ctx context.Context, email string) (string, error) {
	p.closeMu.RLock()
	if p.closed {
		p.closeMu.RUnlock()
		return "", fmt.Errorf("provider is closed")
	}
	p.closeMu.RUnlock()

	if email == "" {
		return "", fmt.Errorf("email is required")
	}

	// Normalize the email
	email = strings.TrimSpace(email)

	// Get a token for Microsoft Graph API
	token, err := p.credential.GetToken(ctx, policy.TokenRequestOptions{
		Scopes: []string{"https://graph.microsoft.com/.default"},
	})
	if err != nil {
		return "", fmt.Errorf("failed to get Graph API token: %w", err)
	}

	// Call Microsoft Graph API to get user by email
	url := fmt.Sprintf("https://graph.microsoft.com/v1.0/users/%s", email)
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+token.Token)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to get user: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("failed to get user: %s", resp.Status)
	}

	var result struct {
		ID string `json:"id"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", fmt.Errorf("failed to decode response: %w", err)
	}

	return result.ID, nil
}

// getRoleAssignmentsClient returns a role assignments client for the given subscription.
// It caches clients to avoid creating new ones for each operation.
//
// Parameters:
//   - subscriptionID: The Azure subscription ID
//
// Returns:
//   - *armauthorization.RoleAssignmentsClient: The role assignments client
//   - error: Any error that occurred during client creation
func (p *AzureProvider) getRoleAssignmentsClient(subscriptionID string) (*armauthorization.RoleAssignmentsClient, error) {
	p.clientMu.RLock()
	if client, exists := p.subscriptionClients[subscriptionID]; exists {
		p.clientMu.RUnlock()
		return client, nil
	}
	p.clientMu.RUnlock()

	// Create a new client
	client, err := armauthorization.NewRoleAssignmentsClient(subscriptionID, p.credential, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create role assignments client: %w", err)
	}

	// Cache the client
	p.clientMu.Lock()
	p.subscriptionClients[subscriptionID] = client
	p.clientMu.Unlock()

	return client, nil
}

// getRoleDefinitionsClient returns a role definitions client.
// It caches the client to avoid creating new ones for each operation.
//
// Returns:
//   - *armauthorization.RoleDefinitionsClient: The role definitions client
//   - error: Any error that occurred during client creation
func (p *AzureProvider) getRoleDefinitionsClient() (*armauthorization.RoleDefinitionsClient, error) {
	p.clientMu.RLock()
	if client, exists := p.roleDefinitionsClients["default"]; exists {
		p.clientMu.RUnlock()
		return client, nil
	}
	p.clientMu.RUnlock()

	// Create a new client
	client, err := armauthorization.NewRoleDefinitionsClient(p.credential, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create role definitions client: %w", err)
	}

	// Cache the client
	p.clientMu.Lock()
	p.roleDefinitionsClients["default"] = client
	p.clientMu.Unlock()

	return client, nil
}

// getRoleDefinitionID retrieves the role definition ID for a given role name.
// It caches the mapping to avoid repeated API calls.
//
// Parameters:
//   - ctx: Context for cancellation and timeout
//   - subscriptionID: The Azure subscription ID
//   - roleName: The name of the role
//
// Returns:
//   - string: The role definition ID
//   - error: Any error that occurred during the lookup
func (p *AzureProvider) getRoleDefinitionID(ctx context.Context, subscriptionID, roleName string) (string, error) {
	// Check cache first
	cacheKey := fmt.Sprintf("%s:%s", subscriptionID, roleName)
	p.cacheMu.RLock()
	if roleID, exists := p.roleDefinitionCache[cacheKey]; exists {
		p.cacheMu.RUnlock()
		return roleID, nil
	}
	p.cacheMu.RUnlock()

	// Get the role definitions client
	roleClient, err := p.getRoleDefinitionsClient()
	if err != nil {
		return "", err
	}

	// List role definitions and find the matching one
	scope := fmt.Sprintf("/subscriptions/%s", subscriptionID)
	pager := roleClient.NewListPager(scope, nil)
	
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return "", fmt.Errorf("failed to list role definitions: %w", err)
		}

		// Local cache to batch updates and avoid lock contention
		localCache := make(map[string]string)
		var foundRoleID string

		for _, role := range page.Value {
			if role.Properties != nil && role.Properties.RoleName != nil {
				roleNameStr := *role.Properties.RoleName
				roleIDStr := *role.ID
				
				roleCacheKey := fmt.Sprintf("%s:%s", subscriptionID, roleNameStr)
				localCache[roleCacheKey] = roleIDStr

				if roleNameStr == roleName {
					foundRoleID = roleIDStr
				}
			}
		}

		// Batch update the global cache
		if len(localCache) > 0 {
			p.cacheMu.Lock()
			for k, v := range localCache {
				p.roleDefinitionCache[k] = v
			}
			p.cacheMu.Unlock()
		}

		if foundRoleID != "" {
			return foundRoleID, nil
		}
	}

	return "", fmt.Errorf("role '%s' not found in subscription %s", roleName, subscriptionID)
}

// resolveSubscription resolves a subscription identifier (name or ID) to its ID and configuration.
// It handles the mapping from subscription names to IDs and role suffixes.
//
// Parameters:
//   - subscriptionIdentifier: The subscription name or ID
//
// Returns:
//   - subscriptionID: The resolved subscription ID
//   - roleSuffix: The role suffix for this subscription (if any)
//   - error: Any error that occurred during resolution
func (p *AzureProvider) resolveSubscription(subscriptionIdentifier string) (subscriptionID string, roleSuffix string, err error) {
	subscriptionIdentifier = strings.TrimSpace(subscriptionIdentifier)

	// Look up in the subscription map
	if sub, exists := p.subscriptionMap[subscriptionIdentifier]; exists {
		return sub.ID, sub.RoleSuffix, nil
	}

	// If not found in map, assume it's already a subscription ID
	// and return it with no role suffix
	return subscriptionIdentifier, "", nil
}

// applyRoleSuffix applies the subscription-specific role suffix to a role name.
// This handles the Azure-specific role naming convention where different subscriptions
// may have different role name suffixes.
//
// Parameters:
//   - roleName: The base role name
//   - roleSuffix: The suffix to apply
//
// Returns:
//   - string: The role name with suffix applied
func (p *AzureProvider) applyRoleSuffix(roleName, roleSuffix string) string {
	if roleSuffix == "" {
		return roleName
	}
	return roleName + roleSuffix
}

// GetSubscriptionMap returns a copy of the subscription map for testing purposes.
//
// Returns:
//   - map[string]*config.CloudAccount: A copy of the subscription map
func (p *AzureProvider) GetSubscriptionMap() map[string]*config.CloudAccount {
	result := make(map[string]*config.CloudAccount, len(p.subscriptionMap))
	for k, v := range p.subscriptionMap {
		// Create a copy of the CloudAccount
		accountCopy := *v
		result[k] = &accountCopy
	}
	return result
}