no-copy/src/main/java/ru/soune/nocopy/service/file/FileUploadServiceImpl.java

package ru.soune.nocopy.service.file;

import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;
import ru.soune.nocopy.dto.file.UploadProgressResponse;
import ru.soune.nocopy.entity.file.FileUploadSession;
import ru.soune.nocopy.entity.file.UploadStatus;
import ru.soune.nocopy.exception.ChunkSizeExceededException;
import ru.soune.nocopy.exception.FileUploadException;
import ru.soune.nocopy.exception.UploadSessionNotFoundException;
import ru.soune.nocopy.repository.FileUploadSessionRepository;

import java.io.*;
import java.nio.file.*;
import java.nio.file.attribute.PosixFilePermission;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.LocalDateTime;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Stream;

@Slf4j
@Service
@RequiredArgsConstructor
public class FileUploadServiceImpl implements FileUploadService {

    private final FileUploadSessionRepository sessionRepository;

    @Value("${file.storage.base-path}")
    private String basePath;

    @Value("${file.storage.chunk-size}")
    private int chunkSize;

    @Value("${file.storage.max-file-size}")
    private long maxFileSize;

    @Value("${file.storage.max-retry-attempts:3}")
    private int maxRetryAttempts;

    @Value("${file.storage.chunk-timeout-ms:300000}") // 5 минут
    private long chunkTimeoutMs;

    @Value("${file.storage.session-expiry-hours:24}")
    private int sessionExpiryHours;

    private Path storageRoot;

    @Autowired
    private FileEntityService fileEntityService;

    @PostConstruct
    public void init() {
        try {
            storageRoot = Paths.get(basePath).toAbsolutePath().normalize();
            log.info("Initializing file storage at: {}", storageRoot);

            createDirectoryIfNotExists(storageRoot);
            createDirectoryIfNotExists(storageRoot.resolve("temp"));
            createDirectoryIfNotExists(storageRoot.resolve("uploads"));

            if (!System.getProperty("os.name").toLowerCase().contains("win")) {
                setDirectoryPermissions(storageRoot);
            }

            log.info("File storage initialized successfully");

        } catch (IOException e) {
            log.error("Failed to initialize file storage at: {}", basePath, e);
            throw new RuntimeException("Storage initialization failed", e);
        }
    }

    @Override
    @Transactional
    public FileUploadSession initUpload(Long userId, String fileName,
                                        String fileType, String extension, long fileSize) {
        log.info("Initializing upload for user {}: {} ({} bytes, type: {})",
                userId, fileName, fileSize, fileType);

        if (fileSize > maxFileSize) {
            throw new FileUploadException(
                    String.format("File size %d exceeds maximum allowed size %d",
                            fileSize, maxFileSize));
        }

        if (fileSize <= 0) {
            throw new FileUploadException("File size must be positive");
        }

        int totalChunks = (int) Math.ceil((double) fileSize / chunkSize);
        log.debug("File will be split into {} chunks (chunk size: {} bytes)",
                totalChunks, chunkSize);
        Long chunkSize = totalChunks == 1 ? fileSize : 1048576L;

        FileUploadSession session = FileUploadSession.builder()
                .userId(userId)
                .extension(extension)
                .fileName(fileName)
                .fileType(fileType)
                .chunkSize(chunkSize)
                .fileSize(fileSize)
                .totalChunks(totalChunks)
                .chunksUploaded(0)
                .status(UploadStatus.INITIATED)
                .expiresAt(LocalDateTime.now().plusHours(sessionExpiryHours))
                .retryCount(0)
                .build();

        FileUploadSession savedSession = sessionRepository.save(session);
        log.info("Upload session created: {} for file: {}",
                savedSession.getUploadId(), fileName);

        return savedSession;
    }

    @Override
    @Transactional
    public UploadProgressResponse uploadChunk(String uploadId, Integer chunkNumber,
                                              MultipartFile chunkFile) {

        log.info("Processing chunk {} for session {}, size: {} bytes",
                chunkNumber, uploadId, chunkFile.getSize());

        FileUploadSession session = sessionRepository.findById(uploadId)
                .orElseThrow(() -> new UploadSessionNotFoundException(uploadId));

        validateSession(session);

        if (session.getExpiresAt().isBefore(LocalDateTime.now())) {
            handleExpiredSession(session);
            throw new FileUploadException("Upload session expired");
        }

        if (chunkNumber < 0 || chunkNumber >= session.getTotalChunks()) {
            throw new FileUploadException(
                    String.format("Invalid chunk number %d. Expected number: %d",
                            chunkNumber, session.getTotalChunks() - 1));
        }

        if (chunkNumber == 0 && chunkFile.getSize() > chunkSize || chunkNumber + 1 == (session.getTotalChunks())
                && chunkFile.getSize() > chunkSize) {
            throw new ChunkSizeExceededException(chunkFile.getSize(), chunkSize);
        }

        if (chunkNumber > 0 && chunkNumber + 1 < session.getTotalChunks() && chunkFile.getSize() != chunkSize) {
            throw new ChunkSizeExceededException(chunkFile.getSize(), chunkSize);
        }

        return processChunk(session, chunkNumber, chunkFile);
    }

    @Override
    @Transactional
    public void handleExpiredSession(FileUploadSession session) {
        session.setStatus(UploadStatus.FAILED);
        session.setLastError("Upload session expired");
        sessionRepository.save(session);

        CompletableFuture.runAsync(() -> cleanupSessionFiles(session));
    }

    @Async("fileUploadTaskExecutor")
    @Transactional
    public void assembleFileAsync(FileUploadSession session) {
        try {
            assembleFile(session);
            log.info("File assembly completed successfully for session: {}",
                    session.getUploadId());

        } catch (Exception e) {
            log.error("Failed to assemble file for session {}: {}",
                    session.getUploadId(), e.getMessage(), e);

            handleAssemblyFailure(session, e);
        }
    }

    @Override
    @Transactional
    public void cancelUpload(String uploadId) {
        FileUploadSession session = sessionRepository.findById(uploadId)
                .orElseThrow(() -> new UploadSessionNotFoundException(uploadId));

        if (session.getStatus() == UploadStatus.COMPLETED || session.getStatus() == UploadStatus.CANCELLED) {
            throw new FileUploadException("Cannot cancel completed or cancelled upload");
        }

        session.setStatus(UploadStatus.CANCELLED);
        sessionRepository.save(session);

        CompletableFuture.runAsync(() -> cleanupSessionFiles(session));

        log.info("Upload cancelled: {}", uploadId);
    }

    @Override
    @Transactional
    public void retryFailedUpload(String uploadId) {
        FileUploadSession session = sessionRepository.findById(uploadId)
                .orElseThrow(() -> new UploadSessionNotFoundException(uploadId));

        if (session.getStatus() != UploadStatus.FAILED) {
            throw new FileUploadException("Only failed uploads can be retried");
        }

        if (session.getRetryCount() >= maxRetryAttempts) {
            throw new FileUploadException("Max retry attempts exceeded");
        }

        session.setStatus(UploadStatus.UPLOADING);
        session.setRetryCount(session.getRetryCount() + 1);
        session.setLastError(null);

        sessionRepository.save(session);

        if (session.getChunksUploaded().equals(session.getTotalChunks())) {
            log.info("Retrying file assembly for session: {}", uploadId);
            assembleFileAsync(session);
        }

        log.info("Upload retry initiated for session: {} (attempt {})",
                uploadId, session.getRetryCount());
    }

    @Override
    public UploadProgressResponse getUploadProgress(String uploadId) {
        FileUploadSession session = sessionRepository.findById(uploadId)
                .orElseThrow(() -> new UploadSessionNotFoundException(uploadId));

        return UploadProgressResponse.fromSession(session);
    }

    private void validateSession(FileUploadSession session) {
        UploadStatus status = session.getStatus();

        if (status == UploadStatus.FAILED) {
            if (session.getRetryCount() >= maxRetryAttempts) {
                throw new FileUploadException(
                        "Upload failed after maximum retry attempts");
            }
            session.setStatus(UploadStatus.UPLOADING);
            session.setRetryCount(session.getRetryCount() + 1);
            sessionRepository.save(session);
            log.info("Retrying failed upload session: {}, attempt: {}",
                    session.getUploadId(), session.getRetryCount());
        }

        if (status == UploadStatus.COMPLETED) {
            throw new FileUploadException("Upload already completed");
        }

        if (status == UploadStatus.CANCELLED) {
            throw new FileUploadException("Upload was cancelled");
        }

        if (status == UploadStatus.INITIATED) {
            session.setStatus(UploadStatus.UPLOADING);
            sessionRepository.save(session);
        }
    }

    private UploadProgressResponse processChunk(FileUploadSession session,
                                                Integer chunkNumber,
                                                MultipartFile chunkFile) {
        String chunkPath = null;

        try {
            if (session.getChunkPaths().containsKey(chunkNumber)) {
                return handleExistingChunk(session, chunkNumber, chunkFile);
            }

            chunkPath = saveChunkWithIntegrityCheck(session, chunkNumber, chunkFile);

            session.getChunkPaths().put(chunkNumber, chunkPath);
            session.setChunksUploaded(session.getChunksUploaded() + 1);
            session.setStatus(UploadStatus.UPLOADING);
            session.setExpiresAt(LocalDateTime.now().plusMinutes(1));

            log.debug("Chunk {} saved successfully. Uploaded: {}/{}",
                    chunkNumber, session.getChunksUploaded(), session.getTotalChunks());

            if (session.getChunksUploaded().equals(session.getTotalChunks())) {
                log.info("All chunks uploaded for session {}. Starting assembly...",
                        session.getUploadId());
                sessionRepository.save(session);
                assembleFileAsync(session);
            } else {
                sessionRepository.save(session);
            }

            return UploadProgressResponse.fromSession(session);

        } catch (Exception e) {
            if (chunkPath != null) {
                cleanupFailedChunk(chunkPath);
            }

            log.error("Failed to process chunk {} for session {}: {}",
                    chunkNumber, session.getUploadId(), e.getMessage(), e);

            session.setStatus(UploadStatus.FAILED);
            session.setLastError(e.getMessage());
            sessionRepository.save(session);

            throw new FileUploadException("Failed to upload chunk: " + e.getMessage(), e);
        }
    }

    private UploadProgressResponse handleExistingChunk(FileUploadSession session,
                                                       Integer chunkNumber,
                                                       MultipartFile chunkFile) throws IOException {
        String existingPath = session.getChunkPaths().get(chunkNumber);
        Path chunkPath = Paths.get(existingPath);

        if (!Files.exists(chunkPath)) {
            log.warn("Chunk file missing, re-uploading: {}", chunkPath);
            session.getChunkPaths().remove(chunkNumber);
            session.setChunksUploaded(session.getChunksUploaded() - 1);
            sessionRepository.save(session);
            return processChunk(session, chunkNumber, chunkFile);
        }

        long existingSize = Files.size(chunkPath);
        if (existingSize != chunkFile.getSize()) {
            log.warn("Chunk size mismatch, re-uploading: {} != {}",
                    existingSize, chunkFile.getSize());
            Files.deleteIfExists(chunkPath);
            session.getChunkPaths().remove(chunkNumber);
            session.setChunksUploaded(session.getChunksUploaded() - 1);
            sessionRepository.save(session);
            return processChunk(session, chunkNumber, chunkFile);
        }

        log.debug("Chunk {} already uploaded and valid", chunkNumber);
        return UploadProgressResponse.fromSession(session);
    }

    private String saveChunkWithIntegrityCheck(FileUploadSession session,
                                               Integer chunkNumber,
                                               MultipartFile chunkFile) throws IOException {

        Long userId = session.getUserId();
        String uploadId = session.getUploadId();

        Path chunkDir = getChunkDirectory(userId, uploadId);
        Files.createDirectories(chunkDir);

        String tempFileName = String.format("chunk_%04d.%s.tmp",
                chunkNumber, UUID.randomUUID());
        String finalFileName = String.format("chunk_%04d.tmp", chunkNumber);

        Path tempPath = chunkDir.resolve(tempFileName);
        Path finalPath = chunkDir.resolve(finalFileName);

        try {
            chunkFile.transferTo(tempPath.toFile());

            long savedSize = Files.size(tempPath);
            long uploadedSize = chunkFile.getSize();

            if (uploadedSize > 0 && savedSize != uploadedSize) {
                throw new IOException(
                        String.format("Size mismatch: saved %d, uploaded %d",
                                savedSize, uploadedSize));
            }

            Files.move(tempPath, finalPath, StandardCopyOption.ATOMIC_MOVE);

            log.debug("Chunk {} saved successfully: {} bytes",
                    chunkNumber, savedSize);

            return finalPath.toString();

        } finally {
            Files.deleteIfExists(tempPath);
        }
    }

    private void cleanupFailedChunk(String chunkPath) {
        try {
            Files.deleteIfExists(Paths.get(chunkPath));
            log.debug("Cleaned up failed chunk: {}", chunkPath);
        } catch (IOException e) {
            log.warn("Failed to cleanup chunk: {}", chunkPath, e);
        }
    }

    private void assembleFile(FileUploadSession session) throws IOException {
        log.info("Starting file assembly for session: {} ({})",
                session.getUploadId(), session.getFileName());

        Path finalFilePath = null;
        String checksum = null;

        try {
            finalFilePath = prepareFinalFile(session);
            log.info("Final file path: {}", finalFilePath);

            validateAllChunksExist(session);

            mergeChunksToFile(session, finalFilePath);

            validateFinalFile(session, finalFilePath);

            checksum = calculateChecksum(finalFilePath);
            log.debug("File checksum calculated: {}", checksum);

            session.setFilePath(finalFilePath.toString());
            session.setChecksum(checksum);
            session.setStatus(UploadStatus.COMPLETED);
            session.setCompletedAt(LocalDateTime.now());
            sessionRepository.save(session);
            log.info("Upload session updated to COMPLETED: {}", session.getUploadId());

            try {
                fileEntityService.createFromUploadSession(session, checksum);
                log.info("FileEntity successfully created for session: {}",
                        session.getUploadId());
            } catch (Exception e) {
                log.error("Failed to create FileEntity for session {} (file uploaded, but metadata not saved): {}",
                        session.getUploadId(), e.getMessage());
            }

            cleanupSessionFiles(session);

            log.info("File assembly completed successfully: {} -> {} ({} bytes)",
                    session.getFileName(), finalFilePath, session.getFileSize());

        } catch (Exception e) {
            log.error("File assembly failed for session {}: {}",
                    session.getUploadId(), e.getMessage(), e);

            if (finalFilePath != null) {
                try {
                    Files.deleteIfExists(finalFilePath);
                    log.debug("Cleaned up partial file: {}", finalFilePath);
                } catch (IOException ioException) {
                    log.warn("Failed to cleanup partial file: {}", finalFilePath, ioException);
                }
            }

            session.setStatus(UploadStatus.FAILED);
            session.setLastError(e.getMessage());
            sessionRepository.save(session);

            throw new IOException("File assembly failed: " + e.getMessage(), e);
        }
    }

    private Path prepareFinalFile(FileUploadSession session) throws IOException {
        Path userUploadsDir = storageRoot.resolve("uploads")
                .resolve(String.valueOf(session.getUserId()))
                .resolve(extractCategory(session.getFileType()));

        Files.createDirectories(userUploadsDir);

        String safeFileName = generateUniqueFileName(session.getFileName(), session.getExtension());
        Path finalPath = userUploadsDir.resolve(safeFileName);

        if (Files.exists(finalPath)) {
            throw new IOException("File already exists: " + finalPath);
        }

        return finalPath;
    }

    private String generateUniqueFileName(String originalName, String extension) {
        int dotIndex = originalName.lastIndexOf('.');
        log.info("dotIndex: {}", dotIndex);
        log.info("originalName: {}", originalName);
        String nameWithoutExtension;

        if (dotIndex > 0) {
            nameWithoutExtension = originalName.substring(0, dotIndex)
                    .replaceAll("[^a-zA-Z0-9\\-_]", "_");
        } else {
            nameWithoutExtension = originalName.replaceAll("[^a-zA-Z0-9\\-_]", "_");
        }

        String timestamp = String.valueOf(System.currentTimeMillis());
        String uuid = UUID.randomUUID().toString().substring(0, 8);

        return String.format("%s_%s_%s.%s",
                nameWithoutExtension, timestamp, uuid, extension);
    }

    private void validateAllChunksExist(FileUploadSession session) throws IOException {
        Path chunkDir = getChunkDirectory(session.getUserId(), session.getUploadId());

        for (int i = 0; i < session.getTotalChunks(); i++) {
            Path chunkPath = chunkDir.resolve(String.format("chunk_%04d.tmp", i));
            if (!Files.exists(chunkPath)) {
                throw new IOException("Missing chunk: " + chunkPath.getFileName());
            }

            if (Files.size(chunkPath) == 0) {
                throw new IOException("Empty chunk: " + chunkPath.getFileName());
            }
        }
    }

    private void mergeChunksToFile(FileUploadSession session, Path finalPath)
            throws IOException {

        Path chunkDir = getChunkDirectory(session.getUserId(), session.getUploadId());
        byte[] buffer = new byte[8192];

        try (BufferedOutputStream outputStream = new BufferedOutputStream(
                Files.newOutputStream(finalPath, StandardOpenOption.CREATE_NEW))) {

            long totalWritten = 0;

            for (int i = 0; i < session.getTotalChunks(); i++) {
                Path chunkPath = chunkDir.resolve(String.format("chunk_%04d.tmp", i));

                try (BufferedInputStream inputStream = new BufferedInputStream(
                        Files.newInputStream(chunkPath))) {

                    int bytesRead;
                    while ((bytesRead = inputStream.read(buffer)) != -1) {
                        outputStream.write(buffer, 0, bytesRead);
                        totalWritten += bytesRead;
                    }
                }

                if (log.isDebugEnabled()) {
                    log.debug("Processed chunk {}: {} bytes (total: {})",
                            i, Files.size(chunkPath), totalWritten);
                }
            }

            outputStream.flush();
            log.info("Total bytes written to final file: {}", totalWritten);
        }
    }

    private void validateFinalFile(FileUploadSession session, Path finalPath)
            throws IOException {

        long actualSize = Files.size(finalPath);

        if (actualSize != session.getFileSize()) {
            throw new IOException(
                    String.format("File size mismatch: expected %d, got %d",
                            session.getFileSize(), actualSize));
        }

        String checksum = calculateChecksum(finalPath);
        log.debug("Final file checksum: {}", checksum);
    }

    private void cleanupSessionFiles(FileUploadSession session) {
        try {
            Path chunkDir = getChunkDirectory(session.getUserId(), session.getUploadId());
            if (Files.exists(chunkDir)) {
                deleteDirectoryRecursively(chunkDir);
                log.info("Cleaned up temp directory for session: {}",
                        session.getUploadId());
            }
        } catch (IOException e) {
            log.warn("Failed to cleanup temp files for session {}: {}",
                    session.getUploadId(), e.getMessage());
        }
    }

    private void deleteDirectoryRecursively(Path directory) throws IOException {
        try (Stream<Path> walk = Files.walk(directory)) {
            walk.sorted(Comparator.reverseOrder())
                    .forEach(path -> {
                        try {
                            Files.deleteIfExists(path);
                        } catch (IOException e) {
                            log.warn("Failed to delete: {}", path, e);
                        }
                    });
        }
    }

    private void handleAssemblyFailure(FileUploadSession session, Exception e) {
        session.setStatus(UploadStatus.FAILED);
        session.setLastError(e.getMessage());
        sessionRepository.save(session);

        log.error("File assembly failed for session {}: {}",
                session.getUploadId(), e.getMessage());
    }



//    @Override
//    @Transactional
//    public void cleanupExpiredSessions() {
//        log.info("Starting cleanup of expired upload sessions");
//
//        LocalDateTime expiryThreshold = LocalDateTime.now().minusHours(sessionExpiryHours);
//
//        try {
//            var expiredSessions = sessionRepository.findByStatusInAndCreatedAtBefore(
//                    Set.of(UploadStatus.INITIATED, UploadStatus.UPLOADING, UploadStatus.FAILED),
//                    expiryThreshold
//            );
//
//            int cleanedCount = 0;
//            for (FileUploadSession session : expiredSessions) {
//                session.setStatus(UploadStatus.FAILED);
//                session.setLastError("Session expired during cleanup");
//                sessionRepository.save(session);
//
//                CompletableFuture.runAsync(() -> cleanupSessionFiles(session));
//                cleanedCount++;
//            }
//
//            log.info("Cleaned up {} expired upload sessions", cleanedCount);
//
//        } catch (Exception e) {
//            log.error("Error during session cleanup", e);
//        }
//    }

    private Path getChunkDirectory(Long userId, String uploadId) {
        return storageRoot.resolve("temp")
                .resolve("user")
                .resolve(String.valueOf(userId))
                .resolve(uploadId);
    }

    private String extractCategory(String fileType) {
        if (fileType == null) {
            return "other";
        }

        String lowerType = fileType.toLowerCase();
        if (lowerType.startsWith("image")) {
            return "images";
        } else if (lowerType.startsWith("video")) {
            return "videos";
        } else if (lowerType.startsWith("audio")) {
            return "audio";
        } else if (lowerType.contains("document")) {
            return "documents";
        } else {
            return "other";
        }
    }

    private String calculateChecksum(Path filePath) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            byte[] buffer = new byte[8192];

            try (InputStream inputStream = Files.newInputStream(filePath)) {
                int bytesRead;
                while ((bytesRead = inputStream.read(buffer)) != -1) {
                    md.update(buffer, 0, bytesRead);
                }
            }

            byte[] digest = md.digest();
            StringBuilder sb = new StringBuilder();
            for (byte b : digest) {
                sb.append(String.format("%02x", b));
            }

            return sb.toString();

        } catch (NoSuchAlgorithmException | IOException e) {
            log.warn("Failed to calculate checksum for file: {}", filePath, e);
            return "N/A";
        }
    }

    private void createDirectoryIfNotExists(Path directory) throws IOException {
        if (!Files.exists(directory)) {
            Files.createDirectories(directory);
            log.debug("Created directory: {}", directory);
        }
    }

    private void setDirectoryPermissions(Path directory) throws IOException {
        try {
            Set<PosixFilePermission> permissions = new HashSet<>();
            permissions.add(PosixFilePermission.OWNER_READ);
            permissions.add(PosixFilePermission.OWNER_WRITE);
            permissions.add(PosixFilePermission.OWNER_EXECUTE);
            permissions.add(PosixFilePermission.GROUP_READ);
            permissions.add(PosixFilePermission.GROUP_EXECUTE);
            permissions.add(PosixFilePermission.OTHERS_READ);

            Files.setPosixFilePermissions(directory, permissions);
            log.debug("Set permissions for directory: {}", directory);

        } catch (UnsupportedOperationException e) {
            log.debug("Posix permissions not supported on this system");
        }
    }
}