# gemini_nearest_obstacle_name.py
import os
import json
import time
import argparse
from typing import Dict

from PIL import Image
from google import genai
from tqdm import tqdm

SYSTEM_MESSAGE = (
    "You are a mobility assistant who analyzes the scene for safe navigation. "
    "Be concise and accurate."
)

QUESTION = (
    "Identify the nearest obstacle on the sidewalk or walkable path ahead. "
    "Output ONLY the object name. "
    "No punctuation, no explanation, no full sentences."
)

DEFAULT_MODEL = "gemini-3-pro-preview"


def _clean_object_name(text: str) -> str:
    """Keep first line; strip common wrappers/whitespace."""
    if not text:
        return ""
    t = text.strip()

    # Keep only first non-empty line
    lines = [ln.strip() for ln in t.splitlines() if ln.strip()]
    if not lines:
        return ""
    t = lines[0]

    # Strip trivial quotes/backticks
    t = t.strip("`\"' \t")
    return t


def ask_gemini_object_name(client: genai.Client, image_path: str, model_id: str) -> str:
    image = Image.open(image_path).convert("RGB")
    image.thumbnail((768, 768))

    contents = [
        SYSTEM_MESSAGE,
        image,
        QUESTION,
    ]

    resp = client.models.generate_content(
        model=model_id,
        contents=contents,
        # If your installed google-genai supports config, you can uncomment:
        # config={"temperature": 0.0},
    )

    text = getattr(resp, "text", "") or ""
    name = _clean_object_name(text)
    return name


def iter_pngs(folder: str):
    for fname in sorted(os.listdir(folder)):
        if fname.lower().endswith(".png"):
            yield fname, os.path.join(folder, fname)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--folder",
        default="/scratch/ds5725/OOPS/images_resized",
        help="Folder containing PNG images",
    )
    parser.add_argument("--model", default=DEFAULT_MODEL, help="Gemini model id")
    parser.add_argument(
        "--out",
        default=None,
        help="Output path base (without extension). "
             "If omitted, uses <folder>/nearest_object_name_dict",
    )
    parser.add_argument("--sleep", type=float, default=0.2, help="Sleep seconds between requests")
    parser.add_argument("--retries", type=int, default=5, help="Retries per image on failure")
    args = parser.parse_args()

    api_key = "AIzaSyCjz1zbRQ_57ovEBPN2rlbfPYm2qVOEiuY"
    if not api_key:
        raise RuntimeError(
            "Missing GEMINI_API_KEY env var.\n"
            "Do: export GEMINI_API_KEY='...'\n"
        )

    client = genai.Client(api_key=api_key)

    folder = args.folder
    if not os.path.isdir(folder):
        raise FileNotFoundError(f"Folder not found: {folder}")

    out_base = args.out
    if out_base is None:
        out_base = os.path.join(folder, "nearest_object_name_dict")

    results: Dict[str, str] = {}

    # Optional progress bar if available

    iterator = list(iter_pngs(folder))
    pbar = tqdm(iterator, desc="Gemini nearest object", unit="img")
    for fname, fpath in pbar:
        if fname in results:
            continue

        last_err = None
        for attempt in range(1, args.retries + 1):
            try:
                obj = ask_gemini_object_name(client, fpath, args.model)
                results[fname] = obj
                pbar.set_postfix_str(obj[:40] if obj else "EMPTY")
                break
            except Exception as e:
                last_err = e
                # exponential backoff
                backoff = min(8.0, 0.5 * (2 ** (attempt - 1)))
                time.sleep(backoff)
        else:
            # exhausted retries
            results[fname] = ""
            pbar.set_postfix_str(f"FAILED: {type(last_err).__name__}")

        time.sleep(args.sleep)
   

    # Save outputs
    json_path = out_base + ".json"

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"Saved {len(results)} entries")
    print(f"- JSON: {json_path}")
   


if __name__ == "__main__":
    main()