images.py

"""
Image retrieval: filename-based candidate search, scoring, ranking, and browser opening.
"""
import os
import logging
import webbrowser
from pathlib import Path

from langchain_core.documents import Document

from config.settings import RAW_IMAGE_PATH
from config.runtime import is_auto_open_image, get_open_image_count
from retrieval.query import (
    normalize_for_match,
    infer_source_hints,
    path_matches_hints,
    source_matches_hints,
    rewrite_query_for_retrieval,
)

logger = logging.getLogger(__name__)

_SUPPORTED_IMAGE_EXT = (".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff")

_IMAGE_STOPWORDS = {
    "show", "me", "the", "a", "an", "of", "for", "please",
    "photo", "image", "picture", "pic",
}


# ---------------------------------------------------------------------------
# Candidate creation
# ---------------------------------------------------------------------------

def create_image_candidate_doc(path: str) -> Document:
    return Document(
        page_content=f"Image candidate by filename: {os.path.basename(path)}",
        metadata={
            "source": os.path.basename(path),
            "source_type": "image",
            "image_path": path,
        },
    )


def apply_image_source_focus(question: str, docs: list[Document]) -> list[Document]:
    """For image queries that name a product, keep only docs from that product."""
    hints = infer_source_hints(question)
    if not hints:
        return docs

    focused: list[Document] = []
    for doc in docs or []:
        source = doc.metadata.get("source", "")
        if doc.metadata.get("source_type") == "image":
            image_path = doc.metadata.get("image_path") or os.path.join(RAW_IMAGE_PATH, source)
            if path_matches_hints(image_path, hints) or source_matches_hints(source, hints):
                focused.append(doc)
        elif source_matches_hints(source, hints):
            focused.append(doc)
    return focused or docs


# ---------------------------------------------------------------------------
# Filename-based candidate search
# ---------------------------------------------------------------------------

def get_image_candidates(question: str, limit: int = 3) -> list[str]:
    """
    Return up to `limit` image file paths whose filenames best match the query.
    Used as a safety net when BM25 hasn't indexed image OCR text.
    """
    if not os.path.isdir(RAW_IMAGE_PATH):
        return []

    query_norm = normalize_for_match(question)
    query_tokens = set(query_norm.split())
    if not query_tokens:
        return []

    scored: list[tuple[int, str, str]] = []
    for filename in os.listdir(RAW_IMAGE_PATH):
        if not filename.lower().endswith(_SUPPORTED_IMAGE_EXT):
            continue
        path = os.path.join(RAW_IMAGE_PATH, filename)
        name_norm = normalize_for_match(os.path.splitext(filename)[0])
        name_tokens = set(name_norm.split())
        if not name_tokens:
            continue

        overlap = len(query_tokens.intersection(name_tokens))
        contains_boost = 1 if (
            " ".join(query_tokens) in name_norm
            or any(t in name_norm for t in query_tokens if len(t) >= 6)
        ) else 0
        score = overlap * 10 + contains_boost
        if score > 0:
            scored.append((score, filename, path))

    scored.sort(key=lambda x: (-x[0], x[1]))
    return [p for _, _, p in scored[:limit]]


# ---------------------------------------------------------------------------
# Path scoring and ranking
# ---------------------------------------------------------------------------

def score_image_path(question: str, image_path: str) -> int:
    """
    Return a numeric score indicating how well this image matches the query.
    A product-hint match adds a large bonus (40 pts) to avoid cross-product confusion.
    """
    query_tokens = [
        t for t in normalize_for_match(question).split()
        if t and t not in _IMAGE_STOPWORDS
    ]
    name_norm = normalize_for_match(os.path.basename(image_path))
    name_tokens = set(name_norm.split())
    name_compact = name_norm.replace(" ", "")

    if not query_tokens or not name_tokens:
        return 0

    score = 0
    for token in query_tokens:
        if token in name_tokens:
            score += 12
        elif token in name_norm:
            score += 8
        elif token in name_compact:
            # Handles "keychain" vs "key chain" spacing difference
            score += 8

    normalized_question = rewrite_query_for_retrieval(question)
    hints = infer_source_hints(normalized_question)
    if path_matches_hints(image_path, hints):
        score += 40
    return score


def rank_image_hits(question: str, image_hits: list[str]) -> list[str]:
    """Sort image paths by relevance: hint match first, then path score, then name."""
    normalized_question = rewrite_query_for_retrieval(question)
    hints = infer_source_hints(normalized_question)
    return sorted(
        image_hits,
        key=lambda p: (
            -int(path_matches_hints(p, hints)),
            -score_image_path(normalized_question, p),
            p.lower(),
        ),
    )


# ---------------------------------------------------------------------------
# Browser opening
# ---------------------------------------------------------------------------

def open_images_in_browser(image_hits: list[str]) -> None:
    if not is_auto_open_image():
        return
    open_count = get_open_image_count()
    for path in image_hits[:open_count]:
        try:
            uri = Path(path).resolve().as_uri()
            webbrowser.open(uri)
            print(f"[INFO] Opened in browser: {path}")
        except Exception as e:
            logger.warning("Cannot open image in browser: %s (%s)", path, e)
            print(f"[WARN] Cannot open image in browser: {path} ({e})")