package archive

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/aws/smithy-go"
)

type s3Client interface {
	GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
	PutObject(ctx context.Context, params *s3.PutObjectInput, optFns ...func(*s3.Options)) (*s3.PutObjectOutput, error)
}

type s3Storage struct {
	bucket string
	prefix string
	client s3Client
}

type S3Option func(*s3Options)

type s3Options struct {
	httpClient *http.Client
}

func WithS3HTTPClient(httpClient *http.Client) S3Option {
	return func(options *s3Options) {
		options.httpClient = httpClient
	}
}

// NewS3 constructs an S3-backed archive storage.
func NewS3(ctx context.Context, region string, bucket string, prefix string, opts ...S3Option) (Storage, error) {
	trimmedRegion := strings.TrimSpace(region)
	if trimmedRegion == "" {
		return nil, fmt.Errorf("archive S3 region is required")
	}
	trimmedBucket := strings.TrimSpace(bucket)
	if trimmedBucket == "" {
		return nil, fmt.Errorf("archive S3 bucket is required")
	}
	options := s3Options{}
	for _, opt := range opts {
		if opt != nil {
			opt(&options)
		}
	}
	loadOptions := []func(*awsconfig.LoadOptions) error{awsconfig.WithRegion(trimmedRegion)}
	if options.httpClient != nil {
		loadOptions = append(loadOptions, awsconfig.WithHTTPClient(options.httpClient))
	}
	cfg, err := awsconfig.LoadDefaultConfig(ctx, loadOptions...)
	if err != nil {
		return nil, fmt.Errorf("load AWS config: %w", err)
	}
	return newS3Storage(trimmedBucket, prefix, s3.NewFromConfig(cfg)), nil
}

func newS3Storage(bucket string, prefix string, client s3Client) Storage {
	return &s3Storage{
		bucket: strings.TrimSpace(bucket),
		prefix: strings.Trim(strings.TrimSpace(prefix), "/"),
		client: client,
	}
}

func (storage *s3Storage) ReadChunk(ctx context.Context, workspaceID string, threadID string, chunkStartSeq int64) ([]Message, error) {
	key, err := ObjectKey(storage.prefix, workspaceID, threadID, chunkStartSeq)
	if err != nil {
		return nil, err
	}
	output, err := storage.client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: &storage.bucket,
		Key:    &key,
	})
	if err != nil {
		return nil, mapS3ReadError(err)
	}
	defer output.Body.Close()
	messages, err := decodeChunk(output.Body)
	if err != nil {
		return nil, fmt.Errorf("decode archive chunk %s: %w", key, err)
	}
	if err := validateChunkMessages(chunkStartSeq, messages); err != nil {
		return nil, fmt.Errorf("validate archive chunk %s: %w", key, err)
	}
	return messages, nil
}

func (storage *s3Storage) WriteChunk(ctx context.Context, workspaceID string, threadID string, chunkStartSeq int64, messages []Message) error {
	if err := validateChunkMessages(chunkStartSeq, messages); err != nil {
		return err
	}
	key, err := ObjectKey(storage.prefix, workspaceID, threadID, chunkStartSeq)
	if err != nil {
		return err
	}
	payload, err := encodeChunk(messages)
	if err != nil {
		return err
	}
	contentType := "application/x-ndjson+zstd"
	ifNoneMatch := "*"
	_, err = storage.client.PutObject(ctx, &s3.PutObjectInput{
		Bucket:            &storage.bucket,
		Key:               &key,
		Body:              bytes.NewReader(payload),
		ChecksumAlgorithm: s3types.ChecksumAlgorithmSha256,
		ContentType:       &contentType,
		IfNoneMatch:       &ifNoneMatch,
	})
	if err != nil {
		var apiError smithy.APIError
		if errors.As(err, &apiError) && apiError.ErrorCode() == "PreconditionFailed" {
			existing, readErr := storage.ReadChunk(ctx, workspaceID, threadID, chunkStartSeq)
			if readErr != nil {
				return fmt.Errorf("read existing archive chunk %s: %w", key, readErr)
			}
			if chunkMessagesEqual(existing, messages) {
				return nil
			}
			return fmt.Errorf("%w: %s", ErrChunkAlreadyExists, key)
		}
		return fmt.Errorf("put archive chunk %s: %w", key, err)
	}
	return nil
}

func mapS3ReadError(err error) error {
	var noSuchKey *s3types.NoSuchKey
	if errors.As(err, &noSuchKey) {
		return ErrChunkNotFound
	}
	var apiError smithy.APIError
	if errors.As(err, &apiError) && apiError.ErrorCode() == "NoSuchKey" {
		return ErrChunkNotFound
	}
	return fmt.Errorf("get archive chunk from S3: %w", err)
}