package erasure

import (
	"bytes"
	"context"
	"crypto"
	"crypto/rand"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"sort"
	"strings"
	"time"
)

type AWSSwarmSettings struct {
	S3Bucket, CloudFrontDomain, Region, AccessKeyID, SecretAccessKey, KeyPairID, PrivateKeyPEM string
}

func (s AWSSwarmSettings) Enabled() bool {
	return strings.TrimSpace(s.S3Bucket) != "" && strings.TrimSpace(s.CloudFrontDomain) != ""
}
func (s AWSSwarmSettings) CredentialsReady() bool {
	return s.Enabled() && strings.TrimSpace(s.AccessKeyID) != "" && strings.TrimSpace(s.SecretAccessKey) != ""
}
func (s AWSSwarmSettings) SigningReady() bool {
	return s.Enabled() && strings.TrimSpace(s.KeyPairID) != "" && strings.TrimSpace(s.PrivateKeyPEM) != ""
}

func HydrateAWSSwarmFromEnv(s AWSSwarmSettings) AWSSwarmSettings {
	if v := strings.TrimSpace(os.Getenv("AF_AWS_ACCESS_KEY_ID")); v != "" && s.AccessKeyID == "" {
		s.AccessKeyID = v
	}
	if v := strings.TrimSpace(os.Getenv("AF_AWS_SECRET_ACCESS_KEY")); v != "" && s.SecretAccessKey == "" {
		s.SecretAccessKey = v
	}
	if v := strings.TrimSpace(os.Getenv("AF_AWS_REGION")); v != "" && s.Region == "" {
		s.Region = v
	}
	if v := strings.TrimSpace(os.Getenv("AF_CLOUDFRONT_KEY_PAIR_ID")); v != "" && s.KeyPairID == "" {
		s.KeyPairID = v
	}
	if v := strings.TrimSpace(os.Getenv("AF_CLOUDFRONT_PRIVATE_KEY")); v != "" && s.PrivateKeyPEM == "" {
		s.PrivateKeyPEM = strings.ReplaceAll(v, `\n`, "\n")
	}
	if s.Region == "" {
		s.Region = "us-east-1"
	}
	s.S3Bucket = strings.TrimSpace(s.S3Bucket)
	s.CloudFrontDomain = strings.TrimRight(strings.TrimSpace(s.CloudFrontDomain), "/")
	return s
}

type ShardObjectStore interface {
	PutShard(ctx context.Context, bucket, key string, body []byte) error
	HeadBucket(ctx context.Context, bucket string) error
}

type S3HTTPStore struct {
	Settings AWSSwarmSettings
	Endpoint string
	Client   interface{ Do(*http.Request) (*http.Response, error) }
}

func (s *S3HTTPStore) client() interface{ Do(*http.Request) (*http.Response, error) } {
	if s != nil && s.Client != nil {
		return s.Client
	}
	return http.DefaultClient
}

func (s *S3HTTPStore) endpoint(bucket string) string {
	if s != nil && strings.TrimSpace(s.Endpoint) != "" {
		return strings.TrimRight(strings.TrimSpace(s.Endpoint), "/")
	}
	region := strings.TrimSpace(s.Settings.Region)
	if region == "" {
		region = "us-east-1"
	}
	return fmt.Sprintf("https://%s.s3.%s.amazonaws.com", bucket, region)
}

func (s *S3HTTPStore) HeadBucket(ctx context.Context, bucket string) error {
	u := s.endpoint(bucket) + "/"
	req, _ := http.NewRequestWithContext(ctx, http.MethodHead, u, nil)
	signS3Request(req, s.Settings, bucket, "", nil)
	resp, err := s.client().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 300 {
		return nil
	}
	return fmt.Errorf("s3 head bucket HTTP %d", resp.StatusCode)
}

func (s *S3HTTPStore) PutShard(ctx context.Context, bucket, key string, body []byte) error {
	u := s.endpoint(bucket) + "/" + strings.TrimLeft(key, "/")
	req, _ := http.NewRequestWithContext(ctx, http.MethodPut, u, bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/octet-stream")
	signS3Request(req, s.Settings, bucket, key, body)
	resp, err := s.client().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 300 {
		return nil
	}
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
	return fmt.Errorf("s3 put HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
}

type SwarmAttachResult struct {
	EdgeURLs, ShardManifestURLs []string
	SwarmMagnet                 string
}

func AttachS3Swarm(ctx context.Context, cfg AWSSwarmSettings, store ShardObjectStore, token, payloadSHA string, payloadSize int, p Params, shards [][]byte, hashes []string) (*SwarmAttachResult, error) {
	cfg = HydrateAWSSwarmFromEnv(cfg)
	if !cfg.Enabled() {
		return nil, nil
	}
	if store == nil {
		store = &S3HTTPStore{Settings: cfg}
	}
	p, err := p.Normalize()
	if err != nil {
		return nil, err
	}
	region := strings.TrimSpace(cfg.Region)
	if region == "" {
		region = "us-east-1"
	}
	edgeURLs := make([]string, p.TotalShards())
	for i := range shards {
		key := S3ShardKey(token, region, i, hashes[i])
		if err := store.PutShard(ctx, cfg.S3Bucket, key, shards[i]); err != nil {
			return nil, err
		}
		edge, err := CloudFrontSignedShardURL(cfg, "/"+key, 24*time.Hour)
		if err != nil {
			return nil, err
		}
		edgeURLs[i] = edge
	}
	return &SwarmAttachResult{EdgeURLs: edgeURLs, ShardManifestURLs: append([]string(nil), edgeURLs...), SwarmMagnet: SwarmMagnetLinkEx(token, payloadSHA, edgeURLs)}, nil
}

func CloudFrontSignedShardURL(cfg AWSSwarmSettings, objectPath string, ttl time.Duration) (string, error) {
	cfg = HydrateAWSSwarmFromEnv(cfg)
	if !cfg.SigningReady() {
		return "", fmt.Errorf("cloudfront signing not configured")
	}
	domain := cfg.CloudFrontDomain
	if !strings.HasPrefix(domain, "http") {
		domain = "https://" + domain
	}
	return signCloudFrontURL(strings.TrimRight(domain, "/")+"/"+strings.TrimLeft(objectPath, "/"), cfg.KeyPairID, cfg.PrivateKeyPEM, time.Now().Add(ttl).Unix())
}

func signCloudFrontURL(rawURL, keyPairID, pemKey string, expires int64) (string, error) {
	block, _ := pem.Decode([]byte(pemKey))
	if block == nil {
		return "", fmt.Errorf("invalid PEM")
	}
	keyAny, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		keyAny, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return "", err
		}
	}
	priv, ok := keyAny.(*rsa.PrivateKey)
	if !ok {
		return "", fmt.Errorf("not RSA key")
	}
	canned := fmt.Sprintf("%s?Expires=%d", rawURL, expires)
	hash := sha1.Sum([]byte(canned))
	sig, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA1, hash[:])
	if err != nil {
		return "", err
	}
	u, _ := url.Parse(rawURL)
	q := u.Query()
	q.Set("Expires", fmt.Sprintf("%d", expires))
	q.Set("Signature", base64URLEncode(sig))
	q.Set("Key-Pair-Id", keyPairID)
	u.RawQuery = q.Encode()
	return u.String(), nil
}

func MinimalIAMPolicyJSON(bucket string) string {
	if bucket == "" {
		bucket = "YOUR_SHARD_BUCKET"
	}
	return fmt.Sprintf(`{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Action":["s3:PutObject","s3:GetObject","s3:ListBucket"],"Resource":["arn:aws:s3:::%s","arn:aws:s3:::%s/shards/*"]}]}`, bucket, bucket)
}

func MinimalBucketPolicyJSON(bucket, cfARN string) string {
	if cfARN == "" {
		cfARN = "arn:aws:cloudfront::ACCOUNT_ID:distribution/DISTRIBUTION_ID"
	}
	return fmt.Sprintf(`{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"cloudfront.amazonaws.com"},"Action":"s3:GetObject","Resource":"arn:aws:s3:::%s/shards/*","Condition":{"StringEquals":{"AWS:SourceArn":"%s"}}}]}`, bucket, cfARN)
}

func SwarmMagnetLinkEx(token, payloadSHA string, edgeURLs []string) string {
	base := SwarmMagnetLink(token, payloadSHA)
	if base == "" || len(edgeURLs) == 0 {
		return base
	}
	sort.Strings(edgeURLs)
	for _, u := range edgeURLs {
		if strings.TrimSpace(u) != "" {
			base += "&xs=" + url.QueryEscape(strings.TrimSpace(u))
		}
	}
	return base
}

func signS3Request(req *http.Request, cfg AWSSwarmSettings, _, objectKey string, body []byte) {
	region := cfg.Region
	if region == "" {
		region = "us-east-1"
	}
	now := time.Now().UTC()
	amzDate := now.Format("20060102T150405Z")
	dateStamp := now.Format("20060102")
	payloadHash := hex.EncodeToString(hashPayload(body))
	req.Header.Set("Host", req.URL.Host)
	req.Header.Set("X-Amz-Date", amzDate)
	req.Header.Set("X-Amz-Content-Sha256", payloadHash)
	canonicalURI := "/"
	if objectKey != "" {
		canonicalURI = "/" + escapePath(objectKey)
	}
	canonicalHeaders := "host:" + req.URL.Host + "\n" + "x-amz-content-sha256:" + payloadHash + "\n" + "x-amz-date:" + amzDate + "\n"
	canonicalRequest := strings.Join([]string{req.Method, canonicalURI, "", canonicalHeaders, "host;x-amz-content-sha256;x-amz-date", payloadHash}, "\n")
	scope := dateStamp + "/" + region + "/s3/aws4_request"
	stringToSign := strings.Join([]string{"AWS4-HMAC-SHA256", amzDate, scope, hex.EncodeToString(hashString(canonicalRequest))}, "\n")
	sig := hex.EncodeToString(hmacSHA256(deriveSigningKey(cfg.SecretAccessKey, dateStamp, region, "s3"), stringToSign))
	req.Header.Set("Authorization", fmt.Sprintf("AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=host;x-amz-content-sha256;x-amz-date, Signature=%s", cfg.AccessKeyID, scope, sig))
}

func hashPayload(b []byte) []byte {
	if b == nil {
		return hashString("")
	}
	s := sha256.Sum256(b)
	return s[:]
}
func hashString(s string) []byte { h := sha256.Sum256([]byte(s)); return h[:] }
func hmacSHA256(key []byte, msg string) []byte {
	m := hmac.New(sha256.New, key)
	m.Write([]byte(msg))
	return m.Sum(nil)
}
func deriveSigningKey(secret, date, region, svc string) []byte {
	kDate := hmacSHA256([]byte("AWS4"+secret), date)
	kRegion := hmacSHA256(kDate, region)
	kService := hmacSHA256(kRegion, svc)
	return hmacSHA256(kService, "aws4_request")
}
func escapePath(p string) string {
	parts := strings.Split(p, "/")
	for i, s := range parts {
		parts[i] = url.PathEscape(s)
	}
	return strings.Join(parts, "/")
}
func base64URLEncode(b []byte) string {
	return strings.NewReplacer("+", "-", "=", "_", "/", "~").Replace(base64.StdEncoding.EncodeToString(b))
}