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.entity.company.Company;
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.entity.user.User;
import ru.soune.nocopy.exception.DuplicateImageException;
import ru.soune.nocopy.repository.*;
import ru.soune.nocopy.util.FileUtil;

import java.io.File;
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> 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, String mimeType, Long userId) throws Exception {
        FileEntity duplicateByHash = findDuplicateByHash(path, mimeType, userId);
        if (duplicateByHash != null) {
            throw new DuplicateImageException("Duplicate", duplicateByHash.getId(),
                    duplicateByHash.getUserId());
        }
    }

    public FileEntity findDuplicateByHash(String path, String mimeType, Long userId) throws Exception {
        String hash = calculateFileHash(path);

        User user = userRepository.findById(userId).orElseThrow();
        Company company = user.getCompany();
        List<Long> userIds = new ArrayList<>();

        if (company != null) {
            userIds.addAll(company.getUsers().stream().map(User::getId).toList());
        } else {
            userIds.add(userId);
        }

        List<FileEntity> fileEntityList = fileEntityRepository.findByMimeTypeAndUserIds(mimeType, userIds);

        for (FileEntity file : fileEntityList) {
            if (file.getProtectedFilePath() == null) continue;
            File existingFile = new File(file.getProtectedFilePath());
            File newFile = new File(path);

            if (existingFile.length() == newFile.length()) {
                if (calculateFileHash(file.getProtectedFilePath()).equals(hash)) {
                    return file;
                }
            }
        }

        return null;
    }

    private final UserRepository userRepository;

    public Page<SimilarFileDTO> findSimilarFiles(
            String fileId,
            List<String> similarityLevels,
            Pageable pageable,
            String authToken) {

        try {
            var imageHashEntity = hashRepository.findById(fileId)
                    .orElseThrow(() -> new RuntimeException("Hash not found for file: " + fileId));


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

            Long userId = authTokenRepository.findUserIdByToken(authToken);
            User user = userRepository.findById(userId).orElseThrow(() -> new RuntimeException("User not found"));
            Company company = user.getCompany();

            if (company != null) {
                userIds = company.getUsers().stream().map(User::getId).collect(Collectors.toList());
            } else {
                userIds = List.of(userId);
            }

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

            List<SimilarFileDTO> allResults = candidates.stream()
                    .map(c -> createSimilarFileResponse(c,
                            imageHashEntity.getHash64Hi(),
                            imageHashEntity.getHash64Lo(), userIds))
                    .filter(dto -> levels.contains(dto.getSimilarityLevel()))
                    .filter(dto -> !(dto.getFileStatus().equals(FileStatus.MODERATION.name()) && !dto.getOwner()))
                    .sorted(Comparator.comparingInt(SimilarFileDTO::getHammingDistance))
                    .collect(Collectors.toList());

            int start = (int) pageable.getOffset();
            int end = Math.min((start + pageable.getPageSize()), allResults.size());
            List<SimilarFileDTO> pageContent = start < allResults.size() ?
                    allResults.subList(start, end) : Collections.emptyList();

            return new PageImpl<>(pageContent, pageable, allResults.size());

        } catch (Exception e) {
            log.error("Error finding similar files for fileId: {}", fileId, e);
            return new PageImpl<>(Collections.emptyList(), pageable, 0);
        }
    }


    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, List<Long> userIds) {
        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";
        }

        boolean owner = userIds.contains(similarImageProjection.getUserId());


        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())
                .owner(owner)
//                .url(baseUrl + "/api/files/protected/" + similarImageProjection.getId())
                .url(baseUrl + "/api/files/protected/" + similarImageProjection.getId() + "/thumbnail")
                .fileStatus(similarImageProjection.getFileStatus())
                .build();
    }

    public 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())
                .fileStatus(fileEntity.getStatus().name())
                .status(fileEntity.getProtectionStatus())
                .supportId(Long.valueOf(fileEntity.getSupportId()))
                .build();
    }
}