find-image

import { existsSync, readFileSync, statSync } from "fs";
import crypto from "crypto";
import path from "path";
import type { DrawThingsSearchResult, GenerationMatch, GenerationMetadata, ImageResult, IndexedGeneration, MatchType } from "../types.js";
import { MultimodalEmbeddingStore } from "../embeddings/index.js";
import { searchMultimodalEmbeddings, type MultimodalQueryMode } from "./multimodalSearch.js";
import { describeFilterMatches, matchesFilters, parseStructuredQuery, type StructuredFilterTolerances } from "./queryParser.js";
import type { DtcModelMappingSnapshotV1 } from "../helpers/dtcModelMappingSnapshot.js";
import { SUPPORTED_IMAGE_FORMAT_SET } from "../documents/imageFormats.js";

export interface SearchOptions {
  maxResults?: number;
  snapshot?: DtcModelMappingSnapshotV1;
  includeMultimodalSearch?: boolean;
  multimodalQueryEmbedding?: number[];
  multimodalStore?: MultimodalEmbeddingStore;
  multimodalModel?: string;
  multimodalDimension?: number;
  minMultimodalScore?: number;
  embeddingDisplayMultiplier?: number;
  multimodalQueryMode?: MultimodalQueryMode;
  contentOnly?: boolean;
  excludeImageRefs?: string[];
  excludeContentHashes?: string[];
  multimodalExcludeImageRefs?: string[];
  multimodalExcludeContentHashes?: string[];
  disableStructuredFilters?: boolean;
  targetGenerationMetadata?: GenerationMetadata;
  targetContentHash?: string;
  metadataExactMatches?: boolean;
  allowByteIdenticalResults?: boolean;
  structuredFilterTolerances?: StructuredFilterTolerances;
}

const DEFAULT_OPTIONS = {
  maxResults: Number.POSITIVE_INFINITY,
  includeMultimodalSearch: false,
  minMultimodalScore: 50,
  embeddingDisplayMultiplier: 1,
  disableStructuredFilters: false,
  allowByteIdenticalResults: false,
};
function normalizedRef(ref: string): string {
  return ref.startsWith("project://") ? ref : path.resolve(ref);
}

function availableImages(imagePaths: string[] = []): string[] {
  return imagePaths.filter((imagePath) => imagePath.startsWith("project://") || existsSync(imagePath));
}

function toMatch(generation: GenerationMetadata, score: number, matchType: MatchType, matchedTerms?: string[]): GenerationMatch {
  return {
    prompt: generation.prompt,
    latitude: generation.latitude,
    longitude: generation.longitude,
    lensModel: generation.lensModel,
    exposureTime: generation.exposureTime,
    fNumber: generation.fNumber,
    iso: generation.iso,
    exposureCompensation: generation.exposureCompensation,
    focalLength: generation.focalLength,
    focalLength35mm: generation.focalLength35mm,
    exposureProgram: generation.exposureProgram,
    meteringMode: generation.meteringMode,
    whiteBalance: generation.whiteBalance,
    flash: generation.flash,
    orientation: generation.orientation,
    exposureMode: generation.exposureMode,
    negativePrompt: generation.negativePrompt,
    model: generation.model,
    loras: generation.loras?.map((lora) => typeof lora === "string" ? lora : lora.model),
    tags: generation.tags,
    sampler: generation.sampler,
    seed: generation.seed,
    seedMode: generation.seedMode,
    shift: generation.shift,
    strength: generation.strength,
    mode: generation.mode,
    steps: generation.steps,
    cfgScale: generation.cfgScale,
    width: generation.width,
    height: generation.height,
    imagePaths: availableImages(generation.imagePaths),
    contentHash: generation.contentHash,
    httpPreviewUrls: generation.httpPreviewUrls,
    sourceInfo: generation.sourceInfo,
    timestamp: generation.timestamp,
    matchScore: score,
    matchType,
    matchedTerms,
    exactScore: matchType === "exact" ? score : undefined,
    numFrames: generation.numFrames && generation.numFrames > 1 ? generation.numFrames : undefined,
  };
}

function excludeImages(matches: GenerationMatch[], refs: string[] = []): GenerationMatch[] {
  const excluded = new Set(refs.filter(Boolean).map(normalizedRef));
  return matches.flatMap((match) => {
    const imagePaths = match.imagePaths.filter((imagePath) => !excluded.has(normalizedRef(imagePath)));
    return imagePaths.length > 0 ? [{ ...match, imagePaths }] : [];
  });
}

function excludeContentHashes(matches: GenerationMatch[], hashes: string[] = []): GenerationMatch[] {
  const excluded = new Set(hashes.filter(Boolean));
  if (excluded.size === 0) return matches;
  return matches.filter((match) => !excluded.has(match.contentHash ?? ""));
}

function contentHash(imagePath: string): string | null {
  try {
    if (imagePath.startsWith("project://") || !existsSync(imagePath)) return null;
    return crypto.createHash("sha256").update(readFileSync(imagePath)).digest("hex");
  } catch {
    return null;
  }
}

function resultRecency(match: GenerationMatch): number {
  const imagePath = match.imagePaths.length === 1 ? match.imagePaths[0] : undefined;
  if (!imagePath) return Number.NEGATIVE_INFINITY;
  if (imagePath.startsWith("project://")) {
    const thumbnailId = Number(imagePath.slice(imagePath.lastIndexOf("#") + 1));
    return Number.isSafeInteger(thumbnailId) ? thumbnailId : Number.NEGATIVE_INFINITY;
  }
  try {
    return existsSync(imagePath) ? statSync(imagePath).mtimeMs : Number.NEGATIVE_INFINITY;
  } catch {
    return Number.NEGATIVE_INFINITY;
  }
}

function dropByteDuplicates(matches: GenerationMatch[]): GenerationMatch[] {
  const byHash = new Map<string, GenerationMatch>();
  const passthrough: GenerationMatch[] = [];
  for (const match of matches) {
    const hash = match.contentHash;
    if (!hash) {
      passthrough.push(match);
      continue;
    }
    const previous = byHash.get(hash);
    if (!previous
      || resultRecency(match) > resultRecency(previous)) {
      byHash.set(hash, match);
    }
  }
  return [...passthrough, ...byHash.values()];
}

function exactContentHashMatches(generations: IndexedGeneration[], targetHash: string): GenerationMatch[] {
  return generations.flatMap((generation) => {
    return generation.contentHash === targetHash
      ? [toMatch(generation, 100, "exact", ["byte-identical"])]
      : [];
  });
}

function uniqueMatchesByImagePaths(matches: GenerationMatch[]): GenerationMatch[] {
  const seen = new Set<string>();
  return matches.filter((match) => {
    const key = match.imagePaths.map(normalizedRef).sort().join("\u0000");
    if (seen.has(key)) return false;
    seen.add(key);
    return true;
  });
}

function normalizeFilename(value: string): string {
  return path.basename(value.replace(/\\/g, "/").replace(/^['"]|['"]$/g, "")).trim().toLowerCase().normalize("NFKC");
}

function isImageFilenameQuery(query: string): boolean {
  const filename = normalizeFilename(query);
  return filename.includes(".") && SUPPORTED_IMAGE_FORMAT_SET.has(filename.split(".").pop() ?? "");
}

function isProjectFilenameQuery(query: string): boolean {
  return normalizeFilename(query).endsWith(".sqlite3");
}

function sameText(left?: string, right?: string): boolean {
  return (left ?? "").trim() === (right ?? "").trim();
}

function sameNumber(left?: number, right?: number): boolean {
  return (left ?? null) === (right ?? null);
}

function sameLoras(left?: GenerationMetadata["loras"], right?: GenerationMetadata["loras"]): boolean {
  const normalize = (loras?: GenerationMetadata["loras"]) => (loras ?? []).map((lora) => typeof lora === "string" ? lora : lora.model).map((lora) => lora.trim().toLowerCase()).sort();
  const a = normalize(left);
  const b = normalize(right);
  return a.length === b.length && a.every((value, index) => value === b[index]);
}

function sameMetadata(generation: GenerationMetadata, target: GenerationMetadata): boolean {
  return sameText(generation.prompt, target.prompt)
    && sameText(generation.negativePrompt, target.negativePrompt)
    && sameText(generation.model, target.model)
    && sameLoras(generation.loras, target.loras)
    && sameText(generation.sampler, target.sampler)
    && sameNumber(generation.steps, target.steps)
    && sameNumber(generation.cfgScale, target.cfgScale)
    && sameNumber(generation.seed, target.seed)
    && sameNumber(generation.width, target.width)
    && sameNumber(generation.height, target.height);
}

function recallMatches(generations: IndexedGeneration[], opts: SearchOptions): GenerationMatch[] {
  if (!opts.includeMultimodalSearch || !opts.multimodalQueryEmbedding || !opts.multimodalStore || !opts.multimodalModel || !opts.multimodalDimension) return [];
  const excluded = new Set(
    (opts.allowByteIdenticalResults ? [] : [...(opts.excludeImageRefs ?? []), ...(opts.multimodalExcludeImageRefs ?? [])])
      .filter(Boolean)
      .map(normalizedRef),
  );
  const excludedContentHashes = new Set(
    opts.allowByteIdenticalResults ? [] : (opts.multimodalExcludeContentHashes ?? []).filter(Boolean),
  );
  const embeddingDisplayMultiplier = opts.embeddingDisplayMultiplier ?? DEFAULT_OPTIONS.embeddingDisplayMultiplier;
  const hits = searchMultimodalEmbeddings(opts.multimodalQueryEmbedding, opts.multimodalStore, opts.multimodalModel, opts.multimodalDimension, {
    maxResults: Number.POSITIVE_INFINITY,
    minScore: opts.minMultimodalScore ?? DEFAULT_OPTIONS.minMultimodalScore,
  });
  const generationsByImage = new Map<string, IndexedGeneration>();
  for (const generation of generations) for (const imagePath of generation.imagePaths ?? []) generationsByImage.set(imagePath, generation);
  let excludedByReference = 0;
  let excludedByContentHash = 0;
  let missingGeneration = 0;
  let unavailableImage = 0;
  const matches: GenerationMatch[] = [];
  for (const hit of hits) {
    if (excluded.has(normalizedRef(hit.entry.imageRef))) {
      excludedByReference++;
      continue;
    }
    if (excludedContentHashes.has(contentHash(hit.entry.imageRef) ?? "")) {
      excludedByContentHash++;
      continue;
    }
    const generation = generationsByImage.get(hit.entry.imageRef);
    if (!generation) {
      missingGeneration++;
      continue;
    }
    if (availableImages(generation.imagePaths).length === 0) {
      unavailableImage++;
      continue;
    }
    const displayScore = Math.min(100, Math.round(hit.score * embeddingDisplayMultiplier));
    matches.push({
      ...toMatch(generation, displayScore, "hybrid", ["qwen_multimodal"]),
      multimodalImageRef: hit.entry.imageRef,
      multimodalScore: displayScore,
      visualScore: displayScore,
      rankingSimilarity: hit.similarity,
    });
  }
  console.info(
    `[Multimodal recall] model=${opts.multimodalModel} dimension=${opts.multimodalDimension} ` +
    `thresholdHits=${hits.length} returned=${matches.length} excludedByReference=${excludedByReference} ` +
    `excludedByContentHash=${excludedByContentHash} missingGeneration=${missingGeneration} unavailableImage=${unavailableImage}`,
  );
  return matches;
}

function toImageResults(matches: GenerationMatch[]): ImageResult[] {
  return matches.flatMap((match) => match.imagePaths.map((imagePath, index) => ({
    path: imagePath,
    httpPreviewUrl: match.httpPreviewUrls?.[index],
    prompt: match.prompt,
    model: match.model,
    matchType: match.matchType,
    score: match.matchScore,
    timestamp: match.timestamp,
  })));
}

function makeResult(query: string, startedAt: number, exactMatches: GenerationMatch[], contentMatches: GenerationMatch[]): DrawThingsSearchResult {
  const all = [...exactMatches, ...contentMatches];
  return {
    query,
    totalFound: all.length,
    searchTimeMs: Date.now() - startedAt,
    exactMatches,
    semanticMatches: contentMatches,
    imageResults: toImageResults(all),
    semanticSearchEnabled: false,
    multimodalSearchEnabled: contentMatches.length > 0,
    summary: all.length === 0 ? `No generations found matching "${query}".` : `Found ${all.length} generation${all.length === 1 ? "" : "s"}.`,
  };
}

export async function searchGenerations(query: string, generations: IndexedGeneration[], options: SearchOptions = {}): Promise<DrawThingsSearchResult> {
  const startedAt = Date.now();
  const opts = { ...DEFAULT_OPTIONS, ...options };
  const excludedImageRefs = opts.allowByteIdenticalResults ? [] : opts.excludeImageRefs;
  const knownProjects = [...new Set(generations.flatMap((generation) =>
    generation.sourceInfo?.type === "draw_things_project" ? [generation.sourceInfo.projectFile] : [],
  ))];
  const parsed = opts.disableStructuredFilters ? { promptQuery: query, filters: {}, hasFilters: false } : parseStructuredQuery(query, opts.snapshot, knownProjects, opts.structuredFilterTolerances);
  const filtered = parsed.hasFilters ? generations.filter((generation) => matchesFilters(generation, parsed.filters, opts.snapshot)) : generations;

  if (isProjectFilenameQuery(query)) {
    const filename = normalizeFilename(query);
    const matches = filtered
      .filter((generation) => generation.sourceInfo?.type === "draw_things_project" && normalizeFilename(generation.sourceInfo.projectFile) === filename)
      .sort((left, right) => (left.sourceInfo?.type === "draw_things_project" ? left.sourceInfo.logicalTime ?? 0 : 0) - (right.sourceInfo?.type === "draw_things_project" ? right.sourceInfo.logicalTime ?? 0 : 0))
      .map((generation) => toMatch(generation, 100, "exact", [filename]));
    const candidates = excludeImages(matches, excludedImageRefs);
    const deduplicated = opts.allowByteIdenticalResults ? candidates : dropByteDuplicates(candidates);
    return makeResult(query, startedAt, deduplicated.slice(0, opts.maxResults), []);
  }

  if (isImageFilenameQuery(query)) {
    const filename = normalizeFilename(query);
    const matches = filtered.flatMap((generation) => {
      const imageMatch = (generation.imagePaths ?? []).some((imagePath) => normalizeFilename(imagePath) === filename);
      const attachmentMatch = generation.sourceInfo?.type === "attachment" && normalizeFilename(generation.sourceInfo.originalName) === filename;
      return imageMatch || attachmentMatch ? [toMatch(generation, 100, "exact", [filename])] : [];
    });
    const candidates = excludeImages(matches, excludedImageRefs);
    const deduplicated = opts.allowByteIdenticalResults ? candidates : dropByteDuplicates(candidates);
    return makeResult(query, startedAt, deduplicated.slice(0, opts.maxResults), []);
  }

  const contentHashMatches = opts.targetContentHash
    ? excludeImages(exactContentHashMatches(filtered, opts.targetContentHash), excludedImageRefs)
    : [];
  const metadataMatches = opts.metadataExactMatches && opts.targetGenerationMetadata
    ? excludeContentHashes(
      excludeImages(filtered.filter((generation) => sameMetadata(generation, opts.targetGenerationMetadata!)).map((generation) => toMatch(generation, 100, "exact", ["metadata"])), excludedImageRefs),
      opts.allowByteIdenticalResults ? [] : opts.excludeContentHashes,
    )
    : parsed.hasFilters && !parsed.promptQuery.trim() && !opts.multimodalQueryEmbedding
      ? excludeImages(filtered.map((generation) => toMatch(generation, 100, "exact", describeFilterMatches(parsed.filters))), excludedImageRefs)
      : [];
  const exactMatches = uniqueMatchesByImagePaths([...contentHashMatches, ...metadataMatches]);
  const deduplicatedExactMatches = opts.allowByteIdenticalResults ? exactMatches : dropByteDuplicates(exactMatches);
  const contentCandidates = excludeImages(
    recallMatches(filtered, opts),
    [...(excludedImageRefs ?? []), ...deduplicatedExactMatches.flatMap((match) => match.imagePaths)],
  );
  const contentMatches = (opts.allowByteIdenticalResults ? contentCandidates : dropByteDuplicates(contentCandidates))
    .sort((left, right) => (right.rankingSimilarity ?? right.matchScore / 100) - (left.rankingSimilarity ?? left.matchScore / 100) || right.matchScore - left.matchScore)
    .slice(0, Math.max(0, opts.maxResults - deduplicatedExactMatches.length));
  return makeResult(query, startedAt, deduplicatedExactMatches.slice(0, opts.maxResults), contentMatches);
}