ai-teacher/backend/src/main/java/com/aiteacher/figure/S3FigureStorageService.java

package com.aiteacher.figure;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.S3Configuration;
import software.amazon.awssdk.services.s3.model.*;
import software.amazon.awssdk.services.s3.presigner.S3Presigner;
import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest;
import software.amazon.awssdk.services.s3.model.S3Object;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URI;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

@Service
public class S3FigureStorageService implements FigureStorageService {

    private static final Logger log = LoggerFactory.getLogger(S3FigureStorageService.class);

    private final S3Client s3;
    private final S3Presigner presigner;
    private final String bucket;

    public S3FigureStorageService(
            @Value("${app.figure-storage.endpoint}") String endpoint,
            @Value("${app.figure-storage.region}") String region,
            @Value("${app.figure-storage.bucket}") String bucket,
            @Value("${app.figure-storage.access-key-id}") String accessKeyId,
            @Value("${app.figure-storage.secret-access-key}") String secretKey) {
        this.bucket = bucket;
        URI endpointUri = URI.create(endpoint);
        StaticCredentialsProvider credentials = StaticCredentialsProvider.create(
                AwsBasicCredentials.create(accessKeyId, secretKey));
        Region awsRegion = Region.of(region);

        S3Configuration s3Config = S3Configuration.builder()
                .pathStyleAccessEnabled(true)
                .build();

        this.s3 = S3Client.builder()
                .endpointOverride(endpointUri)
                .region(awsRegion)
                .credentialsProvider(credentials)
                .serviceConfiguration(s3Config)
                .build();

        this.presigner = S3Presigner.builder()
                .endpointOverride(endpointUri)
                .region(awsRegion)
                .credentialsProvider(credentials)
                .serviceConfiguration(s3Config)
                .build();
    }

    @Override
    public String save(UUID bookId, String figureId, BufferedImage image) {
        String key = "figures/" + bookId + "/" + figureId + ".png";
        try {
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            ImageIO.write(image, "PNG", out);
            byte[] bytes = out.toByteArray();

            s3.putObject(
                    PutObjectRequest.builder().bucket(bucket).key(key)
                            .contentType("image/png").contentLength((long) bytes.length).build(),
                    RequestBody.fromBytes(bytes));
            return key;
        } catch (IOException ex) {
            throw new RuntimeException("Failed to encode figure " + figureId, ex);
        } catch (S3Exception ex) {
            throw new RuntimeException("Failed to upload figure " + figureId + " to S3", ex);
        }
    }

    @Override
    public byte[] getBytes(String key) {
        try {
            return s3.getObjectAsBytes(
                    GetObjectRequest.builder().bucket(bucket).key(key).build()).asByteArray();
        } catch (S3Exception ex) {
            throw new RuntimeException("Failed to download figure from S3: " + key, ex);
        }
    }

    @Override
    public String presignedUrl(String key) {
        GetObjectPresignRequest request = GetObjectPresignRequest.builder()
                .signatureDuration(Duration.ofHours(1))
                .getObjectRequest(r -> r.bucket(bucket).key(key))
                .build();
        return presigner.presignGetObject(request).url().toString();
    }

    @Override
    public void deleteAll(UUID bookId) {
        String prefix = "figures/" + bookId + "/";
        try {
            List<ObjectIdentifier> toDelete = new ArrayList<>();
            ListObjectsV2Request listRequest = ListObjectsV2Request.builder()
                    .bucket(bucket).prefix(prefix).build();

            s3.listObjectsV2Paginator(listRequest).stream()
                    .flatMap(page -> page.contents().stream())
                    .map(S3Object::key)
                    .map(k -> ObjectIdentifier.builder().key(k).build())
                    .forEach(toDelete::add);

            if (toDelete.isEmpty()) return;

            s3.deleteObjects(DeleteObjectsRequest.builder()
                    .bucket(bucket)
                    .delete(Delete.builder().objects(toDelete).build())
                    .build());
            log.info("Deleted {} figures from S3 for book {}", toDelete.size(), bookId);
        } catch (S3Exception ex) {
            log.warn("Could not fully delete figures for book {} from S3: {}", bookId, ex.getMessage());
        }
    }
}