package storage

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/url"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
	"github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss"
	osscredentials "github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss/credentials"
	alicredentials "github.com/aliyun/credentials-go/credentials"
	"github.com/aws/aws-sdk-go-v2/aws"
	s3sdk "github.com/aws/aws-sdk-go-v2/service/s3"
	"google.golang.org/api/impersonate"
	"google.golang.org/api/option"

	"github.com/tidbcloud/service-proxy/pkg/checker/core"
	"github.com/tidbcloud/service-proxy/pkg/checker/util"
)

func init() {
	core.DefaultRegistry.Register("audit", func(req *core.CheckRequest) (core.Checker, error) {
		if req.Audit == nil {
			return nil, errors.New("audit params required")
		}
		return NewAuditChecker(req.Audit), nil
	})
}

const checkFileName = "_test-permission"

// AuditChecker checks audit log storage connectivity.
type AuditChecker struct {
	params *core.AuditParams
}

// NewAuditChecker creates a new Audit checker.
func NewAuditChecker(params *core.AuditParams) core.Checker {
	return &AuditChecker{params: params}
}

// Type returns the checker type.
func (c *AuditChecker) Type() string {
	return "audit"
}

// Check performs audit storage connectivity check.
func (c *AuditChecker) Check(ctx context.Context) error {
	u, err := url.Parse(c.params.BucketURL)
	if err != nil {
		return fmt.Errorf("failed to parse bucket url: %w", err)
	}

	switch u.Scheme {
	case "s3":
		return c.checkS3(ctx, u)
	case "gcs":
		return c.checkGCS(ctx, u)
	case "oss":
		return c.checkOSS(ctx, u)
	case "https":
		return c.checkAzBlob(ctx, u)
	default:
		return fmt.Errorf("unknown scheme %q in bucket url", u.Scheme)
	}
}

func (c *AuditChecker) checkS3(ctx context.Context, u *url.URL) error {
	if c.params.AWS == nil {
		return errors.New("aws auth params required for s3")
	}

	svc, err := util.NewS3ClientWithAssumeRole(ctx, c.params.Region, []util.AssumeRole{
		{RoleARN: c.params.AWS.RoleARN, ExternalID: c.params.AWS.ExternalID},
	})
	if err != nil {
		return fmt.Errorf("failed to create s3 client: %w", err)
	}

	_, err = svc.PutObject(ctx, &s3sdk.PutObjectInput{
		Bucket: aws.String(u.Host),
		Key:    aws.String(targetKey(u)),
		Body:   bytes.NewBuffer(nil),
	})
	if err != nil {
		return fmt.Errorf("failed to put object to s3: %w", err)
	}
	return nil
}

func (c *AuditChecker) checkGCS(ctx context.Context, u *url.URL) error {
	if c.params.GCP == nil {
		return errors.New("gcp auth params required for gcs")
	}

	// Base credentials sourced from ADC or provided client options.
	ts, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
		TargetPrincipal: c.params.GCP.ServiceAccount,
		Scopes:          []string{"https://www.googleapis.com/auth/cloud-platform"},
	})
	if err != nil {
		return fmt.Errorf("failed to get token source: %w", err)
	}

	client, err := storage.NewClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		return fmt.Errorf("failed to new gcp storage client: %w", err)
	}

	wc := client.Bucket(u.Host).Object(targetKey(u)).NewWriter(ctx)
	return wc.Close()
}

func (c *AuditChecker) checkAzBlob(ctx context.Context, u *url.URL) error {
	if c.params.Azure == nil {
		return errors.New("azure auth params required for azure blob storage")
	}

	containerURL, containerName, err := azBlobContainerURL(u, c.params.Azure.SASToken)
	if err != nil {
		return err
	}

	var containerClient *container.Client
	if strings.TrimSpace(c.params.Azure.SASToken) != "" {
		containerClient, err = container.NewClientWithNoCredential(containerURL, nil)
		if err != nil {
			return fmt.Errorf("failed to build azure container client with SAS: %w", err)
		}
	} else {
		credOpts := &azidentity.ManagedIdentityCredentialOptions{}
		clientID := strings.TrimSpace(c.params.Azure.ClientID)
		if clientID == "" {
			return fmt.Errorf("azure client ID is required when using managed identity auth")
		}
		credOpts.ID = azidentity.ClientID(clientID)

		cred, credErr := azidentity.NewManagedIdentityCredential(credOpts)
		if credErr != nil {
			return fmt.Errorf("failed to build azure managed identity credential: %w", credErr)
		}

		containerClient, err = container.NewClient(containerURL, cred, nil)
		if err != nil {
			return fmt.Errorf("failed to build azure container client with managed identity: %w", err)
		}
	}

	// targetKey starts with container name, blob name should exclude it
	blobName := strings.TrimPrefix(targetKey(u), containerName+"/")
	blobClient := containerClient.NewBlockBlobClient(blobName)
	var buf []byte
	_, err = blobClient.UploadBuffer(ctx, buf, nil)
	if err != nil {
		return fmt.Errorf("failed to upload blob to azure container: %w", err)
	}
	return nil
}

func azBlobContainerURL(u *url.URL, sasToken string) (string, string, error) {
	s := strings.SplitN(u.Path, "/", 3)
	if len(s) < 2 || s[1] == "" {
		return "", "", fmt.Errorf("no container specified in bucket url")
	}

	containerName := s[1]
	containerURL := fmt.Sprintf("%s://%s/%s", u.Scheme, u.Host, containerName)
	if strings.TrimSpace(sasToken) != "" {
		containerURL = fmt.Sprintf("%s?%s", containerURL, sasToken)
	}
	return containerURL, containerName, nil
}

func (c *AuditChecker) checkOSS(ctx context.Context, u *url.URL) error {
	if c.params.AliCloud == nil {
		return errors.New("alicloud auth params required for oss")
	}

	// Force use ECS Role as base credentials for AssumeRole
	provider, err := alicredentials.NewCredential(new(alicredentials.Config).SetType("ecs_ram_role"))
	if err != nil {
		return fmt.Errorf("failed to create ecs role credential: %w", err)
	}
	cred, err := provider.GetCredential()
	if err != nil {
		return fmt.Errorf("failed to get ecs role credential: %w", err)
	}

	sessionName := "audit-check"
	if s := c.params.AliCloud.RoleSessionName; s != "" {
		sessionName = s
	}
	credConfig := new(alicredentials.Config).
		SetType("ram_role_arn").
		SetAccessKeyId(*cred.AccessKeyId).
		SetAccessKeySecret(*cred.AccessKeySecret).
		SetSecurityToken(*cred.SecurityToken).
		SetRoleArn(c.params.AliCloud.RoleARN).
		SetExternalId(c.params.AliCloud.ExternalID).
		SetRoleSessionName(sessionName)
	arnCredential, err := alicredentials.NewCredential(credConfig)
	if err != nil {
		return fmt.Errorf("failed to create arn credential: %w", err)
	}

	credProvider := osscredentials.CredentialsProviderFunc(func(ctx context.Context) (osscredentials.Credentials, error) {
		cred2, err := arnCredential.GetCredential()
		if err != nil {
			return osscredentials.Credentials{}, err
		}
		return osscredentials.Credentials{
			AccessKeyID:     *cred2.AccessKeyId,
			AccessKeySecret: *cred2.AccessKeySecret,
			SecurityToken:   *cred2.SecurityToken,
		}, nil
	})

	// Configure OSS client
	cfg := oss.LoadDefaultConfig().
		WithCredentialsProvider(credProvider).
		WithRegion(c.params.Region)
	client := oss.NewClient(cfg)

	// Test write permission by uploading an empty object
	_, err = client.PutObject(ctx, &oss.PutObjectRequest{
		Bucket: oss.Ptr(u.Host),
		Key:    oss.Ptr(targetKey(u)),
		Body:   bytes.NewBuffer(nil),
	})
	if err != nil {
		return fmt.Errorf("failed to put object to oss: %w", err)
	}
	return nil
}

func targetKey(u *url.URL) string {
	prefix := u.Path
	prefix = strings.TrimLeft(prefix, "/")
	return filepath.Join(prefix, checkFileName)
}