package pimport

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

	"github.com/aws/aws-sdk-go/aws"
	awss3 "github.com/aws/aws-sdk-go/service/s3"
	"github.com/docker/go-units"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/dataflow-service/internal/db"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/pkg/account/accountdb"
	"github.com/tidbcloud/dataflow-service/pkg/s3"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

const (
	// MaxUploadSize defines the maximum file size to be uploaded.
	MaxUploadSize = 50 * units.MiB
	// DefaultRegion is the default region to upload the file into.
	DefaultRegion = "us-east-1"
)

var (
	// ErrLocalFileTooLarge represents errors that the source local file is too large
	ErrLocalFileTooLarge = errors.New("The source local file is too large.")
	// ErrTargetFileNotReady represents errors that the target file for uploading is not ready.
	ErrTargetFileNotReady = errors.New("The target file for uploading is not ready.")
)

// StubManager represents the model for managing upload stubs.
type StubManager interface {
	// CreateStub creates an import upload stub, and gets back the stub ID.
	CreateStub(ctx context.Context, stub *entity.ImportUploadStub) (entity.ID, error)
	// GetStubByID gets the import upload stub information by an upload stub ID.
	GetStubByID(ctx context.Context, stubID entity.ID) (*entity.ImportUploadStub, error)
}

// StubManagerImpl implements the stub manager.
type StubManagerImpl struct{}

// CreateStub creates an import upload stub, and gets back the stub ID.
// It implements the `StubManager` interface.
func (m *StubManagerImpl) CreateStub(ctx context.Context, stub *entity.ImportUploadStub) (entity.ID, error) {
	if err := db.Transaction(func(tx *gorm.DB) error {
		return errors.Annotate(stub.Create(tx), "create import upload stub error")
	}); err != nil {
		return entity.ID(0), err
	}
	return stub.ID, nil
}

// GetStubByID gets the import upload stub information by an upload stub ID.
// It implements the `StubManager` interface.
func (m *StubManagerImpl) GetStubByID(ctx context.Context, stubID entity.ID) (*entity.ImportUploadStub, error) {
	var stub entity.ImportUploadStub
	if err := db.Transaction(func(tx *gorm.DB) error {
		return errors.Annotate(entity.NewImportUploadStubQuerySet(tx).
			IDEq(stubID).
			One(&stub), "get upload stub by ID error")
	}); err != nil {
		return nil, err
	}
	return &stub, nil
}

// PreparedUploadInfo represents the information for an prepared upload file.
// This usually means the remote file has been registered, but the data has not been transferred into the file.
type PreparedUploadInfo struct {
	// OriginalFileName is the original file name at the client side.
	OriginalFileName string
	// TargetFileName is the file name at the remote side.
	TargetFileName string
	// TargetURL is the URL that can be used to retrieve the target.
	TargetURL string
}

// LocalFileUploader represents the operations that are related to uploading a local file.
type LocalFileUploader interface {
	// PrepareTargetFile makes a local file ready to be transferred to the remote.
	// This usually means the remote file has been registered, but the data has not been transferred into the file.
	PrepareTargetFile(ctx context.Context, localFileName string, totalSize uint64) (string, error)
	// UploadData does the local file data transfer.
	UploadData(ctx context.Context, reader io.Reader) error
	// GetPreparedUploadInfo gets the prepared file's upload info.
	GetPreparedUploadInfo(ctx context.Context) (*PreparedUploadInfo, error)
}

// S3TargetObjectInfo represents the information for uploading a local file into the S3 bucket.
type S3TargetObjectInfo struct {
	UploadURL        string
	TotalSize        int64
	TargetFileName   string
	TargetURL        string
	OriginalFileName string
}

// LocalFileS3Uploader is the S3 implementation of the LocalFileUploader.
// It will upload the local file into an S3 bucket.
type LocalFileS3Uploader struct {
	user       *accountdb.User
	region     string
	targetInfo *S3TargetObjectInfo
}

// NewLocalFileS3Uploader creates an `LocalFileS3Uploader`.
func NewLocalFileS3Uploader(user *accountdb.User, region string, targetInfo *S3TargetObjectInfo) *LocalFileS3Uploader {
	return &LocalFileS3Uploader{
		user:       user,
		region:     region,
		targetInfo: targetInfo,
	}
}

// PrepareTargetFile makes a local file ready to be transferred to the remote.
// This usually means the remote file has been registered, but the data has not been transferred into the file.
// It implements the `LocalFileS3Uploader` interface.
// The S3 uploader will create a remote object, and get the the pre-signed upload URL to upload the data.
func (u *LocalFileS3Uploader) PrepareTargetFile(ctx context.Context, localFileName string, totalSize uint64) (string, error) {
	var err error
	logger := log.WithContextL(ctx)
	if totalSize > MaxUploadSize {
		logger.Warn("file size is bigger than 50MiB", zap.Uint64("content_length(MiB)", totalSize/units.MiB))
		return "", ErrLocalFileTooLarge
	}

	// todo: check region valid
	svc, err := s3.CreateClientWithAssumeRole(u.region, "", "")
	if err != nil {
		return "", errors.Annotate(err, "failed to create client with assume role")
	}
	now := time.Now()
	fileName := fmt.Sprintf("%s-%s", now.Format(time.RFC3339), filepath.Base(localFileName))
	bucketNaming := NewLocalImportBucketNaming(u.region, fileName, u.user.ID)
	objReq, _ := svc.PutObjectRequest(&awss3.PutObjectInput{
		Bucket:        aws.String(bucketNaming.GenBucketName()),
		Key:           aws.String(bucketNaming.GenObjectKey()),
		ContentLength: aws.Int64(int64(totalSize)),
	})
	// the front end will use the url immediately after fetching it.
	uploadURL, err := objReq.Presign(3 * time.Minute)
	if err != nil {
		return "", errors.Annotate(err, "failed to put object")
	}
	u.targetInfo = &S3TargetObjectInfo{
		UploadURL:        uploadURL,
		TotalSize:        int64(totalSize),
		OriginalFileName: localFileName,
		TargetFileName:   fileName,
		TargetURL:        bucketNaming.GenBucketUrl(),
	}
	return uploadURL, nil
}

// UploadData does the local file data transfer.
// It implements the `LocalFileS3Uploader` interface.
func (u *LocalFileS3Uploader) UploadData(ctx context.Context, r io.Reader) error {
	logger := log.WithContextL(ctx)
	if u.targetInfo == nil {
		return ErrTargetFileNotReady
	}
	logger.Info("upload URL", zap.String("upload_url", u.targetInfo.UploadURL))
	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPut,
		u.targetInfo.UploadURL,
		io.LimitReader(r, u.targetInfo.TotalSize),
	)
	if err != nil {
		return errors.Annotate(err, "construct HTTP request error")
	}
	req.ContentLength = u.targetInfo.TotalSize
	client := http.DefaultClient
	resp, err := client.Do(req)
	if err != nil {
		return errors.Annotate(err, "send request to the target S3 upload URL failed")
	}
	//nolint:errcheck
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		body := ""
		var b strings.Builder
		if _, err := io.Copy(&b, resp.Body); err == nil {
			body = b.String()
		}
		return fmt.Errorf("PUT data to target has error: unexpected status code: %d.  Body: %s", resp.StatusCode, body)
	}
	return nil
}

// GetPreparedUploadInfo gets the prepared file's upload info.
// It implements the `LocalFileS3Uploader` interface.
func (u *LocalFileS3Uploader) GetPreparedUploadInfo(ctx context.Context) (*PreparedUploadInfo, error) {
	if u.targetInfo == nil {
		return nil, ErrTargetFileNotReady
	}

	return &PreparedUploadInfo{
		TargetFileName:   u.targetInfo.TargetFileName,
		TargetURL:        u.targetInfo.TargetURL,
		OriginalFileName: u.targetInfo.OriginalFileName,
	}, nil
}