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.ImageHashEntity;
import ru.soune.nocopy.repository.AuthTokenRepository;
import ru.soune.nocopy.repository.ImageHashRepository;
import ru.soune.nocopy.repository.ImageSimilarityRepository;
import ru.soune.nocopy.repository.SimilarImageProjection;
import ru.soune.nocopy.util.FileUtil;

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;

    @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 Page<SimilarFileDTO> findSimilarFiles(String fileId, SimilarityFilter filter, Pageable pageable, String authToken) {
        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");

        List<SimilarFileDTO> allResults = candidates.stream()
                .map(c -> createSimilarFileResponse(c, hash64Hi, hash64Lo))
                .filter(response -> similarityLevels.contains(response.getSimilarityLevel()))
                .sorted(Comparator.comparingInt(SimilarFileDTO::getHammingDistance))
                .collect(Collectors.toList());

        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 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())
                .fileSize(similarImageProjection.getFileSize())
                .hammingDistance(hamming)
                .similarityLevel(level)
                .supportId(similarImageProjection.getSupportId())
                .uploadDate(similarImageProjection.getUploadDate())
                .status(similarImageProjection.getProtectionStatus())
                .url(baseUrl + "/api/files/protected/" + similarImageProjection.getId())
                .build();
    }
}