package api

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"time"

	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/credentials"
)

const (
	AuthService = "jcs"
	AuthRegion  = "any"
)

// 对一个请求进行签名，并将签名信息添加到请求头中。
//
// 会读取请求体计算sha256哈希值。如果hash值已知，可以使用SignWithPayloadHash方法。
func Sign(req *http.Request, accessKey, secretKey string) error {
	prod := credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")
	cred, err := prod.Retrieve(context.TODO())
	if err != nil {
		return err
	}

	payloadHash := ""
	if req.Body != nil {
		data, err := io.ReadAll(req.Body)
		if err != nil {
			return err
		}
		req.Body.Close()
		req.Body = io.NopCloser(bytes.NewReader(data))

		hasher := sha256.New()
		hasher.Write(data)
		payloadHash = hex.EncodeToString(hasher.Sum(nil))
	} else {
		hash := sha256.Sum256([]byte(""))
		payloadHash = hex.EncodeToString(hash[:])
	}

	signer := v4.NewSigner()
	err = signer.SignHTTP(context.Background(), cred, req, payloadHash, AuthService, AuthRegion, time.Now())
	if err != nil {
		return err
	}

	return nil
}

// 对一个请求进行签名，不计算请求体的哈希，适合上传文件接口。
func SignWithoutBody(req *http.Request, accessKey, secretKey string) error {
	prod := credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")
	cred, err := prod.Retrieve(context.TODO())
	if err != nil {
		return err
	}

	signer := v4.NewSigner()
	err = signer.SignHTTP(context.Background(), cred, req, "", AuthService, AuthRegion, time.Now())
	if err != nil {
		return err
	}

	return nil
}

// 对一个请求进行签名，签名时使用指定的哈希值作为请求体的哈希值。
//
// 参数payloadHash必须为sha256哈希值的16进制字符串，全小写。
func SignWithPayloadHash(req *http.Request, payloadHash string, accessKey, secretKey string) error {
	prod := credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")
	cred, err := prod.Retrieve(context.TODO())
	if err != nil {
		return err
	}

	signer := v4.NewSigner()
	err = signer.SignHTTP(context.Background(), cred, req, payloadHash, AuthService, AuthRegion, time.Now())
	if err != nil {
		return err
	}

	return nil
}

// 生成一个带签名的URL。
//
// expiration为签名过期时间，单位为秒。
//
// 签名时不会包含请求体的哈希值。注：不要设置任何额外的Header（除了自动添加的Host），以免签名校验不通过
func Presign(req *http.Request, accessKey, secretKey string, expiration int) (string, error) {
	query := req.URL.Query()
	query.Add("X-Expires", fmt.Sprintf("%v", expiration))

	req.URL.RawQuery = query.Encode()

	prod := credentials.NewStaticCredentialsProvider(accessKey, secretKey, "")
	cred, err := prod.Retrieve(context.TODO())
	if err != nil {
		return "", err
	}

	signer := v4.NewSigner()
	signedURL, _, err := signer.PresignHTTP(context.Background(), cred, req, "", AuthService, AuthRegion, time.Now())
	return signedURL, err
}
