aivar/cache.py

"""Perceptual result cache backed by data/memory.db (SQLite).

Goal: re-analyzing the same footage costs ZERO Gemini calls — even when the
file was re-encoded, resized, brightness-tweaked, or slightly trimmed.

How: each input file gets a perceptual signature (pHash of the image, or
pHashes of CACHE_SIG_FRAMES evenly-sampled frames for video). A cached
analysis matches when every query view fuzzy-matches a distinct cached view
(mean Hamming distance <= CACHE_PHASH_THRESHOLD). Exact byte equality is
never required — that is the whole point.
"""
import json
import sqlite3
import time
from pathlib import Path

import cv2
import imagehash
from PIL import Image

from .config import CACHE_DB, CACHE_PHASH_THRESHOLD, CACHE_SIG_FRAMES
from .ingest import IMAGE_EXT, VIDEO_EXT
from .schemas import AnalysisResult

SCHEMA = """
CREATE TABLE IF NOT EXISTS analyses(
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    created_at REAL NOT NULL,
    n_views INTEGER NOT NULL,
    signature TEXT NOT NULL,      -- JSON: list of per-view signatures
    result TEXT NOT NULL          -- JSON: AnalysisResult
);
"""


def connect() -> sqlite3.Connection:
    CACHE_DB.parent.mkdir(parents=True, exist_ok=True)
    con = sqlite3.connect(CACHE_DB)
    con.executescript(SCHEMA)
    return con


# ---------------- signatures ----------------

def _hash_frame(frame) -> str:
    img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    return str(imagehash.phash(img))


def file_signature(path: str) -> dict:
    """{'kind': 'image'|'video', 'hashes': [hex phash, ...]}"""
    ext = Path(path).suffix.lower()
    if ext in IMAGE_EXT:
        frame = cv2.imread(path)
        if frame is None:
            raise ValueError(f"unreadable image: {path}")
        return {"kind": "image", "hashes": [_hash_frame(frame)]}
    if ext in VIDEO_EXT:
        cap = cv2.VideoCapture(path)
        n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1
        hashes = []
        for k in range(CACHE_SIG_FRAMES):
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(k * n / CACHE_SIG_FRAMES))
            ok, frame = cap.read()
            if ok:
                hashes.append(_hash_frame(frame))
        cap.release()
        if not hashes:
            raise ValueError(f"unreadable video: {path}")
        return {"kind": "video", "hashes": hashes}
    raise ValueError(f"unsupported file type: {ext}")


def inputs_signature(paths: list[str]) -> list[dict]:
    return [file_signature(p) for p in paths]


# ---------------- fuzzy matching ----------------

def _dist(a: str, b: str) -> int:
    return imagehash.hex_to_hash(a) - imagehash.hex_to_hash(b)


def _views_match(q: dict, c: dict) -> bool:
    if q["kind"] != c["kind"]:
        return False
    qh, ch = q["hashes"], c["hashes"]
    n = min(len(qh), len(ch))
    if n == 0:
        return False
    # aligned comparison tolerates trims; mean distance tolerates re-encoding
    mean = sum(_dist(qh[i], ch[i]) for i in range(n)) / n
    return mean <= CACHE_PHASH_THRESHOLD


def _all_match(query: list[dict], cached: list[dict]) -> bool:
    """Every query view must match a distinct cached view (order-free)."""
    if len(query) != len(cached):
        return False
    remaining = list(range(len(cached)))
    for q in query:
        hit = next((j for j in remaining if _views_match(q, cached[j])), None)
        if hit is None:
            return False
        remaining.remove(hit)
    return True


# ---------------- public API ----------------

def lookup(paths: list[str]) -> AnalysisResult | None:
    sig = inputs_signature(paths)
    con = connect()
    try:
        rows = con.execute(
            "SELECT signature, result FROM analyses WHERE n_views=? ORDER BY id DESC",
            (len(sig),)).fetchall()
        for row_sig, row_result in rows:
            if _all_match(sig, json.loads(row_sig)):
                result = AnalysisResult.model_validate_json(row_result)
                result.from_cache = True
                result.gemini_calls = 0
                return result
        return None
    finally:
        con.close()


def store(paths: list[str], result: AnalysisResult) -> None:
    sig = inputs_signature(paths)
    con = connect()
    try:
        with con:
            con.execute(
                "INSERT INTO analyses(created_at, n_views, signature, result) VALUES(?,?,?,?)",
                (time.time(), len(sig), json.dumps(sig), result.model_dump_json()))
    finally:
        con.close()