#!/usr/bin/env python3
"""
Generates unique branded cinematic images for models from two sources:

  - Cloud models:  https://api.inferenceport.ai/v1/models   (OpenAI-style {"data":[{"id": ...}]})
  - Local models:  https://api.inferenceport.ai/models       (Ollama-style [{"name": ...}, ...])

using the Pollinations image API (model=zimage).

Naming convention:
  Cloud model "foo/bar":
    {OUTPUT_DIR}/foo_bar.jpg
    {OUTPUT_DIR}/compressed/foo_bar.webp
  Local model "foo/bar":
    {OUTPUT_DIR}/foo_bar.local.jpg
    {OUTPUT_DIR}/compressed/foo_bar.local.webp

Each model gets a deterministically-generated, distinct cinematic prompt:
a bold branded scene (fog, neon, metallic text, surreal landscapes,
cosmic backdrops, etc.) with the model's name rendered as giant stylized
text integrated into the scene. The prompt generator is seeded off the
model's own name, so the same model always gets the same scene, and
different models get different combinations of setting / lighting /
text-material / camera style / palette.

USAGE
-----
Scan both endpoints and fill in anything missing (default command):
    python tools/createpics.py
    python tools/createpics.py scan

Create an image for a single new model:
    python tools/createpics.py create -n "some-model-id" -t cloud
    python tools/createpics.py create -n "llama3:8b" -t local

Regenerate (overwrite) an existing model's image, e.g. if the text
came out warped:
    python tools/createpics.py regenerate -n "some-model-id" -t cloud
"""

import os
import re
import time
import random
import hashlib
import argparse
import urllib.parse
from pathlib import Path
from huggingface_hub import sync_bucket
import requests
from rich.progress import (
    Progress,
    BarColumn,
    TextColumn,
    TimeElapsedColumn,
    MofNCompleteColumn,
)
from rich.console import Console

REPO_ROOT = Path(__file__).resolve().parent.parent

# --------------------------------------------------------------------------
# CONFIG
# --------------------------------------------------------------------------
DEFAULT_API_KEY = "sk_w1UrhHef25QNpFGEFnOGanL4WIaJLbVI"
MODELS_ENDPOINT = "https://api.inferenceport.ai/models"
V1_MODELS_ENDPOINT = "https://us-east-2.api.inferenceport.ai/v1/models"
POLLINATIONS_URL_TEMPLATE = "https://gen.pollinations.ai/image/{prompt}?model=zimage&key={key}&seed={seed}"
REQUEST_TIMEOUT = 60
RETRY_COUNT = 2
RETRY_DELAY = 3
WEBP_QUALITY = 82

console = Console()

# --------------------------------------------------------------------------
# PROMPT COMPONENT POOLS
# Each model's prompt is assembled from one entry per pool, chosen with a
# random.Random() seeded from a hash of the model's own name, so results
# are stable across runs but distinct across models.
# --------------------------------------------------------------------------

SETTINGS = [
    "a fog-drenched pine forest at blue hour, mist curling between towering trunks",
    "a neon-soaked cyberpunk alley slick with rain, steam rising from grates",
    "a colossal cosmic nebula swirling with violet and teal gas clouds",
    "a molten volcanic canyon with rivers of lava cutting through black rock",
    "a vast crystal cave glittering with thousands of embedded gemstones",
    "a chain of floating islands drifting above an endless cloud sea",
    "a cracking glacier under a sky streaked with green aurora",
    "an endless desert of wind-sculpted dunes under a blood-orange sunset",
    "sunken underwater ruins overgrown with coral and bioluminescent algae",
    "a derelict space station drifting past a ringed gas giant",
    "an ancient stone temple half-swallowed by jungle vines",
    "an infinite mirror maze reflecting fractured shards of light",
    "a churning storm-tossed sea beneath towering black clouds",
    "a frozen arctic plain lit by dancing curtains of aurora",
    "a rain-slicked subway tunnel lit by flickering fluorescent tubes",
    "a greenhouse jungle overtaken by glowing bioluminescent flora",
    "a colossal clocktower interior filled with spinning brass gears",
    "a marble colosseum ruin bathed in golden late-afternoon light",
    "a roaring industrial forge with sparks cascading from molten metal",
    "a star-forged gothic cathedral with stained glass depicting galaxies",
    "a wind-swept mountain summit above a sea of clouds at sunrise",
    "a bustling night market alley strung with paper lanterns",
    "a vast salt flat perfectly mirroring a star-filled sky",
    "an abandoned amusement park reclaimed by fog and neon decay",
    "a deep canyon lit by shafts of light piercing storm clouds",
    "a zero-gravity chamber filled with drifting shards of broken glass",
    "a bamboo forest at dawn with shafts of golden light breaking through",
    "a lava-lit obsidian throne room deep beneath a dormant volcano",
    "a vast library of floating books orbiting a soft golden light source",
    "a windswept lighthouse cliff battered by crashing waves at dusk",
    "a shattered glacier cavern lit from within by trapped blue fire",
    "an endless server-farm hallway stretching into infinite reflections",
    "a surreal upside-down city hanging from a cracked stone sky",
    "a canyon of stacked shipping containers glowing with graffiti neon",
    "a vast wheat field rippling under a sky full of shooting stars",
    "a glass observatory dome overlooking a spiral galaxy at close range",
    "a coral-pink slot canyon carved by centuries of desert wind",
    "an orbital ring city seen from a docking bay window",
    "a moss-covered clockwork ruin swallowed by an ancient forest",
    "a vast obsidian mirror lake reflecting a double moon",
    "a subterranean bazaar lit by hanging lanterns and glowing runes",
    "a wind farm silhouetted against a violent orange sunset",
    "a crumbling cathedral of ice deep in a polar cavern",
    "a floating night market suspended above a bottomless chasm",
    "a bioluminescent mangrove swamp glowing teal in the darkness",
    "a vast concrete brutalist plaza swept by dust and wind",
]

LIGHTING = [
    "harsh electric-blue neon glow",
    "warm golden god-rays breaking through haze",
    "eerie pulsing bioluminescent light",
    "flickering violet lightning strikes",
    "soft silver moonlight filtering through mist",
    "rippling aurora-colored light bands",
    "glowing embers and firelight",
    "cool holographic shimmer",
    "dramatic single-source spotlight",
    "deep crimson volcanic glow",
    "cold cyan starlight",
    "warm amber lantern light",
    "stark strobing emergency-red light",
    "diffuse overcast twilight",
    "blazing white magnesium flare light",
    "shifting prismatic refracted light",
    "sodium-vapor orange streetlight glow",
    "deep-space starfield ambient glow",
    "candlelit flickering warmth",
    "harsh midday desert glare",
]

TEXT_MATERIALS = [
    "liquid chrome letters, mirror-polished and dripping light",
    "molten gold lettering, glowing and semi-fluid",
    "frosted glass letters lit from within",
    "carved obsidian letters etched with faint neon veins",
    "circuit-etched neon tubing bent into giant letters",
    "hand-forged iron letters radiating heat",
    "crystalline ice letters refracting rainbow light",
    "bioluminescent organic script pulsing softly",
    "brushed titanium letters catching hard highlights",
    "marble letters carved with gold-leaf inlay",
    "holographic letters flickering between solid and translucent",
    "letters built from swirling starlight and stardust",
    "letters woven from living vines and blooming flowers",
    "polished copper letters streaked with verdigris",
    "letters made of stacked weathered stone blocks",
    "letters formed from suspended droplets of mercury",
    "letters etched into a cliff face by ancient hands",
    "letters built from interlocking clockwork gears",
    "letters made of compressed storm clouds and lightning",
    "letters carved from a single glowing block of amber",
]

CAMERA_STYLES = [
    "ultra-wide cinematic lens, dramatic rim light, high production value",
    "anamorphic lens flare, shallow depth of field, epic scale",
    "symmetrical composition, painterly atmosphere, hyperreal detail",
    "low-angle heroic framing, volumetric fog, blockbuster lighting",
    "aerial establishing shot, vast scale, crisp atmospheric haze",
    "macro-to-wide hybrid framing, intricate detail, dreamlike surrealism",
    "dutch-angle tension shot, high contrast, gritty texture",
    "telephoto compression, layered depth, soft bokeh backdrop",
    "fisheye distortion, immersive foreground, hyper-detailed",
    "long-exposure motion blur, streaking light trails, cinematic grade",
]

COLOR_PALETTES = [
    "teal and orange cinematic color grade",
    "deep violet and acid green duotone",
    "monochrome silver with a single accent color",
    "warm sepia and burnt umber tones",
    "icy cyan and stark white palette",
    "crimson and charcoal high-contrast palette",
    "pastel dawn palette of pink and lavender",
    "saturated jewel-tone palette of emerald and sapphire",
    "desaturated muted earth-tone palette",
    "hyper-saturated synthwave magenta and blue palette",
]


def build_prompt(display_name: str) -> str:
    """Deterministically build a unique cinematic branded prompt for a model name."""
    seed = int(hashlib.sha256(display_name.encode("utf-8")).hexdigest(), 16) + random.randint(0, 200)
    rng = random.Random(seed)

    setting = rng.choice(SETTINGS)
    lighting = rng.choice(LIGHTING)
    material = rng.choice(TEXT_MATERIALS)
    camera = rng.choice(CAMERA_STYLES)
    palette = rng.choice(COLOR_PALETTES)
    upper_name = display_name.upper()

    return (
        f"Cinematic branded scene: {setting}, illuminated by {lighting}. "
        f'Towering, giant-scale {material} spelling out "{upper_name}" '
        f"integrated seamlessly into the environment as the focal point. "
        f"{camera}. {palette}."
    )


def safe_filename(name: str) -> str:
    """Turn a model name (which may contain slashes, dots, colons, etc.) into a safe filename."""
    return re.sub(r"[^A-Za-z0-9._-]", "_", name)


def paths_for(model_id: str, model_type: str, output_dir: str):
    """Return (jpg_path, webp_path) for a given model id/type."""
    safe = safe_filename(model_id)
    if model_type == "local":
        jpg_name = f"{safe}.local.jpg"
        webp_name = f"{safe}.local.webp"
    else:
        jpg_name = f"{safe}.jpg"
        webp_name = f"{safe}.webp"
    jpg_path = os.path.join(output_dir, jpg_name)
    webp_path = os.path.join(output_dir, "compressed", webp_name)
    return jpg_path, webp_path


# --------------------------------------------------------------------------
# FETCHING MODEL LISTS
# --------------------------------------------------------------------------

def fetch_local_models() -> list[str]:
    """Fetch the /models (Ollama-style) list and return model names."""
    response = requests.get(MODELS_ENDPOINT, timeout=REQUEST_TIMEOUT)
    response.raise_for_status()
    data = response.json()

    names = []
    for entry in data:
        name = entry.get("name")
        if name:
            names.append(name)
    return names


def fetch_cloud_models() -> list[str]:
    """Fetch the /v1/models (OpenAI-style) list and return model ids."""
    response = requests.get(V1_MODELS_ENDPOINT, timeout=REQUEST_TIMEOUT)
    response.raise_for_status()
    data = response.json()

    items = data.get("data", data) if isinstance(data, dict) else data

    ids = []
    for entry in items:
        model_id = entry.get("id") or entry.get("name")
        if model_id:
            ids.append(model_id)
    return ids


# --------------------------------------------------------------------------
# IMAGE GENERATION / CONVERSION
# --------------------------------------------------------------------------

def fetch_and_save_image(model_id: str, prompt: str, jpg_path: str, api_key: str, seed: int = 0) -> bool:
    """Fetch one image from Pollinations and save it to jpg_path. Returns success bool."""
    encoded_prompt = urllib.parse.quote(prompt, safe="")
    url = POLLINATIONS_URL_TEMPLATE.format(prompt=encoded_prompt, key=api_key, seed=seed)

    os.makedirs(os.path.dirname(jpg_path) or ".", exist_ok=True)

    last_error = None
    for attempt in range(1, RETRY_COUNT + 2):  # initial try + retries
        try:
            response = requests.get(url, timeout=REQUEST_TIMEOUT)
            response.raise_for_status()

            content_type = response.headers.get("Content-Type", "")
            if "image" not in content_type and len(response.content) < 1000:
                raise ValueError(
                    f"Unexpected response (content-type={content_type!r}, "
                    f"{len(response.content)} bytes): {response.content[:200]!r}"
                )

            with open(jpg_path, "wb") as f:
                f.write(response.content)
            return True

        except Exception as e:  # noqa: BLE001 - catch & report everything
            last_error = e
            if attempt <= RETRY_COUNT:
                time.sleep(RETRY_DELAY)

    console.print(f"[bold red]✗ Failed to generate:[/bold red] {model_id} — {last_error}")
    return False


def convert_to_webp(jpg_path: str, webp_path: str, quality: int = WEBP_QUALITY) -> bool:
    """Convert an existing jpg to webp. Returns success bool."""
    try:
        from PIL import Image
    except ImportError:
        console.print(
            "[bold red]Pillow is required for webp conversion.[/bold red] "
            "Install it with: pip install pillow --break-system-packages"
        )
        return False

    os.makedirs(os.path.dirname(webp_path) or ".", exist_ok=True)
    try:
        with Image.open(jpg_path) as im:
            im = im.convert("RGB")
            im.save(webp_path, "WEBP", quality=quality)
        return True
    except Exception as e:  # noqa: BLE001
        console.print(f"[bold red]✗ WebP conversion failed:[/bold red] {jpg_path} — {e}")
        return False


def process_model(model_id: str, model_type: str, output_dir: str, api_key: str, overwrite: bool = False, seed: int = 0) -> bool:
    """Generate (if needed) and convert (if needed) images for a single model."""
    jpg_path, webp_path = paths_for(model_id, model_type, output_dir)
    prompt = build_prompt(model_id)

    need_jpg = overwrite or not os.path.exists(jpg_path)
    if need_jpg:
        ok = fetch_and_save_image(model_id, prompt, jpg_path, api_key, seed=seed)
        if not ok:
            return False

    need_webp = overwrite or not os.path.exists(webp_path)
    if need_webp:
        convert_to_webp(jpg_path, webp_path)

    return True


# --------------------------------------------------------------------------
# COMMANDS
# --------------------------------------------------------------------------

def cmd_scan(output_dir: str, api_key: str):
    console.print("[bold cyan]Fetching model lists...[/bold cyan]")

    try:
        cloud_models = fetch_cloud_models()
    except Exception as e:
        console.print(f"[bold red]Failed to fetch /v1/models:[/bold red] {e}")
        cloud_models = []

    try:
        local_models = fetch_local_models()
    except Exception as e:
        console.print(f"[bold red]Failed to fetch /models:[/bold red] {e}")
        local_models = []

    console.print(
        f"[bold cyan]Found {len(cloud_models)} cloud models and "
        f"{len(local_models)} local models.[/bold cyan]"
    )

    all_jobs = [(m, "cloud") for m in cloud_models] + [(m, "local") for m in local_models]

    # Only queue work for models missing a jpg and/or webp.
    pending = []
    for model_id, model_type in all_jobs:
        jpg_path, webp_path = paths_for(model_id, model_type, output_dir)
        if not os.path.exists(jpg_path) or not os.path.exists(webp_path):
            pending.append((model_id, model_type))

    console.print(f"[bold cyan]{len(pending)} image(s) need generating/converting.[/bold cyan]\n")

    if not pending:
        console.print("[bold green]Everything is already up to date.[/bold green]")
        return

    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        MofNCompleteColumn(),
        TimeElapsedColumn(),
        console=console,
    ) as progress:
        task = progress.add_task("Generating images", total=len(pending))

        for model_id, model_type in pending:
            progress.update(task, description=f"[green]{model_id} ({model_type})[/green]")
            process_model(model_id, model_type, output_dir, api_key)
            progress.advance(task)

    console.print(f"\n[bold green]Done.[/bold green] Images saved under [bold]{output_dir}[/bold]")


def cmd_create(name: str, model_type: str, output_dir: str, api_key: str):
    jpg_path, webp_path = paths_for(name, model_type, output_dir)

    if os.path.exists(jpg_path) and os.path.exists(webp_path):
        console.print(
            f"[bold yellow]Already exists:[/bold yellow] {jpg_path}\n"
            f"Use the 'regenerate' command if you want to overwrite it."
        )
        return

    console.print(f"[bold cyan]Creating image for[/bold cyan] {name} [dim]({model_type})[/dim]")
    ok = process_model(name, model_type, output_dir, api_key, overwrite=False)
    if ok:
        console.print(f"[bold green]Saved:[/bold green] {jpg_path}\n[bold green]Saved:[/bold green] {webp_path}")
    else:
        console.print("[bold red]Failed to create image.[/bold red]")


def cmd_regenerate(name: str, model_type: str, output_dir: str, api_key: str):
    jpg_path, webp_path = paths_for(name, model_type, output_dir)

    console.print(f"[bold cyan]Regenerating image for[/bold cyan] {name} [dim]({model_type})[/dim]")
    ok = process_model(name, model_type, output_dir, api_key, overwrite=True, seed=random.randint(1, 200))
    if ok:
        console.print(f"[bold green]Overwrote:[/bold green] {jpg_path}\n[bold green]Overwrote:[/bold green] {webp_path}")
    else:
        console.print("[bold red]Failed to regenerate image.[/bold red]")


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------

def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate branded images for inference models (cloud + local)."
    )
    parser.add_argument("--api-key", help="Your Pollinations API key.")
    parser.add_argument("--output-dir", default=str(REPO_ROOT / "modelpics"), help="Directory to save generated images.")

    subparsers = parser.add_subparsers(dest="command")

    subparsers.add_parser(
        "scan",
        help="Scan both endpoints and generate any missing images/webp conversions (default).",
    )

    create_p = subparsers.add_parser("create", help="Create an image for a single new model.")
    create_p.add_argument("-n", "--name", required=True, help="The model ID/name.")
    create_p.add_argument("-t", "--type", required=True, choices=["cloud", "local"], help="Model source type.")

    regen_p = subparsers.add_parser("regenerate", help="Regenerate (overwrite) an existing model's image.")
    regen_p.add_argument("-n", "--name", required=True, help="The model ID/name.")
    regen_p.add_argument("-t", "--type", required=True, choices=["cloud", "local"], help="Model source type.")

    subparsers.add_parser("sync", help="Sync the local output directory with the HF bucket.")
    return parser

def get_md5(file_path):
    """Generate a quick hash of the file content."""
    hasher = hashlib.md5()
    with open(file_path, 'rb') as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hasher.update(chunk)
    return hasher.hexdigest()
def clean_local_renames_before_sync(local_dir):
    """Finds duplicate files by content and removes the older-named version."""
    seen_hashes = {}
    for root, _, files in os.walk(local_dir):
        for file in files:
            path = os.path.join(root, file)
            f_hash = get_md5(path)
            
            if f_hash in seen_hashes:
                # We found a rename/duplicate! Delete the ghost/old copy
                print(f"⚠️ Detected rename conflict. Removing local ghost file: {path}")
                os.remove(path)
            else:
                seen_hashes[f_hash] = path

def sync(output_dir):
    clean_local_renames_before_sync(output_dir)

    sync_bucket("hf://buckets/inferenceport-ai/model-info", output_dir, delete=False)
    
    sync_bucket(output_dir, "hf://buckets/inferenceport-ai/model-info", delete=False)

def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    api_key = args.api_key or DEFAULT_API_KEY
    output_dir = args.output_dir

    if api_key in ("YOUR_POLLINATIONS_API_KEY_HERE", "", None):
        console.print(
            "[bold yellow]Warning:[/bold yellow] API_KEY is still a placeholder. "
            "Pass --api-key or set DEFAULT_API_KEY before running.\n"
        )

    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, "compressed"), exist_ok=True)

    command = args.command or "scan"

    if command == "scan":
        cmd_scan(output_dir, api_key)
    elif command == "create":
        cmd_create(args.name, args.type, output_dir, api_key)
    elif command == "regenerate":
        cmd_regenerate(args.name, args.type, output_dir, api_key)
    elif command == "sync":
        sync(output_dir)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()