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

package ru.soune.nocopy.handler;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.vrt.NoCopyFileService;
import com.vrt.fileprotection.FileProtector;
import com.vrt.fileprotection.NoCopyCheckResult;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import ru.soune.nocopy.configuration.search.SearchProperties;
import ru.soune.nocopy.dto.BaseRequest;
import ru.soune.nocopy.dto.BaseResponse;
import ru.soune.nocopy.dto.MessageCode;
import ru.soune.nocopy.dto.file.ImageSearchRequest;
import ru.soune.nocopy.dto.file.YandexSearchResponse;
import ru.soune.nocopy.entity.file.FileEntity;
import ru.soune.nocopy.entity.file.ImageHashEntity;
import ru.soune.nocopy.entity.tokenoperation.OperationType;
import ru.soune.nocopy.exception.NotValidFieldException;
import ru.soune.nocopy.repository.FileEntityRepository;
import ru.soune.nocopy.repository.ImageHashRepository;
import ru.soune.nocopy.repository.SimilarImageProjection;
import ru.soune.nocopy.service.FileSimilarityService;
import ru.soune.nocopy.service.ImageHashService;
import ru.soune.nocopy.service.file.CheckCounterService;
import ru.soune.nocopy.service.file.FileUploadService;
import ru.soune.nocopy.service.search.SearchImageService;
import ru.soune.nocopy.service.tariff.TariffConstants;
import ru.soune.nocopy.service.tariff.TariffInfoService;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.TimeoutException;

@Slf4j
@Component
@RequiredArgsConstructor
public class ImageFoundRequestHandler implements RequestHandler {

    private final ObjectMapper objectMapper;

    private final SearchImageService searchImageService;

    private final CheckCounterService checkCounterService;

    private final FileEntityRepository fileEntityRepository;

    private final TariffInfoService tariffInfoService;

    private final SearchProperties searchProperties;

    private final FileUploadService fileUploadService;

    private final ImageHashService imageHashService;

    private final ImageHashRepository imageHashRepository;

    private final FileSimilarityService fileSimilarityService;

    private final NoCopyFileService noCopyFileService;

    @Override
    public BaseResponse handle(BaseRequest request) throws Exception {
        ImageSearchRequest imageSearchRequest = objectMapper.convertValue(request.getMessageBody(),
                ImageSearchRequest.class);

        String fileId = imageSearchRequest.getFileId();

        FileEntity fileEntity = fileEntityRepository.findById(fileId)
                .orElseThrow(() -> {
                    throw new NotValidFieldException("File not found", new BaseResponse(20007,
                            MessageCode.FILE_NOT_FOUND.getCode(), MessageCode.FILE_NOT_FOUND.getDescription(),
                            Map.of("fileId", fileId)));});

        List<YandexSearchResponse.ImageResult> allYandexImages = new ArrayList<>();
        List<YandexSearchResponse.ImageResult> allGoogleImages = new ArrayList<>();
        boolean hasTimeout = false;

        boolean useGoogle = searchProperties.getEngines().getOrDefault("google",
                new SearchProperties.EngineConfig()).isEnabled();
        boolean useYandex = searchProperties.getEngines().getOrDefault("yandex",
                new SearchProperties.EngineConfig()).isEnabled();

        log.info("Search settings: useGoogle={}, useYandex={}", useGoogle, useYandex);

        if (useGoogle) {
            try {
                String searchResponseGoogle = searchImageService.searchReverseByPublicUrl(fileEntity, "google_lens",
                        "exact_matches");
                log.info("searchResponseGoogle OK");
                allGoogleImages = searchImageService.getAllImagesWithoutPagination(
                        searchResponseGoogle, "exact_matches");
            } catch (TimeoutException e) {
                log.warn("Google search timeout for file {}", fileId);
                hasTimeout = true;
            } catch (IOException e) {
                log.error("Google search failed for file {}", fileId, e);
            }
        } else {
            log.info("Google search is disabled by settings");
        }

        if (useYandex) {
            try {
                String searchResponseYandex = searchImageService.searchReverseByPublicUrl(fileEntity, "yandex_reverse_image",
                        "visual_matches");
                log.info("searchResponseYandex OK");
                allYandexImages = searchImageService.getAllImagesWithoutPagination(
                        searchResponseYandex, "visual_matches");
            } catch (TimeoutException e) {
                log.warn("Yandex search timeout for file {}", fileId);
                hasTimeout = true;
            } catch (IOException e) {
                log.error("Yandex search failed for file {}", fileId, e);
            }
        } else {
            log.info("Yandex search is disabled by settings");
        }

        if (allYandexImages.isEmpty() && allGoogleImages.isEmpty()) {
            log.warn("No results from any search engine");
            String errorMessage = useGoogle || useYandex ? "No results found" : "All search engines are disabled";

            return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                    errorMessage, new YandexSearchResponse());
        }

        List<YandexSearchResponse.ImageResult> allUniqueImages;
        if (useGoogle && useYandex) {
            allUniqueImages = searchImageService.removeDuplicateUrls(allYandexImages, allGoogleImages);
            log.info("Merged results from both engines: {} unique images", allUniqueImages.size());
        } else if (useGoogle) {
            allUniqueImages = allGoogleImages;
            log.info("Results only from Google: {} images", allUniqueImages.size());
        } else {
            allUniqueImages = allYandexImages;
            log.info("Results only from Yandex: {} images", allUniqueImages.size());
        }

        tariffInfoService.writeOffTokens(fileEntity.getUserId(), TariffConstants.TOKEN_VALUE_FOR_SEARCH,
                OperationType.SEARCH);
        checkCounterService.incrementCheckCount(fileEntity.getUserId(), fileEntity.getMimeType());

        int page = imageSearchRequest.getPage() != null ? imageSearchRequest.getPage() : 1;
        int pageSize = 5;

        boolean checkSimilar = imageSearchRequest.getCheckSimilarFirst() != null &&
                imageSearchRequest.getCheckSimilarFirst();

        if (checkSimilar && !allUniqueImages.isEmpty()) {
            YandexSearchResponse.ImageResult image = allUniqueImages.getFirst();
            String url = image.getUrl();
            Path tempFile = null;
            try {
                tempFile = convertToSupportedFormat(fileUploadService.downloadImageFromUrl(url));
                File file = new File(fileEntity.getFilePath());
//                String filePath = tempFile.toAbsolutePath().toString();
                NoCopyCheckResult checkResult = noCopyFileService.checkFile(file,
                        FileProtector.Type.IMAGE);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
                log.info("checkResult: " + checkResult);
//                Map<String, Long> hashSearchApiFile = imageHashService.calculateHash(Paths.get(filePath));
//                ImageHashEntity checkFileHash = imageHashRepository.findByFileId(fileId);

//                boolean hashesAreEqual = hashSearchApiFile.get("hi").equals(checkFileHash.getHash64Hi())
//                        && hashSearchApiFile.get("low").equals(checkFileHash.getHash64Lo());

//                if (!hashesAreEqual) {
//                    return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
//                            "Not found similar pictures", new YandexSearchResponse());
//                }
            } finally {
                if (tempFile != null) {
                    try {
                        Files.deleteIfExists(tempFile);
                    } catch (IOException e) {
                        log.warn("Not delete temp file: {}", tempFile, e);
                    }
                }
            }
        }

        List<YandexSearchResponse.ImageResult> pagedResults = searchImageService.paginateResults(allUniqueImages, page,
                pageSize * 2);

        YandexSearchResponse finalResponse = new YandexSearchResponse();
        finalResponse.setImages(pagedResults);
        finalResponse.setPage(page);
        finalResponse.setPageSize(pageSize * 2);
        finalResponse.setTotalResults(allUniqueImages.size());
        finalResponse.setTotalPages((int) Math.ceil((double) allUniqueImages.size() / (pageSize * 2)));

        String message = hasTimeout ? "Search completed (partial)" : MessageCode.SUCCESS.getDescription();

        return new BaseResponse(request.getMsgId(), MessageCode.SUCCESS.getCode(),
                message, finalResponse);
    }

    public Path convertToSupportedFormat(Path imagePath) throws IOException {
        String fileName = imagePath.getFileName().toString().toLowerCase();

        if (fileName.matches(".*\\.(jpg|jpeg|png|bmp)$")) {
            return imagePath;
        }

        BufferedImage image = ImageIO.read(imagePath.toFile());
        Path pngFile = Files.createTempFile("converted_", ".png");
        ImageIO.write(image, "PNG", pngFile.toFile());

        return pngFile;
    }
}