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

package ru.soune.nocopy.controller;

import com.vrt.NoCopyFileService;
import com.vrt.fileprotection.FileProtector;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.Resource;
import org.springframework.core.io.UrlResource;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.web.PageableDefault;
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.*;
import ru.soune.nocopy.entity.file.FileEntity;
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.AuthToken;
import ru.soune.nocopy.entity.file.FileEntity;
import ru.soune.nocopy.entity.file.FileStatus;
import ru.soune.nocopy.entity.file.ProtectionStatus;
import ru.soune.nocopy.entity.file.UploadStatus;
import ru.soune.nocopy.exception.*;
import ru.soune.nocopy.handler.*;
import ru.soune.nocopy.repository.AuthTokenRepository;
import ru.soune.nocopy.repository.FileEntityRepository;
import ru.soune.nocopy.repository.FileEntityRepository;
import ru.soune.nocopy.service.FileSimilarityService;
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.Optional;
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;

    private final FileSimilarityService fileSimilarityService;

    private final FileEntityRepository fileEntityRepository;

    private final NoCopyFileService noCopyFileService;

    @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 buildErrorResponse(uploadId, chunkNumber, "Chunk file null or empty");
            }
            if (uploadId == null || uploadId.isBlank()) {
                return buildErrorResponse(uploadId, chunkNumber, "Upload ID is required");
            }
            if (chunkNumber == null || chunkNumber < 0) {
                return buildErrorResponse(uploadId, chunkNumber, "Valid chunk number is required");
            }

            fileUploadService.uploadChunk(uploadId, chunkNumber, chunk);

            return buildSuccessResponse(uploadId, chunkNumber, chunk);
        } catch (DuplicateImageException e) {
            return ResponseEntity.ok().body(new BaseResponse(20004, MessageCode.DUPLICATE_FILE_UPLOAD.getCode(),
                    MessageCode.DUPLICATE_FILE_UPLOAD.getDescription(),
                    Map.of("duplicateOwnerId", e.userId(), "duplicateFileId", e.duplicateFileId())));
        } catch (Exception e) {
            log.error("Error uploading chunk", e);
            return buildErrorResponse(uploadId, chunkNumber, "Failed to upload chunk: " + e.getMessage());
        }
    }

    @GetMapping("/v{version}/files/{fileId}/similar")
    public ResponseEntity<BaseResponse> findSimilarFiles(
            @PathVariable("version") int version,
            @PathVariable String fileId,
            @RequestParam(required = false) List<String> similarityLevels,
            @PageableDefault(size = 20, sort = "hammingDistance") Pageable pageable) {
        SimilarityFilter filter = SimilarityFilter.builder()
                .similarityLevels(similarityLevels)
                .build();
        Page<SimilarFileResponse> similarFiles = fileSimilarityService.findSimilarFiles(fileId, filter, pageable);

        String messageDesc;
        MessageCode success;

        if (similarFiles.isEmpty()) {
            messageDesc = MessageCode.FILE_NOT_FOUND.getDescription();
            success = MessageCode.FILE_NOT_FOUND;
        } else {
            messageDesc = MessageCode.SIMILAR_FILES_FOUND.getDescription();
            success = MessageCode.SIMILAR_FILES_FOUND;
        }

        Map<String, Object> responseData = new HashMap<>();
        responseData.put("content", similarFiles.getContent());
        responseData.put("page", similarFiles.getNumber());
        responseData.put("size", similarFiles.getSize());
        responseData.put("totalElements", similarFiles.getTotalElements());
        responseData.put("totalPages", similarFiles.getTotalPages());
        responseData.put("hasNext", similarFiles.hasNext());
        responseData.put("hasPrevious", similarFiles.hasPrevious());

        return ResponseEntity.ok()
                .body(new BaseResponse(20004, success.getCode(), messageDesc, responseData));
    }

    @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));
        }
    }

    @GetMapping("/protect/{fileId}")
    public ResponseEntity<?> protect( @PathVariable(required = false) String fileId) {
        Optional<FileEntity> fileEntity = fileEntityRepository.findById(fileId);

        if (!fileEntity.isPresent()) {
            return ResponseEntity.notFound().build();
        }

        FileProtector.FileInfo fileInfo = createFileInfo(fileEntity.get());

        noCopyFileService.addFile(fileInfo);

        fileEntity.get().setProtectionStatus(ProtectionStatus.PROCESSING);

        fileEntityRepository.save(fileEntity.get());

        return ResponseEntity.ok().build();
    }

    private FileProtector.FileInfo createFileInfo(FileEntity fileEntity) {
        FileProtector.Type type = determineFileType(fileEntity.getMimeType());

        return new FileProtector.FileInfo(type, fileEntity.getId(), String.valueOf(fileEntity.getUserId()));
    }

    private FileProtector.Type determineFileType(String mimeType) {
        if (mimeType == null) {
            return FileProtector.Type.IMAGE;
        }

        if (mimeType.startsWith("image")) {
            return FileProtector.Type.IMAGE;
        } else if (mimeType.startsWith("video")) {
            return FileProtector.Type.VIDEO;
        } else if (mimeType.startsWith("audio")) {
            return FileProtector.Type.AUDIO;
        } else {
            return FileProtector.Type.IMAGE;
        }
    }

    private ResponseEntity<BaseResponse> checkForDuplicates(String uploadId) throws IOException {
        Optional<FileEntity> uploadedFile = fileEntityRepository.findByUploadSessionId(uploadId);

        if (uploadedFile.isEmpty() || !uploadedFile.get().getMimeType().equals("image")) {
            return null;
        }

        FileEntity fileEntity = uploadedFile.get();
        List<SimilarFileResponse> similarFiles = fileSimilarityService.findSimilarFiles(fileEntity.getId());

        if (hasDuplicate(similarFiles)) {
            return handleDuplicate(fileEntity, similarFiles);
        }

        return null;
    }

    private boolean hasDuplicate(List<SimilarFileResponse> similarFiles) {
        return similarFiles.stream().anyMatch(f -> f.getHammingDistance() <= 5);
    }

    private ResponseEntity<BaseResponse> handleDuplicate(FileEntity fileEntity, List<SimilarFileResponse> similarFiles)
            throws IOException {
        fileEntityService.deleteFromDisk(fileEntity);

        fileEntityService.markAsDeleted(fileEntity);

        Optional<FileEntity> originalFile = fileEntityRepository.findById(similarFiles.get(0).getFileId());

        if (originalFile.isPresent()) {
            Map<String, String> duplicateInfo = Map.of(
                    "duplicate_file_id", originalFile.get().getId(),
                    "owner_user_id", String.valueOf(originalFile.get().getUserId()));

            return ResponseEntity.ok().body(new BaseResponse(
                    20004,
                    MessageCode.DUPLICATE_FILE_UPLOAD.getCode(),
                    "Failed to upload chunk, duplicate",
                    duplicateInfo));
        }

        return null;
    }

    private ResponseEntity<BaseResponse> buildSuccessResponse(String uploadId, Integer chunkNumber, MultipartFile chunk) {
        ChunkUploadResponse responseBody = ChunkUploadResponse.builder()
                .uploadId(uploadId)
                .chunkNumber(chunkNumber)
                .chunkSize(chunk.getSize())
                .message("Chunk uploaded successfully")
                .build();

        return ResponseEntity.ok().body(new BaseResponse(
                20000,
                MessageCode.SUCCESS.getCode(),
                "Chunk uploaded successfully",
                responseBody));
    }

    private ResponseEntity<BaseResponse> buildErrorResponse(String uploadId, Integer chunkNumber, String errorMessage) {
        ChunkUploadResponse responseBody = ChunkUploadResponse.builder()
                .uploadId(uploadId)
                .chunkNumber(chunkNumber)
                .build();

        return ResponseEntity.ok().body(new BaseResponse(
                20004,
                MessageCode.FILE_UPLOAD_ERROR.getCode(),
                errorMessage,
                responseBody));
    }

    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 Long getUserIdFromToken(String tokenHeader) {
        String token = tokenHeader.replace("Bearer ", "");
        AuthToken authToken = authTokenRepository.findByToken(token)
                .orElseThrow(() -> new NotFoundAuthToken("Token not found"));
        return authToken.getUser().getId();
    }

    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");
    }
}