no-copy/src/main/java/ru/soune/nocopy/handler/FileUploadHandler.java

package ru.soune.nocopy.handler;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.validation.BeanPropertyBindingResult;
import org.springframework.validation.BindingResult;
import org.springframework.validation.FieldError;
import ru.soune.nocopy.dto.*;
import ru.soune.nocopy.dto.file.*;
import ru.soune.nocopy.entity.AuthToken;
import ru.soune.nocopy.entity.file.FileType;
import ru.soune.nocopy.entity.file.FileUploadSession;
import ru.soune.nocopy.exception.*;
import ru.soune.nocopy.handler.validator.FileUploadRequestValidator;
import ru.soune.nocopy.repository.AuthTokenRepository;
import ru.soune.nocopy.repository.FileUploadSessionRepository;
import ru.soune.nocopy.service.file.FileUploadService;

import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Component
@RequiredArgsConstructor
public class FileUploadHandler implements RequestHandler {
    private final FileUploadService fileUploadService;
    private final AuthTokenRepository authTokenRepository;
    private final FileUploadSessionRepository fileUploadSessionRepository;
    private final ObjectMapper objectMapper;
    private final FileUploadRequestValidator fileUploadRequestValidator;

    @Override
    public BaseResponse handle(BaseRequest request) {
        try {
            FileUploadRequest fileRequest = objectMapper.convertValue(
                    request.getMessageBody(), FileUploadRequest.class);

            String action = fileRequest.getAction();

            switch (action) {
                case "init":
                    return handleInitUpload(request, fileRequest);
                case "file_types":
                    return handleGetFileTypes(request);
                case "file_extension":
                    return handleGetFileExtensions(fileRequest, request);
                case "chunks":
                    return handleGetChunkStatus(request, fileRequest);
                case "cancel":
                    return handleCancelUpload(request, fileRequest);
                case "progress":
                    return handleGetProgress(request, fileRequest);
                case "retry":
                    return handleRetryUpload(request, fileRequest);
                default:
                    ActionResponse response = ActionResponse.builder()
                            .action(action)
                            .availableActions(Arrays.asList("init", "file_types", "chunks", "cancel", "progress", "retry"))
                            .build();

                    return new BaseResponse(request.getMsgId(), MessageCode.INVALID_ACTION.getCode(),
                            "Invalid action: " + action, response);
            }
        } catch (Exception e) {
            log.error("Error in FileUploadHandler", e);
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Handler error: " + e.getMessage(), null);
        }
    }

    private BaseResponse handleInitUpload(BaseRequest request, FileUploadRequest fileRequest) {
            String token = fileRequest.getToken();

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

            if (tokenOptional.isEmpty()) {
                return new BaseResponse(request.getMsgId(), MessageCode.INVALID_TOKEN.getCode(),
                        MessageCode.INVALID_TOKEN.getDescription(), Map.of("token", token));
            }

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

            BindingResult bindingResult = new BeanPropertyBindingResult(fileRequest, "fileRequest");
            fileUploadRequestValidator.validate(fileRequest, bindingResult);

            if (bindingResult.hasErrors()) {
                Map<String, String> fieldErrors = bindingResult.getFieldErrors()
                        .stream()
                        .collect(Collectors.toMap(
                                FieldError::getField,
                                fieldError -> fieldError.getDefaultMessage() != null
                                        ? fieldError.getDefaultMessage()
                                        : "Validation error"));

                return new BaseResponse(request.getMsgId(), MessageCode.INVALID_FIELD.getCode(),
                        MessageCode.INVALID_FIELD.getDescription(), fieldErrors);
            }

            FileUploadSession session = fileUploadService.initUpload(
                    authToken.getUser().getId(),
                    fileRequest.getFileName(),
                    fileRequest.getFileType(),
                    fileRequest.getExtension(),
                    fileRequest.getFileSize());

            InitFileResponse fileResponse = InitFileResponse.builder()
                    .uploadId(session.getUploadId())
                    .fileName(session.getFileName())
                    .totalChunks(session.getTotalChunks())
                    .chunkSize(session.getChunkSize())
                    .status(session.getStatus().toString())
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    MessageCode.SUCCESS.getDescription(), fileResponse);
    }

    private BaseResponse handleGetProgress(BaseRequest request, FileUploadRequest fileRequest) {
        try {
            ru.soune.nocopy.dto.file.UploadProgressResponse progress =
                    fileUploadService.getUploadProgress(fileRequest.getUploadId());

            UploadProgressResponse responseDto = UploadProgressResponse.builder()
                    .uploadId(progress.getUploadId())
                    .fileName(progress.getFileName())
                    .totalChunks(progress.getTotalChunks())
                    .uploadedChunks(progress.getUploadedChunks())
                    .status(progress.getStatus())
                    .progressPercentage(progress.getProgressPercentage())
                    .filePath(progress.getFilePath())
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    MessageCode.SUCCESS.getDescription(), responseDto);

        } catch (UploadSessionNotFoundException e) {
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_NOT_FOUND.getCode(),
                    MessageCode.FILE_NOT_FOUND.getDescription(), null);
        } catch (Exception e) {
            log.error("Error getting progress", e);
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    MessageCode.FILE_UPLOAD_ERROR.getDescription(), null);
        }
    }

    private BaseResponse handleCancelUpload(BaseRequest request, FileUploadRequest fileRequest) {
        try {
            fileUploadService.cancelUpload(fileRequest.getUploadId());

            CanceledUploadResponse response = CanceledUploadResponse.builder()
                    .uploadId(fileRequest.getUploadId())
                    .message("Upload cancelled successfully")
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    "Upload cancelled successfully", response);

        } catch (UploadSessionNotFoundException e) {
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_NOT_FOUND.getCode(),
                    "Upload session not found", null);
        } catch (FileUploadException e) {
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(), e.getMessage(), null);
        } catch (Exception e) {
            log.error("Error cancelling upload", e);
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to cancel upload", null);
        }
    }

    private BaseResponse handleRetryUpload(BaseRequest request, FileUploadRequest fileRequest) {
        try {
            fileUploadService.retryFailedUpload(fileRequest.getUploadId());
            ru.soune.nocopy.dto.file.UploadProgressResponse progress =
                    fileUploadService.getUploadProgress(fileRequest.getUploadId());

            RetryUploadResponse response = RetryUploadResponse.builder()
                    .uploadId(progress.getUploadId())
                    .message("Upload retry initiated")
                    .status(progress.getStatus().toString())
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(), "Upload retry initiated",
                    response);

        } catch (UploadSessionNotFoundException e) {
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_NOT_FOUND.getCode(),
                    "Upload session not found", null);
        } catch (FileUploadException e) {
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(), e.getMessage(),
                    null);
        } catch (Exception e) {
            log.error("Error retrying upload", e);
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to retry upload", null);
        }
    }

    private BaseResponse handleGetChunkStatus(BaseRequest request, FileUploadRequest fileRequest) {
        try {
            FileUploadSession session = fileUploadSessionRepository.findById(fileRequest.getUploadId())
                    .orElseThrow(() -> new UploadSessionNotFoundException(fileRequest.getUploadId()));

            Map<String, Boolean> chunkStatus = new HashMap<>();
            for (int i = 0; i < session.getTotalChunks(); i++) {
                chunkStatus.put("chunk_" + i, session.getChunkPaths().containsKey(i));
            }

            ChunkStatusResponse response = ChunkStatusResponse.builder()
                    .uploadId(fileRequest.getUploadId())
                    .totalChunks(session.getTotalChunks())
                    .uploadedChunks(session.getChunksUploaded())
                    .missingChunks(session.getTotalChunks() - session.getChunksUploaded())
                    .chunkStatus(chunkStatus)
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    MessageCode.SUCCESS.getDescription(), response);

        } catch (UploadSessionNotFoundException e) {
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_NOT_FOUND.getCode(),
                    "Upload session not found", null);
        } catch (Exception e) {
            log.error("Error getting chunk status", e);
            return new BaseResponse(request.getMsgId(),
                    MessageCode.FILE_UPLOAD_ERROR.getCode(), "Failed to get chunk status",
                    null);
        }
    }

    private BaseResponse handleGetFileTypes(BaseRequest request) {
        try {
            FileTypesResponse response = FileTypesResponse.builder()
                    .fileTypes(Arrays.asList(FileType.values()))
                    .count(FileType.values().length)
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    MessageCode.SUCCESS.getDescription(), response);

        } catch (Exception e) {
            log.error("Error getting file types", e);
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to get file types", null);
        }
    }

    private BaseResponse handleGetFileExtensions(FileUploadRequest fileRequest, BaseRequest request) {
        try {
            FileType fileType = FileType.valueOf(fileRequest.getFileType().toUpperCase());
            List<String> allowedExtensions = fileType.getAllowedExtensions();

            FileExtensionResponse response = FileExtensionResponse.builder()
                    .extension(allowedExtensions)
                    .count(allowedExtensions.size())
                    .build();

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    MessageCode.SUCCESS.getDescription(), response);
        } catch (Exception e) {
            log.error("Error getting file extensions", e);
            return new BaseResponse(request.getMsgId(), MessageCode.FILE_UPLOAD_ERROR.getCode(),
                    "Failed to get file extensions", null);
        }
    }
}