no-copy/src/main/java/ru/soune/nocopy/controller/FileController.java

package ru.soune.nocopy.controller;

import com.google.common.io.ByteStreams;
import com.vrt.NoCopyFileService;
import com.vrt.fileprotection.FileProtector;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.http.*;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import ru.soune.nocopy.dto.file.CheckStatus;
import ru.soune.nocopy.dto.file.FileEntityResponse;
import ru.soune.nocopy.entity.file.FileEntity;
import ru.soune.nocopy.entity.file.ProtectionStatus;
import ru.soune.nocopy.entity.user.ProtectedFileCheck;
import ru.soune.nocopy.entity.user.User;
import ru.soune.nocopy.repository.FileEntityRepository;
import ru.soune.nocopy.repository.UserRepository;
import ru.soune.nocopy.service.file.CheckCounterService;
import ru.soune.nocopy.service.file.ZipService;
import ru.soune.nocopy.service.file.*;
import ru.soune.nocopy.service.file.cloud.CloudStorageService;
import ru.soune.nocopy.util.FileUtil;
import software.amazon.awssdk.services.mq.model.NotFoundException;
import software.amazon.awssdk.services.s3.model.NoSuchKeyException;

import java.io.*;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;

@RestController
@RequestMapping("/api/files")
@Slf4j
@AllArgsConstructor
public class FileController {

    private FileEntityRepository fileRepository;

    private CheckCounterService checkCounterService;

    private UserRepository userRepository;

    private final FileEntityService fileEntityService;

    private NoCopyFileService noCopyFileService;

    private FileUtil fileUtil;

    private ZipService zipService;

    private CloudStorageService cloudStorageService;

    @GetMapping("/public/{fileId}")
    public ResponseEntity<StreamingResponseBody> getPublicFile(@PathVariable String fileId) {
        try {
            FileEntity fileEntity = fileRepository.findById(fileId)
                    .orElseThrow(() -> new FileNotFoundException("Not found in DB: " + fileId));

            MediaType mediaType = getMediaType(fileEntity);

            ContentDisposition contentDisposition = ContentDisposition.inline()
                    .filename(fileEntity.getOriginalFileName(), StandardCharsets.UTF_8)
                    .build();

            StreamingResponseBody responseBody = outputStream -> {
                try (InputStream inputStream = cloudStorageService.readFileFromStorage(fileEntity.getFilePath())) {
                    ByteStreams.copy(inputStream, outputStream);
                } catch (NoSuchKeyException e) {
                    log.error("File disappeared from S3 during streaming: {}", fileId, e);
                    throw new IOException("File not found in storage", e);
                }
            };

            return ResponseEntity.ok()
                    .contentType(mediaType)
                    .header(HttpHeaders.CONTENT_DISPOSITION, contentDisposition.toString())
                    .header(HttpHeaders.CACHE_CONTROL, "public, max-age=3600")
                    .body(responseBody);
        } catch (FileNotFoundException e) {
            log.warn("File not found in DB: {}", fileId);
            return ResponseEntity.notFound().build();
        } catch (NoSuchKeyException e) {
            log.warn("File not found in S3: {}", fileId);
            return ResponseEntity.status(HttpStatus.NOT_FOUND).build();
        } catch (Exception e) {
            log.error("Error reading file: {}", fileId, e);
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).build();
        }
    }

    @GetMapping("/public/file-info/{fileId}")
    public ResponseEntity<FileEntityResponse> getPublicFileInfo(@PathVariable String fileId) {
        return handleGetFileInfo(fileId);
    }

    private ResponseEntity<FileEntityResponse>  handleGetFileInfo(String fileId) {
        try {
            FileEntityResponse fileInfo = fileEntityService.getById(fileId, 1);

            return ResponseEntity
                    .ok()
                    .body(fileInfo);
        } catch (NotFoundException e) {
            return ResponseEntity
                    .notFound().build();
        }
    }

    private MediaType getMediaType(FileEntity fileEntity) {
        try {
            return MediaType.parseMediaType(fileEntity.getMimeType());
        } catch (InvalidMediaTypeException e) {
            String extension = fileEntity.getFileExtension().toLowerCase();
            switch (extension) {
                case "jpg":
                case "jpeg":
                    return MediaType.IMAGE_JPEG;
                case "png":
                    return MediaType.IMAGE_PNG;
                case "pdf":
                    return MediaType.APPLICATION_PDF;
                default:
                    return MediaType.APPLICATION_OCTET_STREAM;
            }
        }
    }

    @GetMapping("/protected/{fileId}/{type}")
    public ResponseEntity<Resource> getProtectedFile(
            @PathVariable String fileId, @PathVariable String type) throws IOException {

        FileEntity fileEntity = fileRepository.findById(fileId)
                .orElseThrow(() -> new RuntimeException("File not found"));

        String filePath = type.equals("thumbnail") ? fileEntity.getThumbnailPath(): fileEntity.getFilePath();

        if (fileEntity.getProtectionStatus() == ProtectionStatus.NOT_PROTECTED ||
                fileEntity.getProtectionStatus() == ProtectionStatus.PROCESSING) {
            throw new RuntimeException("File is not protected");
        }

        String fileName = fileEntity.getOriginalFileName().replace("." +
                fileEntity.getFileExtension(), "") + "_nocopy_protected" +
                "." + fileEntity.getFileExtension();
        String encodedFileName = URLEncoder.encode(fileName, StandardCharsets.UTF_8)
                .replace("+", "%20");

        InputStreamResource resource = new InputStreamResource(cloudStorageService.readFileFromStorage(filePath));

        return ResponseEntity.ok()
                .contentType(new MediaType(fileEntity.getMimeType(), fileEntity.getFileExtension()))
                .header(HttpHeaders.CONTENT_DISPOSITION,
                        "inline; filename*=UTF-8''" + encodedFileName)
                .body(resource);
    }

    @GetMapping("check-file/status/{userId}")
    public ResponseEntity<CheckStatus> getStatus(@PathVariable Long userId) {
        CheckStatus status = checkCounterService.getCurrentStatus(userId);

        return ResponseEntity.ok(status);
    }

    @PostMapping("check-file/update-limit/{userId}/{limit}")
    public ResponseEntity<ProtectedFileCheck> getStatus(@PathVariable Long userId, @PathVariable Integer limit) {
        User user = userRepository.findById(userId).orElseThrow();
        ProtectedFileCheck protectedFileCheck = checkCounterService.updateLimit(user, limit);

        return ResponseEntity
                .ok()
                .body(protectedFileCheck);
    }

    @PostMapping("protect/{fileId}/{convertTo}")
    public ResponseEntity<FileProtector.FileInfo> convert(@PathVariable String fileId, @PathVariable String convertTo) {
        FileEntity file = fileRepository.findByFileId(fileId);
        FileProtector.FileInfo fileInfo = fileUtil.createFileInfo(file, convertTo);

        noCopyFileService.addFile(fileInfo);

        return ResponseEntity
                .ok()
                .body(fileInfo);
    }

    @GetMapping("/download/user-passport/archive/{userId}")
    public ResponseEntity<byte[]> downloadPassportZip(@PathVariable Long userId) {
        try {
            byte[] zipBytes = zipService.assembleZipFromSplitFiles(userId);

            return ResponseEntity.ok()
                    .header(HttpHeaders.CONTENT_DISPOSITION,
                            "attachment; filename=passport_" + userId + ".zip")
                    .contentType(MediaType.APPLICATION_OCTET_STREAM)
                    .contentLength(zipBytes.length)
                    .body(zipBytes);
        } catch (IOException e) {
            log.error("Failed to assemble passport for user {}", userId, e);
            return ResponseEntity.internalServerError().build();
        }
    }
}