package assume_role

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"github.com/pingcap-inc/pchub/pkg/log"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization"
	"github.com/pkg/errors"
)

// AzureRoleManager 管理 Azure subscription 角色分配
type AzureRoleManager struct {
	client         *armauthorization.RoleAssignmentsClient
	subscriptionID string
}

// RoleAssignment 角色分配信息
type RoleAssignment struct {
	PrincipalID    string // 用户、组或服务主体的 Object ID
	RoleDefinition string // 角色定义 ID 或名称
	Scope          string // 分配范围
}

// AzureCredentialConfig 凭据配置
type AzureCredentialConfig struct {
	// 用于服务主体认证
	ClientID     string
	ClientSecret string
	TenantID     string

	// 用于托管身份认证
	UseManagedIdentity bool
	ManagedIdentityID  string // 可选，用户分配的托管身份 ID

	// 用于默认凭据链
	UseDefaultCredential bool
}

// AzureUser represents an Azure AD user
type AzureUser struct {
	ObjectId          string
	Email             string
	UserPrincipalName string
}

// GraphUsersResponse represents the response from Microsoft Graph API
type GraphUsersResponse struct {
	ODataContext  string      `json:"@odata.context"`
	ODataNextLink string      `json:"@odata.nextLink"`
	Value         []GraphUser `json:"value"`
}

// GraphUser represents a user from Microsoft Graph API response
type GraphUser struct {
	ID                string `json:"id"`
	Mail              string `json:"mail"`
	UserPrincipalName string `json:"userPrincipalName"`
	DisplayName       string `json:"displayName"`
	AccountEnabled    bool   `json:"accountEnabled"`
}

// NewAzureRoleManager 创建新的 Azure 角色管理器
func NewAzureRoleManager(subscriptionID string) (*AzureRoleManager, error) {
	return NewAzureRoleManagerWithConfig(subscriptionID, AzureCredentialConfig{
		UseDefaultCredential: true,
	})
}

// NewAzureRoleManagerWithConfig 使用指定配置创建 Azure 角色管理器
func NewAzureRoleManagerWithConfig(subscriptionID string, config AzureCredentialConfig) (*AzureRoleManager, error) {
	if subscriptionID == "" {
		return nil, errors.New("subscription ID cannot be empty")
	}

	// 根据配置选择合适的凭据提供者
	cred, err := getAzureCredential(config)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create Azure credentials")
	}

	// 创建角色分配客户端
	client, err := armauthorization.NewRoleAssignmentsClient(subscriptionID, cred, nil)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create role assignments client")
	}

	return &AzureRoleManager{
		client:         client,
		subscriptionID: subscriptionID,
	}, nil
}

// getAzureCredential 根据配置获取 Azure 凭据
func getAzureCredential(config AzureCredentialConfig) (azcore.TokenCredential, error) {
	// 1. 优先使用托管身份
	if config.UseManagedIdentity {
		opts := &azidentity.ManagedIdentityCredentialOptions{}
		if config.ManagedIdentityID != "" {
			opts.ID = azidentity.ClientID(config.ManagedIdentityID)
		}
		cred, err := azidentity.NewManagedIdentityCredential(opts)
		if err != nil {
			return nil, errors.Wrap(err, "failed to create managed identity credential")
		}
		return cred, nil
	}

	// 2. 使用服务主体
	if config.ClientID != "" && config.ClientSecret != "" && config.TenantID != "" {
		cred, err := azidentity.NewClientSecretCredential(
			config.TenantID,
			config.ClientID,
			config.ClientSecret,
			nil,
		)
		if err != nil {
			return nil, errors.Wrap(err, "failed to create client secret credential")
		}
		return cred, nil
	}

	// 3. 使用默认凭据链（适用于开发环境和多种托管环境）
	if config.UseDefaultCredential {
		cred, err := azidentity.NewDefaultAzureCredential(nil)
		if err != nil {
			return nil, errors.Wrap(err, "failed to create default Azure credential")
		}
		return cred, nil
	}

	return nil, errors.New("no valid credential configuration provided")
}

// AssignSubscriptionRole 给用户在订阅级别分配角色
func (arm *AzureRoleManager) AssignSubscriptionRole(ctx context.Context, assignment RoleAssignment) (*armauthorization.RoleAssignment, error) {
	if assignment.PrincipalID == "" {
		return nil, errors.New("principal ID cannot be empty")
	}
	if assignment.RoleDefinition == "" {
		return nil, errors.New("role definition ID cannot be empty")
	}

	// 设置作用域为订阅级别
	scope := fmt.Sprintf("/subscriptions/%s", arm.subscriptionID)
	if assignment.Scope != "" {
		scope = assignment.Scope
	}

	// 生成唯一的分配 ID
	assignmentName := generateUniqueID()

	// 创建角色分配参数
	params := armauthorization.RoleAssignmentCreateParameters{
		Properties: &armauthorization.RoleAssignmentProperties{
			PrincipalID:      to.Ptr(assignment.PrincipalID),
			RoleDefinitionID: to.Ptr(assignment.RoleDefinition),
		},
	}

	// 执行角色分配
	result, err := arm.client.Create(ctx, scope, assignmentName, params, nil)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create role assignment")
	}

	return &result.RoleAssignment, nil
}

// GetUserRoleAssignments 获取用户的角色分配
func (arm *AzureRoleManager) GetUserRoleAssignments(ctx context.Context, principalID string) ([]*armauthorization.RoleAssignment, error) {
	if principalID == "" {
		return nil, errors.New("principal ID cannot be empty")
	}

	scope := fmt.Sprintf("/subscriptions/%s", arm.subscriptionID)
	filter := fmt.Sprintf("principalId eq '%s'", principalID)

	var assignments []*armauthorization.RoleAssignment
	pager := arm.client.NewListForScopePager(scope, &armauthorization.RoleAssignmentsClientListForScopeOptions{
		Filter: to.Ptr(filter),
	})

	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return nil, errors.Wrap(err, "failed to get role assignments page")
		}

		for _, assignment := range page.Value {
			assignments = append(assignments, assignment)
		}
	}

	return assignments, nil
}

// RemoveSubscriptionRole 移除用户的订阅角色
func (arm *AzureRoleManager) RemoveSubscriptionRole(ctx context.Context, roleAssignmentID string) error {
	if roleAssignmentID == "" {
		return errors.New("role assignment ID cannot be empty")
	}

	_, err := arm.client.DeleteByID(ctx, roleAssignmentID, nil)
	if err != nil {
		return errors.Wrap(err, "failed to delete role assignment")
	}

	return nil
}

// generateUniqueID 生成唯一的分配 ID
func generateUniqueID() string {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}

// 常用的 Azure 内置角色定义 ID
const (
	// 订阅级别常用角色
	RoleOwner               = "/subscriptions/{subscription-id}/providers/Microsoft.Authorization/roleDefinitions/8e3af657-a8ff-443c-a75c-2fe8c4bcb635"
	RoleContributor         = "/subscriptions/{subscription-id}/providers/Microsoft.Authorization/roleDefinitions/b24988ac-6180-42a0-ab88-20f7382dd24c"
	RoleReader              = "/subscriptions/{subscription-id}/providers/Microsoft.Authorization/roleDefinitions/acdd72a7-3385-48ef-bd42-f606fba81ae7"
	RoleUserAccessAdmin     = "/subscriptions/{subscription-id}/providers/Microsoft.Authorization/roleDefinitions/18d7d88d-d35e-4fb5-a5c3-7773c20a72d9"
	RoleReaderAndDataAccess = "/subscriptions/{subscription-id}/providers/Microsoft.Authorization/roleDefinitions/c12c1c16-33a1-487b-954d-41c89c60f349"
)

// GetBuiltInRoleDefinitionID 获取内置角色定义 ID
func (arm *AzureRoleManager) GetBuiltInRoleDefinitionID(roleName string) string {
	subscriptionID := arm.subscriptionID

	switch roleName {
	case "Owner":
		return fmt.Sprintf("/subscriptions/%s/providers/Microsoft.Authorization/roleDefinitions/8e3af657-a8ff-443c-a75c-2fe8c4bcb635", subscriptionID)
	case "Contributor":
		return fmt.Sprintf("/subscriptions/%s/providers/Microsoft.Authorization/roleDefinitions/b24988ac-6180-42a0-ab88-20f7382dd24c", subscriptionID)
	case "Reader":
		return fmt.Sprintf("/subscriptions/%s/providers/Microsoft.Authorization/roleDefinitions/acdd72a7-3385-48ef-bd42-f606fba81ae7", subscriptionID)
	case "User Access Administrator":
		return fmt.Sprintf("/subscriptions/%s/providers/Microsoft.Authorization/roleDefinitions/18d7d88d-d35e-4fb5-a5c3-7773c20a72d9", subscriptionID)
	case "Reader and Data Access":
		return fmt.Sprintf("/subscriptions/%s/providers/Microsoft.Authorization/roleDefinitions/c12c1c16-33a1-487b-954d-41c89c60f349", subscriptionID)
	default:
		return roleName // 如果不是内置角色名称，假设它已经是完整的角色定义 ID
	}
}

// AzureTokenResponse Azure Token 响应结构
type AzureTokenResponse struct {
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
	AccessToken string `json:"access_token"`
}

// getAzureAccessToken 获取 Azure Access Token
func GetAzureAccessToken(ctx context.Context, tenantID, clientID, clientSecret string) (string, error) {
	tokenURL := fmt.Sprintf("https://login.microsoftonline.com/%s/oauth2/v2.0/token", tenantID)

	data := url.Values{}
	data.Set("grant_type", "client_credentials")
	data.Set("client_id", clientID)
	data.Set("client_secret", clientSecret)
	data.Set("scope", "https://graph.microsoft.com/.default")

	req, err := http.NewRequestWithContext(ctx, "POST", tokenURL, strings.NewReader(data.Encode()))
	if err != nil {
		return "", fmt.Errorf("failed to create token request: %w", err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

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

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("token request failed with status %d: %s", resp.StatusCode, string(body))
	}

	var tokenResp AzureTokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
		return "", fmt.Errorf("failed to decode token response: %w", err)
	}

	return tokenResp.AccessToken, nil
}

// getAzureUserByEmail 通过邮箱获取单个用户的 Azure Object ID
func GetAzureUserByEmail(ctx context.Context, email, accessToken string) (string, error) {
	// 使用 Microsoft Graph API 查找用户
	graphURL := fmt.Sprintf("https://graph.microsoft.com/v1.0/users/%s", url.QueryEscape(email))

	req, err := http.NewRequestWithContext(ctx, "GET", graphURL, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create graph request: %w", err)
	}

	req.Header.Set("Authorization", "Bearer "+accessToken)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to call graph api: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusNotFound {
		return "", fmt.Errorf("user not found: %s", email)
	}

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("graph api request failed with status %d: %s", resp.StatusCode, string(body))
	}

	var user GraphUser
	if err := json.NewDecoder(resp.Body).Decode(&user); err != nil {
		return "", fmt.Errorf("failed to decode graph response: %w", err)
	}

	return user.ID, nil
}

// getAzureUserObjectIDs 通过邮箱获取Azure用户的Object ID
func GetAzureUserObjectIDs(ctx context.Context, emails []string, tenantID, clientID, clientSecret string) (map[string]string, error) {
	if len(emails) == 0 {
		return make(map[string]string), nil
	}

	// 获取 Access Token
	accessToken, err := GetAzureAccessToken(ctx, tenantID, clientID, clientSecret)
	if err != nil {
		return nil, fmt.Errorf("failed to get access token: %w", err)
	}

	userObjectIDs := make(map[string]string)

	// 为每个邮箱获取 Object ID
	for _, email := range emails {
		objectID, err := GetAzureUserByEmail(ctx, email, accessToken)
		if err != nil {
			log.L(ctx).Warnf("Failed to get Azure Object ID for user %s: %v", email, err)
			// 如果获取失败，可以选择跳过这个用户或者返回错误
			// 这里选择跳过，但记录警告
			continue
		}

		userObjectIDs[email] = objectID
		log.L(ctx).Infof("Successfully got Azure Object ID for user %s: %s", email, objectID)
	}

	return userObjectIDs, nil
}

// buildAzureRoleDefinitionPath 构建完整的 Azure 角色定义路径
func BuildAzureRoleDefinitionPath(subscriptionID, roleGUID string) string {
	return fmt.Sprintf("/subscriptions/%s/providers/Microsoft.Authorization/roleDefinitions/%s", subscriptionID, roleGUID)
}