package server

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

	"cloud.google.com/go/storage"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/samber/lo"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/api/impersonate"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/utils/ptr"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	"github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/cloud/gcp"
)

const (
	AWSS3Scheme    = "s3"
	GCPGCSScheme   = "gcs"
	AzureABSScheme = "https"
)

var (
	preSignedURLExpireDuration             = time.Minute * 10
	maxObjectCountInOneListOperation int32 = 1000
	systemFilenameSet                      = map[string]bool{
		"_test-permission": true,
	}
)

type AuditLogViewer interface {
	List(ctx context.Context, cluster *cloudv1.Cluster, req *apiv1.ListAuditLogsRequest) (*apiv1.ListAuditLogsResponse, error)
	PreSignForDownloading(ctx context.Context, cluster *cloudv1.Cluster, uri *url.URL) (*apiv1.GetAuditLogDownloadURLResponse, error)
	ParseURI(uri string) (*url.URL, error)
}

var (
	_ AuditLogViewer = &awsAuditLogViewer{}
	_ AuditLogViewer = &gcpAuditLogViewer{}
	_ AuditLogViewer = &azureAuditLogViewer{}
	_ AuditLogViewer = &dummyAuditLogViewer{}
)

func (s *ClusterServer) ListAuditLogs(ctx context.Context, req *apiv1.ListAuditLogsRequest) (*apiv1.ListAuditLogsResponse, error) {
	cluster, err := s.checkSharedBucket(ctx, req.ClusterId)
	if err != nil {
		return nil, err
	}
	if err := s.checkAuditLogEnabled(ctx, cluster); err != nil {
		return nil, err
	}

	// As all files of one pager may all be filtered out, so that we do iterate manually here.
	out := &apiv1.ListAuditLogsResponse{}
	for {
		// list one page
		resp, err := s.auditLogViewer.List(ctx, cluster, req)
		if err != nil {
			return nil, err
		}
		// do filter locally
		resp.Files = lo.Filter(resp.Files, func(item *apiv1.AuditLog, index int) bool {
			if req.Filter == nil {
				return true
			}
			if req.Filter.CreatedAtGte != nil && item.CreatedAt.AsTime().Before(req.Filter.CreatedAtGte.AsTime()) {
				return false
			}
			if req.Filter.CreatedAtLt != nil && item.CreatedAt.AsTime().After(req.Filter.CreatedAtLt.AsTime()) {
				return false
			}

			return true
		})

		// save results
		out.ContinuationToken = resp.ContinuationToken
		out.Files = append(out.Files, resp.Files...)

		// break if no more pages or reaching the max limit
		if len(out.Files) >= int(maxObjectCountInOneListOperation) || lo.FromPtr(resp.ContinuationToken) == "" {
			break
		}
		// reset ContinuationToken
		req.ContinuationToken = resp.ContinuationToken
	}

	return out, nil
}

func (s *ClusterServer) GetAuditLogDownloadURL(ctx context.Context, req *apiv1.GetAuditLogDownloadURLRequest) (*apiv1.GetAuditLogDownloadURLResponse, error) {
	cluster, err := s.checkSharedBucket(ctx, req.ClusterId)
	if err != nil {
		return nil, err
	}
	if err := s.checkAuditLogEnabled(ctx, cluster); err != nil {
		return nil, err
	}

	u, err := s.auditLogViewer.ParseURI(req.Uri)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid uri: %s", err)
	}
	if err := s.validateDBAuditLogPrefix(ctx, cluster, u); err != nil {
		return nil, err
	}

	return s.auditLogViewer.PreSignForDownloading(ctx, cluster, u)
}

func (s *ClusterServer) validateDBAuditLogPrefix(ctx context.Context, cluster *cloudv1.Cluster, uri *url.URL) error {
	logger := tilog.WithContextL(ctx)

	// ensure the requested URI is under the subdirectory of bucket URL for audit log
	if !strings.HasPrefix(uri.String(), cluster.Spec.Cluster.TiDB.Audit.BucketURL) {
		logger.Info("mismatch uri with bucketUrl", zap.String("uri", uri.String()), zap.String("bucketUrl", cluster.Spec.Cluster.TiDB.Audit.BucketURL))
		return status.Errorf(codes.InvalidArgument, "mismatch prefix with your provided URI")
	}
	return nil
}

func (s *ClusterServer) checkSharedBucket(ctx context.Context, clusterId string) (*cloudv1.Cluster, error) {
	cluster, err := s.clusterReader.Get(ctx, s.KubeClient, clusterId)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil, status.Errorf(codes.NotFound, "Cluster %s not found", clusterId)
		}
		return nil, status.Errorf(codes.InvalidArgument, "Failed to get cluster %s: %s", clusterId, err)
	}
	if !cluster.SharedBucketEnabled() {
		return nil, status.Errorf(codes.InvalidArgument, "Shared Bucket is not enabeld for cluster: %s", clusterId)
	}

	return cluster, nil
}

func (s *ClusterServer) checkAuditLogEnabled(ctx context.Context, cluster *cloudv1.Cluster) error {
	if cluster.Spec.Cluster.TiDB.Audit == nil ||
		cluster.Spec.Cluster.TiDB.Audit.BucketURL == "" {
		return status.Errorf(codes.InvalidArgument, "DB audit log is not enabeld for cluster: %s", cluster.Name)
	}

	return nil
}

func getAuditLogListPrefix(ctx context.Context, bucketUrl *url.URL) string {
	// the prefix to list objects from bucket should be trim leading slash
	withoutLeadingSlash, _ := strings.CutPrefix(path.Clean(bucketUrl.Path), "/")
	withTailingSlash := withoutLeadingSlash + "/"
	return withTailingSlash
}

type awsAuditLogViewer struct {
	awsCli aws.Client
}

func NewAwsAuditLogViewer(awsCli aws.Client) *awsAuditLogViewer {
	return &awsAuditLogViewer{awsCli: awsCli}
}

func (v *awsAuditLogViewer) List(ctx context.Context, cluster *cloudv1.Cluster, req *apiv1.ListAuditLogsRequest) (*apiv1.ListAuditLogsResponse, error) {
	logger := tilog.WithContextL(ctx)
	if cluster.Status.SharedBucket == nil ||
		cluster.Status.SharedBucket.AWS == nil ||
		cluster.Status.SharedBucket.AWS.ReadonlyRoleARN == "" {
		return nil, status.Errorf(codes.InvalidArgument, "readonly IAM role of shared bucket is not ready")
	}

	s3Cli := v.awsCli.S3(&aws.AssumedRole{
		RoleARN:    cluster.Status.SharedBucket.AWS.ReadonlyRoleARN,
		ExternalID: ptr.To(cluster.Status.NodeIdentity.AWS.ExternalID),
	})

	u, err := v.ParseURI(cluster.Spec.Cluster.TiDB.Audit.BucketURL)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "invalid bucket URL: %s", cluster.Spec.Cluster.TiDB.Audit.BucketURL)
	}
	bucketName := u.Host
	params := &s3.ListObjectsV2Input{
		Bucket:  ptr.To(bucketName),
		Prefix:  ptr.To(getAuditLogListPrefix(ctx, u)),
		MaxKeys: ptr.To(maxObjectCountInOneListOperation),
	}
	if lo.FromPtr(req.ContinuationToken) != "" {
		params.ContinuationToken = req.ContinuationToken
	}
	logger.Debug("inputs for list objects", zap.Any("inputs", params))
	response, err := s3Cli.ListObjectsV2(ctx, params)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "list objects from AWS: %s", err)
	}

	out := &apiv1.ListAuditLogsResponse{
		ContinuationToken: response.NextContinuationToken,
	}
	for _, content := range response.Contents {
		_, filename := filepath.Split(ptr.Deref(content.Key, ""))
		uri := fmt.Sprintf("%s://%s/%s", AWSS3Scheme, bucketName, ptr.Deref(content.Key, ""))

		// fileter out system files
		if systemFilenameSet[filename] {
			continue
		}

		out.Files = append(out.Files, &apiv1.AuditLog{
			Filename:   filename,
			Uri:        uri,
			BucketName: bucketName,
			Size:       ptr.Deref(content.Size, 0),
			CreatedAt:  timestamppb.New(lo.FromPtr(content.LastModified)),
		})
	}

	return out, nil
}

func (v *awsAuditLogViewer) PreSignForDownloading(ctx context.Context, cluster *cloudv1.Cluster, u *url.URL) (*apiv1.GetAuditLogDownloadURLResponse, error) {
	if cluster.Status.SharedBucket == nil ||
		cluster.Status.SharedBucket.AWS == nil ||
		cluster.Status.SharedBucket.AWS.ReadonlyRoleARN == "" {
		return nil, status.Errorf(codes.InvalidArgument, "readonly IAM role of shared bucket is not ready")
	}

	s3Cli := v.awsCli.S3(&aws.AssumedRole{
		RoleARN:    cluster.Status.SharedBucket.AWS.ReadonlyRoleARN,
		ExternalID: ptr.To(cluster.Status.NodeIdentity.AWS.ExternalID),
	})
	presignClient := s3.NewPresignClient(s3Cli)

	bucketName := u.Host

	params := &s3.GetObjectInput{
		Bucket: ptr.To(bucketName),
		Key:    ptr.To(strings.TrimLeft(u.Path, "/")),
	}
	response, err := presignClient.PresignGetObject(ctx, params, func(opts *s3.PresignOptions) {
		opts.Expires = preSignedURLExpireDuration
	})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "pre sign object for AWS: %s", err)
	}

	return &apiv1.GetAuditLogDownloadURLResponse{
		PresignedUrl: response.URL,
	}, nil
}

func (v *awsAuditLogViewer) ParseURI(uri string) (*url.URL, error) {
	u, err := url.Parse(uri)
	if err != nil {
		return nil, fmt.Errorf("invalid uri: %s", err)
	}
	if u.Scheme != AWSS3Scheme {
		return nil, fmt.Errorf("invalid scheme: %s", u.Scheme)
	}
	return u, nil
}

type gcpAuditLogViewer struct {
	gcpClient gcp.Client
}

func NewGcpAuditLogViewer(gcpCli gcp.Client) *gcpAuditLogViewer {
	return &gcpAuditLogViewer{gcpClient: gcpCli}
}

func (v *gcpAuditLogViewer) List(ctx context.Context, cluster *cloudv1.Cluster, req *apiv1.ListAuditLogsRequest) (*apiv1.ListAuditLogsResponse, error) {
	logger := tilog.WithContextL(ctx)
	if cluster.Status.SharedBucket == nil ||
		cluster.Status.SharedBucket.GCP == nil ||
		cluster.Status.SharedBucket.GCP.ReadonlySA == "" {
		return nil, status.Errorf(codes.InvalidArgument, "readonly IAM SA of shared bucket is not ready")
	}

	tokenSource, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
		TargetPrincipal: cluster.Status.SharedBucket.GCP.ReadonlySA,
		Scopes:          []string{"https://www.googleapis.com/auth/devstorage.read_only"},
	})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "generate token source")
	}
	gcpCli, err := v.gcpClient.NewStorageCli(ctx, option.WithTokenSource(tokenSource))
	if err != nil {
		return nil, status.Errorf(codes.Internal, "new client to GCP storage")
	}
	defer gcpCli.Close() //nolint:errcheck // TODO: handle GCP client close error

	u, err := v.ParseURI(cluster.Spec.Cluster.TiDB.Audit.BucketURL)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "invalid bucket URL: %s", cluster.Spec.Cluster.TiDB.Audit.BucketURL)
	}
	bucketName := u.Host
	query := &storage.Query{
		Prefix: getAuditLogListPrefix(ctx, u),
	}
	if req.ContinuationToken != nil {
		query.StartOffset = ptr.Deref(req.ContinuationToken, "")
	}
	logger.Debug("inputs for list objects", zap.String("bucket", bucketName), zap.Any("inputs", query))
	out := &apiv1.ListAuditLogsResponse{}
	it := gcpCli.Bucket(bucketName).Objects(ctx, query)
	for {
		attrs, err := it.Next()
		if errors.Is(err, iterator.Done) {
			break
		}
		if err != nil {
			return nil, status.Errorf(codes.Internal, "list in bucket: %s", err)
		}

		_, filename := filepath.Split(attrs.Name)
		uri := fmt.Sprintf("%s://%s/%s", GCPGCSScheme, bucketName, attrs.Name)

		// fileter out system files
		if systemFilenameSet[filename] {
			continue
		}

		out.Files = append(out.Files, &apiv1.AuditLog{
			Filename:   filename,
			Uri:        uri,
			BucketName: bucketName,
			Size:       attrs.Size,
			CreatedAt:  timestamppb.New(attrs.Updated),
		})
		pi := it.PageInfo()
		if token := pi.Token; token != "" {
			out.ContinuationToken = &token
		}
		// break when the count retrieved objects reach the limit and all objects in the current page have been retrieved.
		// Retrieving remaining objects in current page even if reached the maxObjectCountInOneListOperation
		// because the token is created per page.
		if len(out.Files) >= int(maxObjectCountInOneListOperation) && pi.Remaining() == 0 {
			break
		}
	}

	return out, nil
}

func (v *gcpAuditLogViewer) PreSignForDownloading(ctx context.Context, cluster *cloudv1.Cluster, u *url.URL) (*apiv1.GetAuditLogDownloadURLResponse, error) {
	if cluster.Status.SharedBucket == nil ||
		cluster.Status.SharedBucket.GCP == nil ||
		cluster.Status.SharedBucket.GCP.ReadonlySA == "" {
		return nil, status.Errorf(codes.InvalidArgument, "readonly IAM SA of shared bucket is not ready")
	}

	gcpCli := v.gcpClient.StorageCli()
	bucketName := u.Host
	objectFullPath := strings.TrimLeft(u.Path, "/")

	opts := &storage.SignedURLOptions{
		// Refer to https://pkg.go.dev/cloud.google.com/go/storage#hdr-Credential_requirements_for_signing
		// Detecting GoogleAccessID may not be possible if you are authenticated using a token source or using
		// option.WithHTTPClient. In this case, you can provide a service account email for GoogleAccessID and
		// the client will attempt to sign the URL or Post Policy using that service account.
		//
		// In other words, the IAM SA used by cluster-server should have `Service Account Token Creator` access to
		// the IAM SA specified by `GoogleAccessID`.
		GoogleAccessID: cluster.Status.SharedBucket.GCP.ReadonlySA,
		Scheme:         storage.SigningSchemeV4,
		Method:         http.MethodGet,
		Expires:        time.Now().Add(preSignedURLExpireDuration),
	}
	preSignedURL, err := gcpCli.Bucket(bucketName).SignedURL(objectFullPath, opts)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "pre sign object for GCP: %s", err)
	}

	return &apiv1.GetAuditLogDownloadURLResponse{
		PresignedUrl: preSignedURL,
	}, nil
}

func (v *gcpAuditLogViewer) ParseURI(uri string) (*url.URL, error) {
	u, err := url.Parse(uri)
	if err != nil {
		return nil, fmt.Errorf("invalid uri: %s", err)
	}
	if u.Scheme != GCPGCSScheme {
		return nil, fmt.Errorf("invalid scheme: %s", u.Scheme)
	}
	return u, nil
}

type azureAuditLogViewer struct {
}

func NewAzureAuditLogViewer() *azureAuditLogViewer {
	return &azureAuditLogViewer{}
}

func (v *azureAuditLogViewer) List(ctx context.Context, cluster *cloudv1.Cluster, req *apiv1.ListAuditLogsRequest) (*apiv1.ListAuditLogsResponse, error) {
	return nil, status.Errorf(codes.InvalidArgument, "Azure is not supported yet")
}

func (v *azureAuditLogViewer) PreSignForDownloading(ctx context.Context, cluster *cloudv1.Cluster, u *url.URL) (*apiv1.GetAuditLogDownloadURLResponse, error) {
	return nil, status.Errorf(codes.InvalidArgument, "Azure is not supported yet")
}

func (v *azureAuditLogViewer) ParseURI(uri string) (*url.URL, error) {
	u, err := url.Parse(uri)
	if err != nil {
		return nil, fmt.Errorf("invalid uri: %s", err)
	}
	if u.Scheme != AzureABSScheme {
		return nil, fmt.Errorf("invalid scheme: %s", u.Scheme)
	}
	return u, nil
}

type dummyAuditLogViewer struct {
}

func NewDummyAuditLogViewer() *dummyAuditLogViewer {
	return &dummyAuditLogViewer{}
}

func (v *dummyAuditLogViewer) List(ctx context.Context, cluster *cloudv1.Cluster, req *apiv1.ListAuditLogsRequest) (*apiv1.ListAuditLogsResponse, error) {
	return &apiv1.ListAuditLogsResponse{}, nil
}

func (v *dummyAuditLogViewer) PreSignForDownloading(ctx context.Context, cluster *cloudv1.Cluster, uri *url.URL) (*apiv1.GetAuditLogDownloadURLResponse, error) {
	return &apiv1.GetAuditLogDownloadURLResponse{PresignedUrl: "url"}, nil
}

func (v *dummyAuditLogViewer) ParseURI(uri string) (*url.URL, error) {
	u, err := url.Parse(uri)
	if err != nil {
		return nil, fmt.Errorf("invalid uri: %s", err)
	}
	return u, nil
}