#!/usr/bin/env python3
"""Rename image files to descriptive snake_case names using a local vision model.

The model (default: llava) is queried through a running Ollama server. This
script ensures the server is up and the model is pulled before processing.
"""
import argparse
import base64
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time

IMG_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
PROMPT = (
    "You are naming an image file. Look at the image and respond with ONLY a "
    "filename, nothing else. Rules: all lowercase; snake_case (underscores only, "
    "no spaces, hyphens, or special characters); start with 'a_' or 'an_'; describe "
    "the main subject, its state, and the background; 5 to 10 words; end with the "
    "correct extension (.jpg/.png/.jpeg). Example: a_snowy_mountain_with_grey_sky.jpg"
)


def run(cmd, **kw):
    return subprocess.run(cmd, capture_output=True, text=True, **kw)


def ollama_ready():
    r = run(["curl", "-s", "-o", "/dev/null", "-w", "%{http_code}",
             "--max-time", "3", "http://localhost:11434/api/tags"])
    return r.returncode == 0 and r.stdout.strip() == "200"


def ensure_server(timeout=90):
    if ollama_ready():
        return
    print("Starting ollama server...", flush=True)
    proc = subprocess.Popen(["ollama", "serve"],
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    deadline = time.time() + timeout
    while time.time() < deadline:
        if ollama_ready():
            return
        time.sleep(1)
    raise SystemExit("ollama server did not become ready in time")


def ensure_model(model, timeout=600):
    listing = run(["ollama", "list"], timeout=30)
    have = set()
    for line in listing.stdout.splitlines()[1:]:
        have.add(line.split()[0].replace(":latest", ""))
    if model.replace(":latest", "") in have:
        return
    print(f"Pulling model '{model}' (one-time download)...", flush=True)
    r = run(["ollama", "pull", model], timeout=timeout)
    if r.returncode != 0:
        raise SystemExit("failed to pull model:\n" + r.stderr)


def resize(path):
    out = tempfile.mktemp(suffix=".jpg")
    run(["ffmpeg", "-y", "-loglevel", "error", "-i", path,
         "-vf", "scale='min(768,iw)':'min(768,ih)'", "-q:v", "4", out],
        timeout=120)
    if not os.path.exists(out):
        # fall back to the original file if ffmpeg chokes
        return path
    return out


def describe(path, model):
    r = resize(path)
    with open(r, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": PROMPT, "images": [b64]}],
        "stream": False,
        "options": {"temperature": 0.2},
    }
    pf = tempfile.NamedTemporaryFile("w", suffix=".json", delete=False)
    json.dump(payload, pf)
    pf.close()
    last = None
    for _ in range(3):
        try:
            p = run(["curl", "-s", "http://localhost:11434/api/chat",
                     "-H", "Content-Type: application/json", "-d", f"@{pf.name}"],
                    timeout=240)
            return json.loads(p.stdout)["message"]["content"].strip()
        except Exception as e:  # noqa: BLE001
            last = e
    raise last


def sanitize(model_out, ext):
    m = re.search(r"[a-z0-9_]+\.(?:jpg|jpeg|png|gif|webp)", model_out.lower())
    name = m.group(0) if m else re.sub(r"[^a-z0-9_]", "_", model_out.lower())
    stem, _ = os.path.splitext(name)
    stem = re.sub(r"[^a-z0-9_]", "_", stem)
    stem = re.sub(r"_+", "_", stem).strip("_")
    if not (stem.startswith("a_") or stem.startswith("an_")):
        stem = "a_" + stem
    return stem + ext.lower()


def already_named(fn):
    s = fn.lower()
    return s.startswith("a_") or s.startswith("an_") or s == "retro_mountain_red.jpg"


def find_images(directory, recursive):
    found = []
    walker = os.walk(directory) if recursive else [(directory, [], os.listdir(directory))]
    for root, _, files in walker:
        for fn in files:
            if os.path.splitext(fn)[1].lower() in IMG_EXTS:
                found.append(os.path.join(root, fn))
    return sorted(found)


def pick_with_fzf(paths):
    # Feed full paths but display only the basename (--with-nth); fzf still
    # outputs the full path, and the preview renders an image thumbnail.
    by_full = {p: p for p in paths}
    by_base = {os.path.basename(p): p for p in paths}
    preview_cmd = os.environ.get(
        "RENAME_PREVIEW_CMD",
        'chafa --size 50x32 --color-space rgb {}')
    r = run(["fzf", "--multi", "--prompt", "pick images to rename> ",
             "--delimiter", "/", "--with-nth", "-1",
             "--preview", preview_cmd,
             "--preview-window", "right:50%:wrap",
             "--border", "--height", "100%"],
            input="\n".join(sorted(paths)), timeout=None)
    if r.returncode != 0:
        return []
    chosen = [line for line in r.stdout.splitlines() if line]
    return [by_full.get(c, by_base.get(os.path.basename(c), c)) for c in chosen]


def process(path, model, force, dry_run):
    fn = os.path.basename(path)
    ext = os.path.splitext(fn)[1].lower()
    if not force and already_named(fn):
        print(f"SKIP {fn}", flush=True)
        return None
    try:
        raw = describe(path, model)
        new = sanitize(raw, ext)
    except Exception as e:  # noqa: BLE001
        print(f"FAIL {fn}: {e}", flush=True)
        return None
    target = path if force else os.path.join(os.path.dirname(path), new)
    i = 1
    while os.path.exists(target) and os.path.abspath(target) != os.path.abspath(path):
        target = os.path.join(os.path.dirname(path),
                              f"{os.path.splitext(new)[0]}_{i}{ext}")
        i += 1
    base = os.path.basename(target)
    if dry_run:
        print(f"{fn} => {base}   (raw: {raw})", flush=True)
    else:
        shutil.move(path, target)
        print(f"{fn} => {base}", flush=True)
    return base


def main():
    ap = argparse.ArgumentParser(
        description="Rename images to descriptive snake_case names via a local vision model.")
    ap.add_argument("files", nargs="*", help="image files to rename")
    ap.add_argument("--dir", default=None, help="directory to scan (default: current dir)")
    ap.add_argument("--recursive", action="store_true", help="scan directories recursively")
    ap.add_argument("--pick", action="store_true", help="interactively pick with fzf")
    ap.add_argument("--model", default=os.environ.get("RENAME_MODEL", "llava"),
                    help="ollama vision model (default: llava)")
    ap.add_argument("--force", action="store_true", help="rename even already-named files")
    ap.add_argument("--dry-run", action="store_true", help="print names without renaming")
    args = ap.parse_args()

    ensure_server()
    ensure_model(args.model)

    paths = list(args.files)
    if args.dir:
        paths += find_images(args.dir, args.recursive)
    if not paths:
        paths = find_images(args.dir or ".", args.recursive)
    # normalise and dedupe
    paths = sorted({os.path.abspath(p) for p in paths})

    if args.pick:
        paths = pick_with_fzf(paths)
        if not paths:
            print("Nothing selected.", flush=True)
            return

    if not paths:
        print("No images found.", flush=True)
        return

    done = 0
    for p in paths:
        if process(p, args.model, args.force, args.dry_run):
            done += 1
    print(f"\nDone. {done} image(s) renamed"
          + (" (dry-run)" if args.dry_run else "") + ".", flush=True)


if __name__ == "__main__":
    main()