package store

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/pingcap/log"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/util"
	"go.uber.org/zap"
)

type S3RemoteStorage struct {
	s3Client *s3.Client
	bucket   string
	hostname string
}

func NewS3RemoteStorage(hostname string) (RemoteStorage, error) {
	cfg := config.Global()
	optFns := []func(*awsconfig.LoadOptions) error{
		awsconfig.WithRegion(cfg.Region),
	}
	if cfg.Stage == "local" {
		optFns = append(optFns, awsconfig.WithSharedConfigProfile("dbaas-dev"))
	}
	awscfg, err := awsconfig.LoadDefaultConfig(context.TODO(), optFns...)
	if err != nil {
		return nil, fmt.Errorf("failed to load aws config: %s", err)
	}

	cli := s3.NewFromConfig(awscfg)

	return &S3RemoteStorage{
		s3Client: cli,
		bucket:   cfg.Store.S3Bucket,
		hostname: hostname,
	}, nil
}

func (s *S3RemoteStorage) Upload(ctx context.Context, file *LocalFile) error {
	f, err := os.Open(file.Path)
	if err != nil {
		return fmt.Errorf("failed to read local metric file: %s", err)
	}
	defer f.Close()

	dir := util.LocateLatestOnTheHour(file.End)
	key := fmt.Sprintf("data/%d/%s/%s", dir, s.hostname, filepath.Base(file.Path))
	_, err = s.s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
		Body:   f,
	})
	if err != nil {
		return fmt.Errorf("failed to upload metric file to s3: %s", err)
	}
	return nil
}

func (s *S3RemoteStorage) cleanupMultipartUploadGarbage(bucket, key, uploadId string) error {
	_, err := s.s3Client.AbortMultipartUpload(context.TODO(), &s3.AbortMultipartUploadInput{
		Bucket:   aws.String(bucket),
		Key:      aws.String(key),
		UploadId: aws.String(uploadId),
	})
	if err != nil {
		return fmt.Errorf("failed to abort multipart upload: %s", err)
	}
	return nil
}

func (s *S3RemoteStorage) MultipartUpload(ctx context.Context, file *LocalFile) error {
	f, err := os.Open(file.Path)
	if err != nil {
		return fmt.Errorf("failed to read local metric file: %s", err)
	}
	defer f.Close()

	dir := util.LocateLatestOnTheHour(file.End)
	key := fmt.Sprintf("data/%d/%s/%s", dir, s.hostname, filepath.Base(file.Path))
	createResp, err := s.s3Client.CreateMultipartUpload(context.TODO(), &s3.CreateMultipartUploadInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return fmt.Errorf("failed to upload metric file to s3: %s", err)
	}

	const partSize = 5 * 1024 * 1024 //5MB
	var concurrency = 5

	partNumber := int32(1)
	var mu sync.Mutex
	var wg sync.WaitGroup
	var completedParts []types.CompletedPart

	partCh := make(chan struct {
		Body       []byte
		PartNumber int32
	}, concurrency)

	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for part := range partCh {
				partResp, err := s.s3Client.UploadPart(context.TODO(), &s3.UploadPartInput{
					Bucket:     aws.String(s.bucket),
					Key:        aws.String(key),
					UploadId:   createResp.UploadId,
					PartNumber: &part.PartNumber,
					Body:       bytes.NewReader(part.Body),
				})
				if err != nil {
					continue
				}

				mu.Lock()
				completedParts = append(completedParts, types.CompletedPart{
					ETag:       partResp.ETag,
					PartNumber: &part.PartNumber,
				})
				mu.Unlock()
			}
		}()
	}

	buf := make([]byte, partSize)
	for {
		n, err := f.Read(buf)
		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 {
			break
		}
		part := make([]byte, n)
		copy(part, buf[:n])
		partCh <- struct {
			Body       []byte
			PartNumber int32
		}{
			Body:       part,
			PartNumber: partNumber,
		}
		partNumber++
	}

	close(partCh)
	wg.Wait()

	if len(completedParts) != int(partNumber-1) {
		goto cleanup
	}

	_, err = s.s3Client.CompleteMultipartUpload(context.TODO(), &s3.CompleteMultipartUploadInput{
		Bucket:          aws.String(s.bucket),
		Key:             aws.String(key),
		UploadId:        createResp.UploadId,
		MultipartUpload: &types.CompletedMultipartUpload{Parts: completedParts},
	})
	if err != nil {
		goto cleanup
	}
	return nil

cleanup:
	err = s.cleanupMultipartUploadGarbage(s.bucket, key, *createResp.UploadId)
	if err != nil {
		return fmt.Errorf("failed to cleanup multipart upload: %s", err)
	}
	return fmt.Errorf("failed to complete multipart upload")
}

func (s *S3RemoteStorage) BuildScanner(ctx context.Context, ts int64) (RemoteScanner, error) {
	dir := util.LocateLatestOnTheHour(ts)
	// TODO: page
	output, err := s.s3Client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
		Bucket: aws.String(s.bucket),
		Prefix: aws.String(fmt.Sprintf("data/%d", dir)),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to build s3 scanner: %s", err)
	}

	ft := GetFileType()
	switch ft {
	case FileTypeJson:
		return &S3JsonScanner{
			s3Client: s.s3Client,
			bucket:   s.bucket,
			objects:  output.Contents,
			cursor:   0,
		}, nil
	case FileTypeParquet:
		return &S3ParquetScanner{
			s3Client: s.s3Client,
			bucket:   s.bucket,
			objects:  output.Contents,
			cursor:   0,
		}, nil
	default:
		return nil, fmt.Errorf("unknown file type")
	}
}

func (s *S3RemoteStorage) PutData(ctx context.Context, key string, data []byte) error {
	_, err := s.s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
		Body:   bytes.NewReader(data),
	})
	if err != nil {
		return fmt.Errorf("failed to upload data to s3: %s", err)
	}
	return nil
}

func (s *S3RemoteStorage) GetData(ctx context.Context, key string) ([]byte, error) {
	output, err := s.s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to get data from s3: %s", err)
	}
	defer output.Body.Close()

	data, err := io.ReadAll(output.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read s3 data: %s", err)
	}

	return data, nil
}

type S3JsonScanner struct {
	s3Client *s3.Client
	bucket   string
	objects  []types.Object
	cursor   int
}

func (s *S3JsonScanner) ScanMetrics(ctx context.Context) ([]*Metric, bool, error) {
	if !s.hasMore() {
		return nil, false, nil
	}
	obj := s.objects[s.cursor]

	log.Debug("get s3 object", zap.String("key", *obj.Key))
	output, err := s.s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    obj.Key,
	})
	if err != nil {
		return nil, false, fmt.Errorf("failed to get metrics file from s3: %s", err)
	}
	log.Debug("get s3 object finished", zap.String("key", *obj.Key))
	defer output.Body.Close()

	metrics := make([]*Metric, 0)
	scanner := bufio.NewScanner(output.Body)
	for scanner.Scan() {
		var metric Metric
		data := scanner.Bytes()
		err = json.Unmarshal(data, &metric)
		if err != nil {
			return nil, false, fmt.Errorf("failed to unmarshal metric file: %s", err)
		}
		metrics = append(metrics, &metric)
	}
	s.cursor++
	log.Debug("decode s3 object finished", zap.String("key", *obj.Key))

	return metrics, true, nil
}

func (s *S3JsonScanner) ScanRemoteFiles(ctx context.Context) (*RemoteFile, bool, error) {
	return nil, false, errors.New("not implemented")
}

func (s *S3JsonScanner) hasMore() bool {
	return s.cursor < len(s.objects)
}

type S3ParquetScanner struct {
	s3Client *s3.Client
	bucket   string
	objects  []types.Object
	cursor   int
}

func (s *S3ParquetScanner) ScanMetrics(ctx context.Context) ([]*Metric, bool, error) {
	return nil, false, errors.New("not implemented")
}

func (s *S3ParquetScanner) ScanRemoteFiles(ctx context.Context) (*RemoteFile, bool, error) {
	if !s.hasMore() {
		return nil, false, nil
	}
	obj := s.objects[s.cursor]

	log.Debug("get s3 object", zap.String("key", *obj.Key))
	output, err := s.s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    obj.Key,
	})
	if err != nil {
		return nil, false, fmt.Errorf("failed to get metrics file from s3: %s", err)
	}
	log.Debug("get s3 object finished", zap.String("key", *obj.Key))

	name := strings.ReplaceAll(*obj.Key, "/", "-")
	s.cursor++
	return &RemoteFile{Body: output.Body, Name: name}, true, nil
}

func (s *S3ParquetScanner) hasMore() bool {
	return s.cursor < len(s.objects)
}