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

package ru.soune.nocopy.service;

import com.vrt.fileprotection.image.phash.PHash;
import com.vrt.fileprotection.image.phash.PerceptualHashHelper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.Pageable;
import org.springframework.stereotype.Service;
import ru.soune.nocopy.dto.file.SimilarFileDTO;
import ru.soune.nocopy.dto.file.SimilarityFilter;
import ru.soune.nocopy.entity.file.FileEntity;
import ru.soune.nocopy.entity.file.FileStatus;
import ru.soune.nocopy.entity.file.ImageHashEntity;
import ru.soune.nocopy.exception.DuplicateImageException;
import ru.soune.nocopy.repository.*;
import ru.soune.nocopy.util.FileUtil;

import java.io.FileInputStream;
import java.io.InputStream;
import java.security.MessageDigest;
import java.util.*;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@Slf4j
public class FileSimilarityService {

    private final ImageSimilarityRepository repository;

    private final ImageHashRepository hashRepository;

    private final FileUtil fileUtil;

    private final AuthTokenRepository authTokenRepository;

    private final FileEntityRepository fileEntityRepository;

    @Value("${server.baseurl}")
    private String baseUrl;

    public List<SimilarFileDTO> findSimilarFiles(String fileId) {
        var imageHashEntity = hashRepository.findById(fileId)
                .orElseThrow(() -> new RuntimeException("Hash not found"));

        Long hash64Hi = imageHashEntity.getHash64Hi();
        Long hash64Lo = imageHashEntity.getHash64Lo();

        List<SimilarImageProjection> candidates = repository.findCandidates(fileId);

        return candidates.stream()
                .map(c -> {
                    Long cHi = c.getHash64Hi();
                    Long cLo = c.getHash64Lo();

                    int hamming = fileUtil.hamming64(hash64Hi, hash64Lo, cHi, cLo);

                    String level;
                    if (hamming <= 5) {
                        level = "DUPLICATE";
                    } else if (hamming <= 12) {
                        level = "SIMILAR";
                    } else {
                        level = "DIFFERENT";
                    }

                    return SimilarFileDTO.builder()
                            .fileId(c.getId())
                            .ownerId(c.getUserId())
                            .originalFileName(c.getOriginalFileName())
                            .fileSize(c.getFileSize())
                            .hammingDistance(hamming)
                            .similarityLevel(level)
                            .build();
                })
                .sorted((a, b) ->
                        Integer.compare(a.getHammingDistance(), b.getHammingDistance()))
                .toList();
    }

    public List<SimilarFileDTO> findDuplicateByHammingDistance(String fileId, int hammingDistance,
                                                               int duplicate, int similar) {
        List<SimilarImageProjection> candidates = repository.findCandidates(fileId);

        Optional<ImageHashEntity> hashOptional = hashRepository.findById(fileId);

        if (hashOptional.isEmpty()) {
            return Collections.emptyList();
        }

        var imageHashEntity = hashOptional.get();

        Long hash64Hi = imageHashEntity.getHash64Hi();
        Long hash64Lo = imageHashEntity.getHash64Lo();

        return candidates.stream()
                .map(c -> {
                    Long cHi = c.getHash64Hi();
                    Long cLo = c.getHash64Lo();

                    int hamming = fileUtil.hamming64(hash64Hi, hash64Lo, cHi, cLo);

                    return new AbstractMap.SimpleEntry<>(c, hamming);
                })
                .filter(entry -> entry.getValue() <= hammingDistance)
                .sorted((a, b) -> Integer.compare(a.getValue(), b.getValue()))
                .map(entry -> {
                    SimilarImageProjection similarImageProjection = entry.getKey();
                    int hamming = entry.getValue();

                    String level;
                    if (hamming <= duplicate) {
                        level = "DUPLICATE";
                    } else if (hamming <= similar) {
                        level = "SIMILAR";
                    } else {
                        level = "DIFFERENT";
                    }

                    return SimilarFileDTO.builder()
                            .fileId(similarImageProjection.getId())
                            .originalFileName(similarImageProjection.getOriginalFileName())
                            .fileSize(similarImageProjection.getFileSize())
                            .hammingDistance(hamming)
                            .similarityLevel(level)
                            .ownerId(similarImageProjection.getUserId())
                            .build();
                })
                .toList();
    }

    public List<SimilarImageProjection> findDuplicatedByHash(Long hash64Hi, Long hash64Lo) {
        List<SimilarImageProjection> duplicates = repository.findExactDuplicates(hash64Hi,hash64Lo);

        if (duplicates.isEmpty()) {
            return new ArrayList<>();
        }

        return duplicates;
    }

    public void hasDuplicatesByHash(String path, Long userId, String mimeType) throws Exception {
        String hash = calculateFileHash(path);
        List<FileEntity> fileEntityList = fileEntityRepository.findByUserIdAndMimeType(userId, mimeType);

        for (FileEntity file : fileEntityList) {
            if (calculateFileHash(file.getFilePath()).equals(hash)) {
                throw new DuplicateImageException("Duplicate", file.getId(),
                        file.getUserId());
            }
        }
    }

    public Page<SimilarFileDTO> findSimilarFiles(String fileId, SimilarityFilter filter, Pageable pageable,
                                                 String authToken) throws Exception {
        FileEntity fileEntity = fileEntityRepository.findByFileId(fileId);
        String mimeType = fileEntity.getMimeType();
        List<SimilarFileDTO> allResults = new ArrayList<>();

        if (mimeType.equals("image")) {
            var imageHashEntity = hashRepository.findById(fileId)
                    .orElseThrow(() -> new RuntimeException("Hash not found"));

            Long hash64Hi = imageHashEntity.getHash64Hi();
            Long hash64Lo = imageHashEntity.getHash64Lo();
            List<SimilarImageProjection> candidates = authToken.equals("all") ? repository.findCandidates(fileId):
                    repository.findCandidatesFromUserFiles(fileId, authTokenRepository.findUserIdByToken(authToken));

            List<String> similarityLevels = (filter != null && filter.getSimilarityLevels() != null)
                    ? filter.getSimilarityLevels()
                    : List.of("DUPLICATE", "SIMILAR", "DIFFERENT");

            allResults = candidates.stream()
                    .map(c -> createSimilarFileResponse(c, hash64Hi, hash64Lo))
                    .filter(response -> similarityLevels.contains(response.getSimilarityLevel()))
                    .sorted(Comparator.comparingInt(SimilarFileDTO::getHammingDistance))
                    .collect(Collectors.toList());
        } else {
            String hash = calculateFileHash(fileEntity.getFilePath());
            Long userId = fileEntity.getUserId();
            List<FileEntity> fileEntityList = fileEntityRepository.findByUserIdAndMimeType(userId, mimeType);

            for (FileEntity file : fileEntityList) {
                if (calculateFileHash(file.getFilePath()).equals(hash) && !file.getStatus().equals(FileStatus.TEMP)) {
                    allResults.add(buildDTO(file));
                }
            }
        }

        int total = allResults.size();
        int page = (pageable != null) ? pageable.getPageNumber() : 0;
        int size = (pageable != null) ? pageable.getPageSize() : 20;

        int fromIndex = Math.min(page * size, total);
        int toIndex = Math.min(fromIndex + size, total);

        List<SimilarFileDTO> pageContent = allResults.subList(fromIndex, toIndex);

        return new PageImpl<>(pageContent, pageable, total);
    }


    private String calculateFileHash(String path) throws Exception {
        MessageDigest digest = MessageDigest.getInstance("SHA-256");

        try (InputStream is = new FileInputStream(path)) {
            byte[] buffer = new byte[8192];
            int read;
            while ((read = is.read(buffer)) > 0) {
                digest.update(buffer, 0, read);
            }
        }

        byte[] hashBytes = digest.digest();
        return Base64.getEncoder().encodeToString(hashBytes);
    }

    private SimilarFileDTO createSimilarFileResponse(SimilarImageProjection similarImageProjection,
                                                     Long hash64Hi, Long hash64Lo) {
        Long imageProjectionHash64Hi = similarImageProjection.getHash64Hi();
        Long similarImageProjectionHash64Lo = similarImageProjection.getHash64Lo();
        int hamming = PerceptualHashHelper.INSTANCE.hammingDistance(new PHash(hash64Hi, hash64Lo),
                new PHash(imageProjectionHash64Hi, similarImageProjectionHash64Lo));

        String level;
        if (hamming <= 6) {
            level = "DUPLICATE";
        } else if (hamming <= 12) {
            level = "SIMILAR";
        } else {
            level = "DIFFERENT";
        }

        return SimilarFileDTO.builder()
                .fileId(similarImageProjection.getId())
                .ownerId(similarImageProjection.getUserId())
                .originalFileName(similarImageProjection.getOriginalFileName().replace("." +
                        similarImageProjection.getExtension(), "") + "_nocopy_protected" +
                        "." + similarImageProjection.getExtension())
                .fileSize(similarImageProjection.getFileSize())
                .hammingDistance(hamming)
                .similarityLevel(level)
                .supportId(similarImageProjection.getSupportId())
                .uploadDate(similarImageProjection.getUploadDate())
                .status(similarImageProjection.getProtectionStatus())
                .url(baseUrl + "/api/files/protected/" + similarImageProjection.getId())
                .build();
    }

    private SimilarFileDTO buildDTO(FileEntity fileEntity) {
        return SimilarFileDTO.builder()
                .fileId(fileEntity.getId())
                .ownerId(fileEntity.getUserId())
                .originalFileName(fileEntity.getOriginalFileName().replace("." +
                        fileEntity.getFileExtension(), "") + "_nocopy_protected" +
                        "." + fileEntity.getFileExtension())
                .fileSize(fileEntity.getFileSize())
                .uploadDate(fileEntity.getCreatedAt())
                .status(fileEntity.getProtectionStatus())
                .build();
    }
}