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 com.vrt.fileprotection.documents.DocumentCheckResult; import com.vrt.fileprotection.image.ImageCheckResult; 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; import static java.awt.SystemColor.info; @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 allYandexImages = new ArrayList<>(); List 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 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 = tempFile.toFile(); NoCopyCheckResult checkResult = noCopyFileService.checkFile(file, FileProtector.Type.IMAGE); if (checkResult instanceof ImageCheckResult.Success success) { double watermarkCorrelation = success.getWatermarkCorrelation(); log.info("watermarkCorrelation: " + watermarkCorrelation); if (watermarkCorrelation < 0.5) { 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 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; } }