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 com.vrt.fileprotection.NoCopyCheckResult;
import com.vrt.fileprotection.audio.AudioCheckResult;
import com.vrt.fileprotection.documents.DocumentCheckResult;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.FileSystemResource;
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.file.*;
import ru.soune.nocopy.dto.register.RegAnswer;
import ru.soune.nocopy.entity.company.Company;
import ru.soune.nocopy.entity.file.FileEntity;
import ru.soune.nocopy.entity.file.FileStatus;
import ru.soune.nocopy.entity.file.UploadStatus;
import ru.soune.nocopy.entity.user.AuthToken;
import ru.soune.nocopy.entity.user.User;
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.UserRepository;
import ru.soune.nocopy.service.FileSimilarityService;
import ru.soune.nocopy.service.file.ProtectionsLimitService;
import ru.soune.nocopy.service.register.AuthService;
import ru.soune.nocopy.service.file.FileEntityService;
import ru.soune.nocopy.service.file.FileUploadService;
import ru.soune.nocopy.util.FileUtil;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
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 FileUtil fileUtil;

    private final ProtectionsLimitService protectionsLimitService;

    private final AuthTokenRepository authTokenRepository;

    private final UserRepository userRepository;

    @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 (UserNotFoundException e) {
            return ResponseEntity.ok().body(new BaseResponse(msgId, MessageCode.USER_NOT_FOUND.getCode(),
                    MessageCode.USER_NOT_FOUND.getDescription(), new HashMap<>()));
        } 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,
            @RequestParam(value = "findSimilar", required = false, defaultValue = "0") Integer findSimilar) {
        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");
            }

            UploadProgressResponse uploadProgressResponse = fileUploadService.uploadChunk(uploadId, chunkNumber,
                    chunk, findSimilar);

            return buildSuccessResponse(uploadId, chunkNumber, chunk,
                    fileEntityService.findFileIdByPath(uploadProgressResponse.getFilePath()));
        } catch (DuplicateImageException e) {
            Map<String, Object> duplicateData = new HashMap<>();
            duplicateData.put("duplicateFileId", e.duplicateFileId());
            duplicateData.put("userId", e.userId());
            duplicateData.put("message", e.getMessage());

            if (uploadId != null) {
                duplicateData.put("uploadId", uploadId);
            }

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

    private final NoCopyFileService noCopyFileService;

    @GetMapping("/v{version}/files/{fileId}/similar")
    public ResponseEntity<BaseResponse> findSimilarFiles(
            @PathVariable("version") int version,
            @PathVariable String fileId,
            @RequestParam(value = "auth_token", required = false) String authToken,
            @RequestParam(required = false) List<String> similarityLevels,
            @PageableDefault(size = 20, sort = "hammingDistance") Pageable pageable) {

        try {
            FileEntity fileEntity = fileEntityRepository.findByFileId(fileId);
            if (fileEntity == null) {
                return ResponseEntity.ok().body(new BaseResponse(0, MessageCode.INVALID_FIELD.getCode(),
                        MessageCode.INVALID_FIELD.getDescription(), null));
            }

            BaseResponse response;
            if ("image".equals(fileEntity.getMimeType())) {
                Page<SimilarFileDTO> similarFilesPage = fileSimilarityService.findSimilarFiles(
                        fileId, similarityLevels, pageable, authToken);
                response = buildSuccessResponse(similarFilesPage, MessageCode.SIMILAR_FILES_FOUND);
            } else {
                List<SimilarFileDTO> results = processNonImageFile(fileEntity, authToken);
                response = buildSuccessResponse(results, MessageCode.SIMILAR_FILES_FOUND);
            }

            return ResponseEntity.ok(response);

        } catch (Exception e) {
            log.error("Error finding similar files: {}", e.getMessage(), e);
            return ResponseEntity.ok().body(new BaseResponse(0, MessageCode.INVALID_FIELD.getCode(),
                    MessageCode.INVALID_FIELD.getDescription(), null));
        }
    }

    private List<SimilarFileDTO> processNonImageFile(FileEntity fileEntity, String authToken) throws Exception {
        List<SimilarFileDTO> results = new ArrayList<>();

        FileEntity duplicateByHash = fileSimilarityService.findDuplicateByHash(fileEntity.getFilePath(),
                fileEntity.getMimeType(), fileEntity.getUserId());

        if (duplicateByHash != null) {
            results.add(fileSimilarityService.buildDTO(duplicateByHash));

            return results;
        }

        File file = new File(fileEntity.getProtectedFilePath());
        FileProtector.Type type = "document".equals(fileEntity.getMimeType()) ?
                FileProtector.Type.DOC : FileProtector.Type.AUDIO;

        NoCopyCheckResult checkResult = noCopyFileService.checkFile(file, type);

        switch (checkResult) {
            case DocumentCheckResult.Success success -> {
                FileProtector.FileInfo info = success.getInfo();
                FileEntity entity = fileEntityRepository.findByFileId(info.getId());
                if (entity != null) {
                    results.add(fileSimilarityService.buildDTO(entity));
                }
            }
            case AudioCheckResult.Success success -> {
                FileProtector.FileInfo info = success.getInfo();
                FileEntity entity = fileEntityRepository.findByFileId(info.getId());
                if (entity != null) {
                    results.add(fileSimilarityService.buildDTO(entity));
                }
            }
            case DocumentCheckResult.Failed failed ->
                    log.warn("Document check failed: {}", failed.getMessage());
            case AudioCheckResult.Failed failed ->
                    log.warn("Audio check failed: {}", failed.getMessage());
            default -> log.warn("Unexpected result type: {}", checkResult.getClass().getSimpleName());
        }

        return results;
    }

    private BaseResponse buildSuccessResponse(Page<SimilarFileDTO> page, MessageCode messageCode) {
        Map<String, Object> responseData = new HashMap<>();
        responseData.put("content", page.getContent());
        responseData.put("page", page.getNumber());
        responseData.put("size", page.getSize());
        responseData.put("totalElements", page.getTotalElements());
        responseData.put("totalPages", page.getTotalPages());
        responseData.put("hasNext", page.hasNext());
        responseData.put("hasPrevious", page.hasPrevious());

        return new BaseResponse(20004, messageCode.getCode(),
                messageCode.getDescription(), responseData);
    }

    private BaseResponse buildSuccessResponse(List<SimilarFileDTO> list, MessageCode messageCode) {
        Map<String, Object> responseData = new HashMap<>();
        responseData.put("content", list);
        responseData.put("totalElements", list.size());

        return new BaseResponse(20004, messageCode.getCode(),
                messageCode.getDescription(), 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 || tokenHeader.isBlank()) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("token", tokenHeader);

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

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

            User fileUser = userRepository.findById(entityResponse.getUserId()).orElseThrow();
            User user = userRepository.findById(userId).orElseThrow();
            Company companyUser = user.getCompany();

            if (companyUser != null && !fileUser.getCompany().getId().equals(companyUser.getId())) {
                Map<String, Object> errorData = new HashMap<>();
                errorData.put("token", tokenHeader);

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

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

                return ResponseEntity.ok().body(new BaseResponse(20004,
                        MessageCode.USER_NOT_HAD_PERMISSION.getCode(),
                        MessageCode.USER_NOT_HAD_PERMISSION.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_DELETE.getCode(),
                        MessageCode.FILE_DELETE.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_NOT_EXIST.getCode(),
                        MessageCode.FILE_NOT_EXIST.getDescription(),
                        errorData));
            }


            Path filePath = Paths.get(entityResponse.getProtectedFilePath());
            FileSystemResource resource = new FileSystemResource(filePath);
            long fileSize = Files.size(filePath);

            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()
                    .contentLength(fileSize)
                    .contentType(MediaType.parseMediaType(contentType))
                    .header(HttpHeaders.CONTENT_DISPOSITION,
                            "attachment; filename=\"" + entityResponse.getFileName() + "\"")
                    .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 e) {
            Map<String, Object> errorData = new HashMap<>();
            errorData.put("token", tokenHeader);

            return ResponseEntity.ok().body(new BaseResponse(20004,
                    MessageCode.AUTH_TOKEN_NOT_FOUND.getCode(),
                    MessageCode.AUTH_TOKEN_NOT_FOUND.getDescription(),
                    errorData));
        } catch (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("/check/file_stats")
    public ResponseEntity<?> checkFileProtectStats(@RequestHeader("Authorization") String tokenHeader) {
        String token = tokenHeader.replace("Bearer ", "");

        Optional<AuthToken> tokenOptional = authTokenRepository.findByToken(token);

        AuthToken authToken = tokenOptional.orElseThrow(() -> new NotFoundAuthToken("Token: " + token + "not found"));

        Map<String, Object> fileTypeStats = protectionsLimitService.getFileTypeStats(authToken.getUser().getId());

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

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

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

        fileEntityService.softDeleteFileWithHash(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,
                                                              String fileId) {
        ChunkUploadResponse responseBody = ChunkUploadResponse.builder()
                .uploadId(uploadId)
                .chunkNumber(chunkNumber)
                .chunkSize(chunk.getSize())
                .fileId(fileId)
                .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;
    }
}