package storage

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"path"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	awsv2config "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/smithy-go"
)

type s3ObjectAPI interface {
	PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
	GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
	HeadObject(ctx context.Context, params *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
	DeleteObject(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error)
}

type s3PresignAPI interface {
	PresignGetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.PresignOptions)) (*v4.PresignedHTTPRequest, error)
}

type S3Store struct {
	client    s3ObjectAPI
	presigner s3PresignAPI
	bucket    string
	keyPrefix string
}

func NewS3Store(cfg Config) (*S3Store, error) {
	bucket := strings.TrimSpace(cfg.Bucket)
	if bucket == "" {
		return nil, fmt.Errorf("s3 storage bucket cannot be empty")
	}

	region := strings.TrimSpace(cfg.Region)
	if region == "" {
		region = "us-east-1"
	}

	loadOptions := []func(*awsv2config.LoadOptions) error{
		awsv2config.WithRegion(region),
	}
	if strings.TrimSpace(cfg.AccessKeyID) != "" || strings.TrimSpace(cfg.SecretAccessKey) != "" || strings.TrimSpace(cfg.SessionToken) != "" {
		loadOptions = append(loadOptions, awsv2config.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(cfg.AccessKeyID, cfg.SecretAccessKey, cfg.SessionToken),
		))
	}

	awsCfg, err := awsv2config.LoadDefaultConfig(context.Background(), loadOptions...)
	if err != nil {
		return nil, fmt.Errorf("load s3 storage config: %w", err)
	}

	client := s3.NewFromConfig(awsCfg, func(options *s3.Options) {
		options.UsePathStyle = cfg.UsePathStyle
		if endpoint := strings.TrimSpace(cfg.Endpoint); endpoint != "" {
			options.BaseEndpoint = aws.String(endpoint)
		}
	})

	return &S3Store{
		client:    client,
		presigner: s3.NewPresignClient(client),
		bucket:    bucket,
		keyPrefix: strings.Trim(strings.ReplaceAll(cfg.KeyPrefix, "\\", "/"), "/"),
	}, nil
}

func (s *S3Store) PutObject(ctx context.Context, key string, body io.Reader, opts PutObjectOptions) (ObjectMetadata, error) {
	resolvedKey, err := s.resolveKey(key)
	if err != nil {
		return ObjectMetadata{}, err
	}

	// Buffer non-seekable readers so the AWS SDK can retry safely.
	if opts.ContentLength == nil || !isSeekableReader(body) {
		payload, err := io.ReadAll(body)
		if err != nil {
			return ObjectMetadata{}, fmt.Errorf("read object body: %w", err)
		}
		size := int64(len(payload))
		opts.ContentLength = &size
		body = bytes.NewReader(payload)
	}

	_, err = s.client.PutObject(ctx, &s3.PutObjectInput{
		Bucket:        aws.String(s.bucket),
		Key:           aws.String(resolvedKey),
		Body:          body,
		ContentType:   aws.String(strings.TrimSpace(opts.ContentType)),
		ContentLength: opts.ContentLength,
	})
	if err != nil {
		return ObjectMetadata{}, fmt.Errorf("put object: %w", err)
	}

	head, err := s.HeadObject(ctx, key)
	if err != nil {
		return ObjectMetadata{}, err
	}
	return head, nil
}

func isSeekableReader(body io.Reader) bool {
	_, ok := body.(io.Seeker)
	return ok
}

func (s *S3Store) GetObject(ctx context.Context, key string) (io.ReadCloser, ObjectMetadata, error) {
	resolvedKey, err := s.resolveKey(key)
	if err != nil {
		return nil, ObjectMetadata{}, err
	}

	output, err := s.client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(resolvedKey),
	})
	if err != nil {
		if isS3NotFound(err) {
			return nil, ObjectMetadata{}, ErrObjectNotFound
		}
		return nil, ObjectMetadata{}, fmt.Errorf("get object: %w", err)
	}

	contentType := aws.ToString(output.ContentType)
	if strings.TrimSpace(contentType) == "" {
		contentType = "application/octet-stream"
	}

	return output.Body, ObjectMetadata{
		Key:          key,
		Size:         aws.ToInt64(output.ContentLength),
		ContentType:  contentType,
		LastModified: aws.ToTime(output.LastModified),
	}, nil
}

func (s *S3Store) HeadObject(ctx context.Context, key string) (ObjectMetadata, error) {
	resolvedKey, err := s.resolveKey(key)
	if err != nil {
		return ObjectMetadata{}, err
	}

	output, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(resolvedKey),
	})
	if err != nil {
		if isS3NotFound(err) {
			return ObjectMetadata{}, ErrObjectNotFound
		}
		return ObjectMetadata{}, fmt.Errorf("head object: %w", err)
	}

	contentType := aws.ToString(output.ContentType)
	if strings.TrimSpace(contentType) == "" {
		contentType = "application/octet-stream"
	}

	return ObjectMetadata{
		Key:          key,
		Size:         aws.ToInt64(output.ContentLength),
		ContentType:  contentType,
		LastModified: aws.ToTime(output.LastModified),
	}, nil
}

func (s *S3Store) DeleteObject(ctx context.Context, key string) error {
	resolvedKey, err := s.resolveKey(key)
	if err != nil {
		return err
	}

	_, err = s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(resolvedKey),
	})
	if err != nil {
		if isS3NotFound(err) {
			return ErrObjectNotFound
		}
		return fmt.Errorf("delete object: %w", err)
	}
	return nil
}

func (s *S3Store) GeneratePresignedGetURL(ctx context.Context, key string, expiry time.Duration) (string, error) {
	resolvedKey, err := s.resolveKey(key)
	if err != nil {
		return "", err
	}

	request, err := s.presigner.PresignGetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(resolvedKey),
	}, func(options *s3.PresignOptions) {
		options.Expires = expiry
	})
	if err != nil {
		if isS3NotFound(err) {
			return "", ErrObjectNotFound
		}
		return "", fmt.Errorf("presign get object: %w", err)
	}

	return request.URL, nil
}

func (s *S3Store) resolveKey(key string) (string, error) {
	logicalKey, err := normalizeStorageKey(key)
	if err != nil {
		return "", err
	}
	if s.keyPrefix == "" {
		return logicalKey, nil
	}
	return path.Join(s.keyPrefix, logicalKey), nil
}

func normalizeStorageKey(key string) (string, error) {
	cleanKey := strings.TrimSpace(strings.ReplaceAll(key, "\\", "/"))
	cleanKey = strings.TrimPrefix(cleanKey, "/")
	if cleanKey == "" {
		return "", fmt.Errorf("storage key cannot be empty")
	}

	for _, segment := range strings.Split(cleanKey, "/") {
		if segment == ".." {
			return "", fmt.Errorf("invalid storage key %q", key)
		}
	}

	cleanKey = path.Clean(cleanKey)
	if cleanKey == "." || cleanKey == "" || cleanKey == ".." || strings.HasPrefix(cleanKey, "../") {
		return "", fmt.Errorf("invalid storage key %q", key)
	}

	return cleanKey, nil
}

func isS3NotFound(err error) bool {
	var apiErr smithy.APIError
	if !errors.As(err, &apiErr) {
		return false
	}

	switch apiErr.ErrorCode() {
	case "NoSuchKey", "NotFound", "NoSuchBucket":
		return true
	default:
		return false
	}
}