Image Read Cache

#!/usr/bin/env python3
"""Shared utilities for the image-read-cache skill."""
from __future__ import annotations

import hashlib
import os
import struct
import tempfile
import zlib
from pathlib import Path
from typing import Optional

# Unique marker to distinguish our XMP from other XMP data
AI_CACHE_MARKER = "x-ai-cache-v1"
AI_CACHE_MARKER_BYTES = AI_CACHE_MARKER.encode("utf-8")


# ---------------------------------------------------------------------------
# Format detection
# ---------------------------------------------------------------------------

def detect_format(data: bytes, filepath: str | None = None) -> str:
    """Detect image format from magic bytes."""
    if data[:2] == b"\xff\xd8":
        return "jpeg"
    if data[:8] == b"\x89PNG\r\n\x1a\n":
        return "png"
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "webp"
    if data[:6] in (b"GIF87a", b"GIF89a"):
        return "gif"
    if data[:2] == b"BM":
        return "bmp"
    if filepath:
        ext = Path(filepath).suffix.lower()
        return {
            ".jpg": "jpeg", ".jpeg": "jpeg",
            ".png": "png", ".webp": "webp",
            ".gif": "gif", ".bmp": "bmp",
        }.get(ext, "unknown")
    return "unknown"


# ---------------------------------------------------------------------------
# Strip AI cache XMP (for stable hashing)
# ---------------------------------------------------------------------------

def strip_ai_xmp(data: bytes, fmt: str) -> bytes:
    """Remove our AI cache XMP from file bytes for hashing."""
    if AI_CACHE_MARKER_BYTES not in data:
        return data

    if fmt == "jpeg":
        return _strip_jpeg_xmp(data)
    elif fmt == "png":
        return _strip_png_xmp(data)
    elif fmt == "webp":
        return _strip_webp_xmp(data)
    else:
        return _strip_generic_xmp(data)


def _strip_jpeg_xmp(data: bytes) -> bytes:
    XMP_NS = b"http://ns.adobe.com/xap/1.0/\x00"
    pos = 2
    while pos < len(data) - 1:
        if data[pos] != 0xFF:
            break
        m = data[pos + 1]
        if m in (0xDA, 0xD9):
            break
        if 0xD0 <= m <= 0xD7:
            pos += 2
            continue
        if pos + 4 > len(data):
            break
        seg_len = int.from_bytes(data[pos + 2 : pos + 4], "big")
        seg_end = pos + 2 + seg_len
        if m == 0xE1:
            seg_body = data[pos + 4 : seg_end]
            if seg_body.startswith(XMP_NS) and AI_CACHE_MARKER_BYTES in seg_body:
                return data[:pos] + data[seg_end:]
        pos = seg_end
    return data


def _strip_png_xmp(data: bytes) -> bytes:
    pos = 8
    while pos + 8 <= len(data):
        chunk_len = struct.unpack(">I", data[pos : pos + 4])[0]
        chunk_type = data[pos + 4 : pos + 8]
        chunk_end = pos + 12 + chunk_len
        if chunk_type == b"iTXt":
            chunk_data = data[pos + 8 : pos + 8 + chunk_len]
            if AI_CACHE_MARKER_BYTES in chunk_data:
                return data[:pos] + data[chunk_end:]
        pos = chunk_end
    return data


def _strip_webp_xmp(data: bytes) -> bytes:
    pos = 12
    while pos + 8 <= len(data):
        fourcc = data[pos : pos + 4]
        chunk_size = struct.unpack("<I", data[pos + 4 : pos + 8])[0]
        chunk_end = pos + 8 + chunk_size
        if chunk_size % 2 != 0:
            chunk_end += 1
        if fourcc in (b"XMP ", b"XMP\x00"):
            chunk_data = data[pos + 8 : pos + 8 + chunk_size]
            if AI_CACHE_MARKER_BYTES in chunk_data:
                new_data = data[:pos] + data[chunk_end:]
                new_riff_size = len(new_data) - 8
                new_data = new_data[:4] + struct.pack("<I", new_riff_size) + new_data[8:]
                return new_data
        pos = chunk_end
    return data


def _strip_generic_xmp(data: bytes) -> bytes:
    start_tag = b'<?xpacket begin="\xef\xbb\xbf"'
    end_tag = b'<?xpacket end="w"?>'
    search_from = 0
    while True:
        pkt_start = data.find(start_tag, search_from)
        if pkt_start == -1:
            break
        pkt_end = data.find(end_tag, pkt_start)
        if pkt_end == -1:
            break
        pkt_end += len(end_tag)
        if AI_CACHE_MARKER_BYTES in data[pkt_start:pkt_end]:
            return data[:pkt_start] + data[pkt_end:]
        search_from = pkt_end
    return data


# ---------------------------------------------------------------------------
# Hashing
# ---------------------------------------------------------------------------

def image_hash(data: bytes, fmt: str) -> str:
    """Hash of file bytes with AI cache XMP stripped."""
    clean = strip_ai_xmp(data, fmt)
    return hashlib.sha256(clean).hexdigest()[:16]


# ---------------------------------------------------------------------------
# XML helpers
# ---------------------------------------------------------------------------

def escape_xml(text: str) -> str:
    """Escape text for safe embedding in XML."""
    return (
        text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


def unescape_xml(text: str) -> str:
    """Reverse XML escaping."""
    return (
        text.replace("&lt;", "<")
        .replace("&gt;", ">")
        .replace("&quot;", '"')
        .replace("&amp;", "&")
    )


# ---------------------------------------------------------------------------
# XMP packet builder
# ---------------------------------------------------------------------------

def build_xmp_packet(description: str) -> str:
    """Build a complete XMP packet with our AI cache marker."""
    return (
        '<?xpacket begin="\xef\xbb\xbf" id="W5M0MpCehiHzreSzNTczkc9d"?>'
        '<x:xmpmeta xmlns:x="adobe:ns:meta/">'
        '<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">'
        '<rdf:Description rdf:about="" '
        'xmlns:dc="http://purl.org/dc/elements/1.1/" '
        f'xmlns:ai="{AI_CACHE_MARKER}">'
        "<dc:description><rdf:Alt>"
        f'<rdf:li xml:lang="x-default">{escape_xml(description)}</rdf:li>'
        "</rdf:Alt></dc:description>"
        "</rdf:Description></rdf:RDF></x:xmpmeta>"
        '<?xpacket end="w"?>'
    )


# ---------------------------------------------------------------------------
# XMP extraction / validation
# ---------------------------------------------------------------------------

def extract_xmp_description(data: bytes) -> Optional[str]:
    """Extract dc:description from our AI cache XMP packet."""
    if AI_CACHE_MARKER_BYTES not in data:
        return None

    marker_pos = data.find(AI_CACHE_MARKER_BYTES)
    xmp_start = data.rfind(b"<x:xmpmeta", 0, marker_pos)
    if xmp_start == -1:
        return None

    xmp_end = data.find(b"</x:xmpmeta>", marker_pos)
    if xmp_end == -1:
        return None

    xmp = data[xmp_start : xmp_end + len(b"</x:xmpmeta>")].decode(
        "utf-8", errors="ignore"
    )
    start = '<rdf:li xml:lang="x-default">'
    end = "</rdf:li>"
    start_idx = xmp.find(start)
    if start_idx == -1:
        return None
    start_idx += len(start)
    end_idx = xmp.find(end, start_idx)
    if end_idx == -1:
        return None

    content = xmp[start_idx:end_idx].strip()
    return content if content else None


def validate_and_extract(cached: str, data: bytes, fmt: str) -> Optional[str]:
    """Validate hash and extract content from cached string."""
    if not cached.startswith("HASH:") or "|" not in cached:
        return None

    stored_hash, content = cached.split("|", 1)
    if stored_hash[5:] != image_hash(data, fmt):
        return None
    return content


# ---------------------------------------------------------------------------
# Atomic writes
# ---------------------------------------------------------------------------

def atomic_write_bytes(path: str, data: bytes) -> None:
    """Atomically replace a file with bytes written in the same directory."""
    directory = str(Path(path).resolve().parent)
    fd, temp_path = tempfile.mkstemp(prefix=".tmp-ai-cache-", dir=directory)
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temp_path, path)
    except Exception:
        try:
            os.unlink(temp_path)
        except FileNotFoundError:
            pass
        raise


def atomic_write_text(path: str, text: str) -> None:
    """Atomically replace a file with UTF-8 text."""
    atomic_write_bytes(path, text.encode("utf-8"))


# ---------------------------------------------------------------------------
# Chunk builders
# ---------------------------------------------------------------------------

def build_png_itxt_chunk(description: str) -> bytes:
    """Build a PNG iTXt chunk containing our XMP packet."""
    keyword = b"XML:com.adobe.xmp\x00"
    compression_flag = b"\x00"
    compression_method = b"\x00"
    language_tag = b"\x00"
    translated_keyword = b"\x00"
    xmp_text = build_xmp_packet(description).encode("utf-8")
    chunk_data = (
        keyword
        + compression_flag
        + compression_method
        + language_tag
        + translated_keyword
        + xmp_text
    )
    chunk_type = b"iTXt"
    chunk_crc = struct.pack(">I", zlib.crc32(chunk_type + chunk_data) & 0xFFFFFFFF)
    return struct.pack(">I", len(chunk_data)) + chunk_type + chunk_data + chunk_crc

#!/usr/bin/env python3
"""
Check for cached LLM content in image XMP metadata or sidecar file.

Usage: check_cache.py <image-path>

Output:
  CACHED: <content>   — cached content found and still valid
  NO_CACHE             — no cache, or image has changed since caching

Format support:
  JPEG (.jpg/.jpeg) — reads XMP from APP1 segment
  PNG  (.png)       — reads XMP from iTXt chunk
  WebP (.webp)      — reads XMP from RIFF XMP chunk
  GIF  (.gif)       — reads from .ai-cache sidecar file
  BMP  (.bmp)       — reads from .ai-cache sidecar file

Invalidation: The cache stores an image hash computed over the file
bytes with the AI cache XMP stripped out. If the image is modified
(new pixels, re-export, new screenshot), the hash won't match and
the cache is invalidated automatically.
"""
from __future__ import annotations

import sys
from pathlib import Path
from typing import Optional

from cache_common import (
    detect_format,
    extract_xmp_description,
    unescape_xml,
    validate_and_extract,
)


# ---------------------------------------------------------------------------
# Sidecar fallback (GIF, BMP)
# ---------------------------------------------------------------------------

def check_sidecar(filepath: str) -> Optional[str]:
    """Check for a .ai-cache sidecar file."""
    sidecar = Path(filepath + ".ai-cache")
    try:
        content = sidecar.read_text(encoding="utf-8").strip()
    except FileNotFoundError:
        return None
    return content or None


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    if len(sys.argv) < 2:
        print("Usage: check_cache.py <image-path>", file=sys.stderr)
        sys.exit(1)

    filepath = sys.argv[1]
    try:
        data = Path(filepath).read_bytes()
    except FileNotFoundError:
        print(f"Error: file not found: {filepath}", file=sys.stderr)
        sys.exit(1)
    except OSError as exc:
        print(f"Error: failed to read {filepath}: {exc}", file=sys.stderr)
        sys.exit(1)

    fmt = detect_format(data, filepath)

    # GIF/BMP never store embedded XMP in this implementation.
    if fmt not in {"gif", "bmp"}:
        cached = extract_xmp_description(data)
        if cached:
            content = validate_and_extract(cached, data, fmt)
            if content:
                print(f"CACHED: {unescape_xml(content)}")
                return

    cached = check_sidecar(filepath)
    if cached:
        content = validate_and_extract(cached, data, fmt)
        if content:
            print(f"CACHED: {content}")
            return

    print("NO_CACHE")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Write LLM content as XMP metadata into an image file.

Usage: write_cache.py <image-path> <content>

Stores content in XMP dc:description with a file hash prefix for
cache invalidation. The hash is computed over the file bytes with
any existing AI cache XMP stripped, so the hash stays stable across
cache rewrites.

Format support:
  JPEG (.jpg/.jpeg) — XMP injected as APP1 segment
  PNG  (.png)       — XMP injected as iTXt chunk
  WebP (.webp)      — XMP injected as XMP RIFF chunk
  GIF  (.gif)       — sidecar file (no XMP support)
  BMP  (.bmp)       — sidecar file (no XMP support)

Tries exiftool first (all formats), then direct byte injection,
then sidecar file as last resort.

Output:
  OK: <method>
  SIDECAR: <path>
"""
from __future__ import annotations

import struct
import sys
import subprocess
import shutil
import tempfile
from pathlib import Path

from cache_common import (
    atomic_write_bytes,
    atomic_write_text,
    build_png_itxt_chunk,
    build_xmp_packet,
    detect_format,
    image_hash,
    strip_ai_xmp,
)


# ---------------------------------------------------------------------------
# Writers: exiftool -> format-specific -> sidecar
# ---------------------------------------------------------------------------

def write_with_exiftool(filepath: str, description: str) -> bool:
    """Write XMP via exiftool (handles all formats)."""
    if not shutil.which("exiftool"):
        return False

    xmp_packet = build_xmp_packet(description)
    temp_path = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            prefix=".tmp-ai-cache-",
            suffix=".xmp",
            dir=str(Path(filepath).resolve().parent),
            delete=False,
        ) as handle:
            handle.write(xmp_packet)
            handle.flush()
            temp_path = handle.name
        subprocess.run(
            ["exiftool", "-overwrite_original", f"-XMP<={temp_path}", "--", filepath],
            capture_output=True, check=True, timeout=30,
        )
        return True
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError):
        return False
    finally:
        if temp_path is not None:
            Path(temp_path).unlink(missing_ok=True)


def write_xmp_jpeg(filepath: str, data: bytes, description: str) -> bool:
    """Inject XMP APP1 segment into JPEG."""
    xmp_payload = build_xmp_packet(description).encode("utf-8")
    XMP_NS = b"http://ns.adobe.com/xap/1.0/\x00"

    if len(data) < 2 or data[:2] != b"\xff\xd8":
        return False

    # Strip any existing AI cache XMP
    data = strip_ai_xmp(data, "jpeg")

    # Build new JPEG with XMP APP1 inserted after existing APP0/APP1 segments
    output = bytearray(data[:2])  # SOI
    pos = 2
    xmp_inserted = False

    while pos < len(data) - 1:
        if data[pos] != 0xFF:
            if not xmp_inserted:
                seg = XMP_NS + xmp_payload
                output += b"\xff\xe1" + (len(seg) + 2).to_bytes(2, "big") + seg
                xmp_inserted = True
            output += data[pos:]
            break

        m = data[pos + 1]

        if m in (0xDA, 0xD9):  # SOS or EOI
            if not xmp_inserted:
                seg = XMP_NS + xmp_payload
                output += b"\xff\xe1" + (len(seg) + 2).to_bytes(2, "big") + seg
                xmp_inserted = True
            output += data[pos:]
            break

        if 0xD0 <= m <= 0xD7:  # RST markers
            output += data[pos : pos + 2]
            pos += 2
            continue

        if pos + 4 > len(data):
            output += data[pos:]
            break

        seg_len = int.from_bytes(data[pos + 2 : pos + 4], "big")
        seg_end = pos + 2 + seg_len
        output += data[pos:seg_end]
        pos = seg_end

    if not xmp_inserted:
        seg = XMP_NS + xmp_payload
        rest = bytes(output[2:])
        output = bytearray(data[:2])
        output += b"\xff\xe1" + (len(seg) + 2).to_bytes(2, "big") + seg
        output += rest

    atomic_write_bytes(filepath, bytes(output))
    return True


def write_xmp_png(filepath: str, data: bytes, description: str) -> bool:
    """Inject XMP as an iTXt chunk into PNG.

    PNG iTXt chunk structure:
      4 bytes: data length (big-endian)
      4 bytes: chunk type ("iTXt")
      N bytes: chunk data
        - keyword (null-terminated): "XML:com.adobe.xmp"
        - compression flag: 0 (no compression)
        - compression method: 0
        - language tag (null-terminated): ""
        - translated keyword (null-terminated): ""
        - text: the XMP packet
      4 bytes: CRC32 of (chunk type + chunk data)
    """
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        return False

    # Strip existing AI cache iTXt chunk
    data = strip_ai_xmp(data, "png")
    itxt_chunk = build_png_itxt_chunk(description)

    # Insert before IEND chunk (last 12 bytes of a valid PNG)
    # Find IEND
    iend_pos = data.rfind(b"IEND")
    if iend_pos == -1:
        return False
    iend_start = iend_pos - 4  # 4 bytes length before "IEND"

    new_data = data[:iend_start] + itxt_chunk + data[iend_start:]

    atomic_write_bytes(filepath, new_data)
    return True


def write_xmp_webp(filepath: str, data: bytes, description: str) -> bool:
    """Inject XMP as a RIFF chunk into WebP.

    WebP RIFF structure:
      "RIFF" + 4-byte LE size + "WEBP" + chunks...
    Each chunk:
      4-byte FourCC + 4-byte LE size + data (padded to even)
    """
    if len(data) < 12 or data[:4] != b"RIFF" or data[8:12] != b"WEBP":
        return False

    # Strip existing AI cache XMP chunk
    data = strip_ai_xmp(data, "webp")

    # Build XMP chunk
    xmp_data = build_xmp_packet(description).encode("utf-8")
    xmp_chunk = b"XMP " + struct.pack("<I", len(xmp_data)) + xmp_data
    # Pad to even length
    if len(xmp_data) % 2 != 0:
        xmp_chunk += b"\x00"

    # Append chunk to RIFF container
    new_data = data + xmp_chunk

    # Update RIFF container size (bytes 4-8, little-endian)
    new_riff_size = len(new_data) - 8
    new_data = new_data[:4] + struct.pack("<I", new_riff_size) + new_data[8:]

    atomic_write_bytes(filepath, new_data)
    return True


def write_sidecar(filepath: str, description: str) -> str:
    """Last resort for GIF/BMP: write a .ai-cache sidecar file."""
    sidecar_path = filepath + ".ai-cache"
    atomic_write_text(sidecar_path, description)
    return sidecar_path


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    if len(sys.argv) < 3:
        print("Usage: write_cache.py <image-path> <content>", file=sys.stderr)
        sys.exit(1)

    filepath = sys.argv[1]
    content = sys.argv[2]

    try:
        data = Path(filepath).read_bytes()
    except FileNotFoundError:
        print(f"Error: file not found: {filepath}", file=sys.stderr)
        sys.exit(1)
    except OSError as exc:
        print(f"Error: failed to read {filepath}: {exc}", file=sys.stderr)
        sys.exit(1)

    fmt = detect_format(data, filepath)

    # Hash BEFORE writing (strips any existing AI cache XMP)
    current_hash = image_hash(data, fmt)
    full_content = f"HASH:{current_hash}|{content}"

    # Try exiftool first (handles everything)
    if write_with_exiftool(filepath, full_content):
        print("OK: exiftool")
        return

    # Format-specific direct injection
    if fmt == "jpeg" and write_xmp_jpeg(filepath, data, full_content):
        print("OK: jpeg-inject")
    elif fmt == "png" and write_xmp_png(filepath, data, full_content):
        print("OK: png-inject")
    elif fmt == "webp" and write_xmp_webp(filepath, data, full_content):
        print("OK: webp-inject")
    else:
        # GIF, BMP, or injection failed — sidecar
        sidecar = write_sidecar(filepath, full_content)
        print(f"SIDECAR: {sidecar}")


if __name__ == "__main__":
    main()