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

package ru.soune.nocopy.controller;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.Resource;
import org.springframework.core.io.UrlResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.validation.BindingResult;
import org.springframework.validation.FieldError;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import ru.soune.nocopy.dto.BaseRequest;
import ru.soune.nocopy.dto.BaseResponse;
import ru.soune.nocopy.dto.MessageCode;
import ru.soune.nocopy.dto.RegAnswer;
import ru.soune.nocopy.dto.file.ChunkUploadResponse;
import ru.soune.nocopy.dto.file.CompleteUploadResponse;
import ru.soune.nocopy.dto.file.FileEntityResponse;
import ru.soune.nocopy.dto.file.UploadProgress;
import ru.soune.nocopy.entity.file.FileStatus;
import ru.soune.nocopy.entity.file.UploadStatus;
import ru.soune.nocopy.exception.*;
import ru.soune.nocopy.handler.*;
import ru.soune.nocopy.service.auth.AuthService;
import ru.soune.nocopy.service.file.FileEntityService;
import ru.soune.nocopy.service.file.FileUploadService;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("/api")
@RequiredArgsConstructor
public class ApiController {
    private final FileUploadService fileUploadService;

    private final Map<Integer, RequestHandler> handlers;

    private final FileEntityService fileEntityService;

    private final AuthService authService;

    @PostMapping("/v{version}/data")
    public ResponseEntity<?> handlePostRequest(@RequestBody BaseRequest request,
                                               @PathVariable("version") int version) {
        Integer msgId = request.getMsgId();
        BaseResponse response;

        try {
            RequestHandler handler = handlers.get(msgId);

            if (handler == null) {
                response = new BaseResponse(msgId,
                        MessageCode.MSG_ID_NOT_FOUND.getCode(),
                        MessageCode.MSG_ID_NOT_FOUND.getDescription(),
                        new HashMap<>());
            } else {
                response = handler.handle(request);
            }

            return ResponseEntity.ok().body(response);
        } catch (ValidationException e) {
            return createValidationErrorResponse(e.getBindingResult(), e.getMsgId());
        } catch (NotFoundAuthToken e) {
            return ResponseEntity.ok().body(new BaseResponse(msgId, MessageCode.AUTH_TOKEN_NOT_FOUND.getCode(),
                    MessageCode.AUTH_TOKEN_NOT_FOUND.getDescription(), new HashMap<>()));
        } catch (NotValidFieldException e) {
            throw e;
        } catch (Exception e) {
            log.error("Handler execution failed for msgId: {}", msgId, e);

            BaseResponse errorResponse = new BaseResponse(msgId,
                    MessageCode.INVALID_JSON_BODY.getCode(),
                    MessageCode.INVALID_JSON_BODY.getDescription(),
                    new HashMap<>());

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

    @PostMapping("/v{version}/files/chunk")
    public ResponseEntity<BaseResponse> uploadChunk(
            @PathVariable("version") int version,
            @RequestParam(value = "upload_id", required = false) String uploadId,
            @RequestParam(value = "chunk_number", required = false) Integer chunkNumber,
            @RequestParam(value = "chunk", required = false) MultipartFile chunk) {
        try {
            if (chunk == null || chunk.isEmpty()) {
                return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.FILE_UPLOAD_ERROR.getCode(),
                        "Chunk file null or empty", ChunkUploadResponse.builder()
                        .uploadId(uploadId)
                        .chunkNumber(chunkNumber)
                        .build()));
            }

            if (uploadId == null || uploadId.isBlank()) {
                return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.FILE_UPLOAD_ERROR.getCode(),
                        "Upload ID is required", ChunkUploadResponse.builder()
                        .uploadId(uploadId)
                        .chunkNumber(chunkNumber)
                        .build()));
            }

            if (chunkNumber == null || chunkNumber < 0) {
                return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.FILE_UPLOAD_ERROR.getCode(),
                        "Valid chunk number is required", ChunkUploadResponse.builder()
                        .uploadId(uploadId)
                        .chunkNumber(chunkNumber)
                        .build()));
            }

            fileUploadService.uploadChunk(uploadId, chunkNumber, chunk);

            ChunkUploadResponse responseBody = ChunkUploadResponse.builder()
                    .uploadId(uploadId)
                    .chunkNumber(chunkNumber)
                    .chunkSize(chunk.getSize())
                    .message("Chunk uploaded successfully")
                    .build();

            return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.SUCCESS.getCode(),
                    "Chunk uploaded successfully", responseBody));
        } catch (Exception e) {
            log.error("Error uploading chunk", e);

            ChunkUploadResponse responseBody = ChunkUploadResponse.builder()
                    .uploadId(uploadId)
                    .chunkNumber(chunkNumber)
                    .build();

            return ResponseEntity.ok().body(new BaseResponse(
                    20004, MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to upload chunk: " + e.getMessage(), responseBody));
        }
    }

    @GetMapping("/v{version}/files/progress/{uploadId}")
    public ResponseEntity<BaseResponse> getUploadProgress(
            @PathVariable("version") int version,
            @PathVariable String uploadId) {

        log.info("Getting progress for upload session: {}, version: {}", uploadId, version);

        try {
            var progress = fileUploadService.getUploadProgress(uploadId);

            UploadProgress responseBody = UploadProgress.builder()
                    .uploadId(progress.getUploadId())
                    .fileName(progress.getFileName())
                    .totalChunks(progress.getTotalChunks())
                    .uploadedChunks(progress.getUploadedChunks())
                    .status(progress.getStatus().toString())
                    .progressPercentage(progress.getProgressPercentage())
                    .filePath(progress.getFilePath())
                    .remainingChunks(progress.getTotalChunks() - progress.getUploadedChunks())
                    .build();

            return ResponseEntity.ok().body(new BaseResponse(
                    20004, MessageCode.SUCCESS.getCode(), MessageCode.SUCCESS.getDescription(), responseBody));

        } catch (Exception e) {
            log.error("Error getting progress for upload: {}", uploadId, e);

            UploadProgress responseBody = UploadProgress.builder()
                    .uploadId(uploadId)
                    .build();

            return ResponseEntity.ok().body(new BaseResponse(
                    20004, MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to get upload progress: " + e.getMessage(), responseBody));
        }
    }

    @PostMapping("/v{version}/files/{uploadId}/complete")
    public ResponseEntity<BaseResponse> completeUpload(
            @PathVariable("version") int version,
            @PathVariable String uploadId) {

        log.info("Completing upload session: {}, version: {}", uploadId, version);

        try {
            var progress = fileUploadService.getUploadProgress(uploadId);

            if (progress.getStatus() == UploadStatus.COMPLETED) {
                CompleteUploadResponse responseBody = CompleteUploadResponse.builder()
                        .uploadId(progress.getUploadId())
                        .status(progress.getStatus().toString())
                        .uploadedChunks(progress.getUploadedChunks())
                        .totalChunks(progress.getTotalChunks())
                        .message("Upload already completed")
                        .filePath(progress.getFilePath())
                        .build();

                return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.SUCCESS.getCode(),
                        "Upload already completed", responseBody));
            }

            if (progress.getUploadedChunks() < progress.getTotalChunks()) {
                CompleteUploadResponse responseBody = CompleteUploadResponse.builder()
                        .uploadId(progress.getUploadId())
                        .status(progress.getStatus().toString())
                        .uploadedChunks(progress.getUploadedChunks())
                        .totalChunks(progress.getTotalChunks())
                        .message("Not all chunks uploaded")
                        .build();

                return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.INCOMPLETE_UPLOAD.getCode(),
                        "Not all chunks uploaded", responseBody));
            }

            CompleteUploadResponse responseBody = CompleteUploadResponse.builder()
                    .uploadId(progress.getUploadId())
                    .status(progress.getStatus().toString())
                    .uploadedChunks(progress.getUploadedChunks())
                    .totalChunks(progress.getTotalChunks())
                    .message("File assembly in progress")
                    .build();

            return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.SUCCESS.getCode(),
                    "File assembly in progress", responseBody));

        } catch (Exception e) {
            log.error("Error completing upload: {}", uploadId, e);

            CompleteUploadResponse responseBody = CompleteUploadResponse.builder()
                    .uploadId(uploadId)
                    .build();

            return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to complete upload: " + e.getMessage(), responseBody));
        }
    }


    @GetMapping("/v{version}/files/download/{fileId}")
    public ResponseEntity<?> downloadFile(
            @PathVariable(required = false) String fileId,
            @PathVariable(required = false) Integer version,
            @RequestHeader(value = "Authorization", required = false) String tokenHeader) {
        try {
            if (tokenHeader == null) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("token", tokenHeader);
                errorData.put("fileId", fileId);
                errorData.put("version", version);

                return ResponseEntity.ok().body(new BaseResponse(20004,
                        MessageCode.FILE_DOWNLOAD_ERROR_NOT_CORRECT_FIELD.getCode(),
                        MessageCode.FILE_DOWNLOAD_ERROR_NOT_CORRECT_FIELD.getDescription(),
                        errorData));
            }

            Long userId = authService.useUserAuthToken(tokenHeader);
            FileEntityResponse entityResponse = fileEntityService.getById(fileId, version);

            if (!entityResponse.getUserId().equals(userId)) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("token", tokenHeader);

                return ResponseEntity.ok().body(new BaseResponse(20004,
                        MessageCode.FILE_DOWNLOAD_ERROR.getCode(),
                        MessageCode.FILE_DOWNLOAD_ERROR.getDescription(),
                        errorData));
            }

            if (entityResponse.getStatus().equals(FileStatus.DELETED)) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("file_status", entityResponse.getStatus());

                return ResponseEntity.ok().body(new BaseResponse(20004,
                        MessageCode.FILE_DOWNLOAD_ERROR.getCode(),
                        MessageCode.FILE_DOWNLOAD_ERROR.getDescription(),
                        errorData));
            }

            if (!entityResponse.isExistsOnDisk()) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("onDisk", entityResponse.isExistsOnDisk());

                return ResponseEntity.ok().body(new BaseResponse(20004,
                        MessageCode.FILE_DOWNLOAD_ERROR.getCode(),
                        MessageCode.FILE_DOWNLOAD_ERROR.getDescription(),
                        errorData));
            }

            Path filePath = Paths.get(entityResponse.getFilePath());
            Resource resource = new UrlResource(filePath.toUri());

            if (!resource.exists()) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("resource", resource.exists());

                return ResponseEntity.ok().body(new BaseResponse(20004,
                        MessageCode.FILE_DOWNLOAD_ERROR.getCode(),
                        MessageCode.FILE_DOWNLOAD_ERROR.getDescription(),
                        errorData));
            }

            String contentType = determineContentType(filePath);

            return ResponseEntity.ok()
                    .contentType(MediaType.parseMediaType(contentType))
                    .header(HttpHeaders.CONTENT_DISPOSITION,
                            "attachment; filename=\"" + entityResponse.getOriginalFileName() + "\"")
                    .header(HttpHeaders.CONTENT_LENGTH, String.valueOf(entityResponse.getFileSize()))
                    .body(resource);
        } catch (FileEntityNotFoundException e) {
            Map<String, Object> errorData = new HashMap<>();
            errorData.put("fileId", fileId);

            return ResponseEntity.ok().body(new BaseResponse(20004,
                    MessageCode.FILE_NOT_FOUND.getCode(),
                    MessageCode.FILE_NOT_FOUND.getDescription(),
                    errorData));
        } catch (NotFoundAuthToken | IOException e) {
            Map<String, Object> errorData = new HashMap<>();
            errorData.put("token", tokenHeader);
            errorData.put("fileId", fileId);
            errorData.put("version", version);

            return ResponseEntity.ok().body(new BaseResponse(20004,
                    MessageCode.FILE_DOWNLOAD_ERROR_NOT_CORRECT_FIELD.getCode(),
                    MessageCode.FILE_DOWNLOAD_ERROR_NOT_CORRECT_FIELD.getDescription(),
                    errorData));
        }
    }

    private ResponseEntity<BaseResponse> createValidationErrorResponse(BindingResult bindingResult, Integer msgId) {
        List<Map<String, String>> fieldErrors = bindingResult.getFieldErrors()
                .stream()
                .map(this::createErrorDetail)
                .collect(Collectors.toList());

        RegAnswer regAnswer = new RegAnswer();
        regAnswer.setFieldErrors(fieldErrors);

        return ResponseEntity.ok().body(new BaseResponse(msgId, MessageCode.INVALID_FIELD.getCode(),
                MessageCode.INVALID_FIELD.getDescription(), regAnswer));
    }

    private Map<String, String> createErrorDetail(FieldError fieldError) {
        Map<String, String> errorDetail = new HashMap<>();
        errorDetail.put("field", fieldError.getField());
        errorDetail.put("code", fieldError.getCode() != null ? fieldError.getCode() : "VALIDATION_ERROR");
        errorDetail.put("message", fieldError.getDefaultMessage());

        if (fieldError.getRejectedValue() != null &&
                !fieldError.getField().toLowerCase().contains("password")) {
            errorDetail.put("rejected_value", fieldError.getRejectedValue().toString());
        }

        return errorDetail;
    }


    private String determineContentType(Path filePath) throws IOException {
        String contentType = Files.probeContentType(filePath);
        if (contentType == null) {
            contentType = "application/octet-stream";
        }
        return contentType;
    }

    private boolean isPreviewSupported(String mimeType) {
        if (mimeType == null) {
            return false;
        }

        return mimeType.startsWith("image") ||
                mimeType.startsWith("text") ||
                mimeType.equals("pdf") ||
                mimeType.startsWith("video") ||
                mimeType.startsWith("audio");
    }
}