#!/usr/bin/env python3
"""
Caption Tool — caption entire folders of images locally.

Runs 100% on your machine. No API key, no internet after the first model
download. Works on GPUs with as little as 2 GB of VRAM (use Florence-2-base).

For every image it writes a matching <image>.txt caption file next to it.
"""

import argparse
import json
import os
import sys
import time
import warnings

warnings.filterwarnings("ignore")
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")

from pathlib import Path

from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt, Confirm
from rich.progress import (
    Progress,
    SpinnerColumn,
    BarColumn,
    TextColumn,
    TimeElapsedColumn,
    TaskProgressColumn,
)
from rich.live import Live
from rich.table import Table
from rich.text import Text

console = Console()

CONFIG_PATH = Path(__file__).parent / "caption_config.json"

# Model choices. Sizes are approximate fp16 GPU footprint.
MODELS = {
    "florence2-large":    ("microsoft/Florence-2-large",    "~1.5 GB — base, slightly less accurate"),
    "florence2-base":     ("microsoft/Florence-2-base",     "~0.45 GB — fast, great for 2 GB GPUs"),
}

DEFAULT_CONFIG = {
    "model": "microsoft/Florence-2-large",
    "prefix": "",
    "suffix": "",
    "max_new_tokens": 256,
    "min_words": 15,
    "skip_existing": True,
}


def _model_family(model_id: str) -> str:
    return "florence2" if "florence" in model_id.lower() else "unknown"


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif"}


def load_config() -> dict:
    if CONFIG_PATH.exists():
        with open(CONFIG_PATH) as f:
            cfg = json.load(f)
        return {**DEFAULT_CONFIG, **cfg}
    return DEFAULT_CONFIG.copy()


def save_config(cfg: dict):
    with open(CONFIG_PATH, "w") as f:
        json.dump(cfg, f, indent=2)


def find_images(path: Path) -> list[Path]:
    if path.is_file():
        if path.suffix.lower() in IMAGE_EXTENSIONS:
            return [path]
        return []
    return sorted(p for p in path.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS)


def load_model(model_id: str):
    """Load the captioning model onto the GPU, or CPU if no GPU is present."""
    import torch
    from transformers import AutoProcessor, AutoModelForCausalLM

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device_label = (
        f"CUDA ({torch.cuda.get_device_name(0)})" if device.type == "cuda" else "CPU"
    )
    dtype = torch.float16 if device.type == "cuda" else torch.float32
    family = _model_family(model_id)

    with console.status(f"[bold cyan]Loading [white]{model_id}[/white] on {device_label}...", spinner="dots"):
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_id, torch_dtype=dtype, trust_remote_code=True
        ).to(device)
        model.eval()

    return (processor, model, device, dtype, family), device_label


def _generate_caption(processor, model, img, device, dtype, max_new_tokens: int, task: str) -> str:
    """Run one Florence-2 caption generation for the given task prompt."""
    import torch

    inputs = processor(text=task, images=img, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    if "pixel_values" in inputs:
        inputs["pixel_values"] = inputs["pixel_values"].to(dtype)
    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, num_beams=3)
    raw = processor.batch_decode(output_ids, skip_special_tokens=False)[0]
    parsed = processor.post_process_generation(raw, task=task, image_size=(img.width, img.height))
    return parsed[task].strip()


def caption_image(model_bundle, image_path: Path, prefix: str, suffix: str,
                 max_new_tokens: int, min_words: int = 0) -> str:
    import torch
    from PIL import Image

    processor, model, device, dtype, family = model_bundle
    img = Image.open(image_path).convert("RGB")

    with torch.no_grad():
        caption = _generate_caption(
            processor, model, img, device, dtype, max_new_tokens, "<MORE_DETAILED_CAPTION>"
        )

        # Florence-2 has no native minimum length, so re-run with a stronger
        # prompt if the first caption came out too short.
        if min_words and len(caption.split()) < min_words:
            stronger = (
                "Describe this image in extensive detail for a training caption. "
                "Cover the subject, clothing, hair, lighting, pose, expression, "
                "background, and camera angle."
            )
            caption = _generate_caption(
                processor, model, img, device, dtype, max_new_tokens, stronger
            )

    parts = [p for p in (prefix.strip(), caption, suffix.strip()) if p]
    return ", ".join(parts) if prefix or suffix else caption


def select_model(current: str) -> str:
    console.print("\n[bold]Available models:[/bold]")
    table = Table(show_header=False, box=None, padding=(0, 2))
    for key, (model_id, desc) in MODELS.items():
        marker = "→" if model_id == current else " "
        table.add_row(f"[cyan]{marker} {key}[/cyan]", model_id, f"[dim]{desc}[/dim]")
    console.print(table)

    choice = Prompt.ask(
        "\nModel shortname or full HF model ID",
        default=next((k for k, (m, _) in MODELS.items() if m == current), current),
    )
    if choice in MODELS:
        return MODELS[choice][0]
    return choice  # allow any Hugging Face model ID


def parse_args():
    p = argparse.ArgumentParser(description="Caption Tool — local batch image captioner.")
    p.add_argument("--path", help="Image folder or file (skips the path prompt)")
    p.add_argument("--model", help="Override the model (shortname or HF ID)")
    p.add_argument("--max-new-tokens", type=int, help="Override max new tokens")
    p.add_argument("--force", action="store_true", help="Re-caption even if a .txt already exists")
    return p.parse_args()


def _hsl_hex(hue: float) -> str:
    """Hue (0-360) → hex color string, full saturation, mid lightness."""
    hue %= 360
    s, l = 1.0, 0.55
    c = (1 - abs(2 * l - 1)) * s
    hp = hue / 60
    x = c * (1 - abs((hp % 2) - 1))
    m = l - c / 2
    if hp < 1:
        r, g, b = c, x, 0
    elif hp < 2:
        r, g, b = x, c, 0
    elif hp < 3:
        r, g, b = 0, c, x
    elif hp < 4:
        r, g, b = 0, x, c
    elif hp < 5:
        r, g, b = x, 0, c
    else:
        r, g, b = c, 0, x
    R = int((r + m) * 255)
    G = int((g + m) * 255)
    B = int((b + m) * 255)
    return f"#{R:02x}{G:02x}{B:02x}"


def _header_text(offset: int) -> "Text":
    """Build the title panel content with a rainbow-shimmer oohfixer.com."""
    t = Text()
    t.append("Caption Tool", style="bold white")
    t.append("\n")
    t.append("Caption entire image folders locally — no API key required", style="dim")
    t.append("\n")
    for i, ch in enumerate("oohfixer.com"):
        hue = (i * 18 + offset) % 360
        t.append(ch, style=f"bold {_hsl_hex(hue)}")
    return t


def show_header():
    """Print the title. Animate a rainbow shimmer on a real terminal; static otherwise."""
    panel = lambda off: Panel(_header_text(off), style="bold blue", expand=False)  # noqa: E731
    if console.is_terminal:
        with Live(console=console, refresh_per_second=24, transient=False) as live:
            for frame in range(40):
                live.update(panel(frame * 14))
                time.sleep(0.03)
    else:
        console.print(panel(0))


def main():
    args = parse_args()

    show_header()

    cfg = load_config()

    # --- Path input (use --path if given) ---
    if args.path:
        target = Path(args.path).expanduser().resolve()
    else:
        raw = Prompt.ask("\n[bold]Image folder or file path[/bold]")
        target = Path(raw).expanduser().resolve()

    if not target.exists():
        console.print(f"[red]✗ Path not found:[/red] {target}")
        sys.exit(1)

    images = find_images(target)
    if not images:
        console.print("[red]✗ No images found.[/red]")
        sys.exit(1)

    # CLI / config overrides
    if args.model:
        cfg["model"] = args.model
    if args.max_new_tokens:
        cfg["max_new_tokens"] = args.max_new_tokens
    if args.force:
        cfg["skip_existing"] = False

    # Filter already-captioned unless skipping is disabled
    pending = images
    if cfg["skip_existing"]:
        pending = [p for p in images if not p.with_suffix(".txt").exists()]
        skipped = len(images) - len(pending)
    else:
        skipped = 0

    console.print(f"\n[green]✓[/green] Found [bold]{len(images)}[/bold] image(s)", end="")
    if skipped:
        console.print(f"  [dim]({skipped} already captioned, skipping)[/dim]", end="")
    console.print()

    if not pending:
        console.print("[yellow]All images already have captions. Run with --force to redo them.[/yellow]")
        sys.exit(0)

    # --- Settings summary + optional changes (skip prompts when --path given) ---
    console.print(Panel(
        f"[cyan]Model:[/cyan]       {cfg['model']}\n"
        f"[cyan]Prefix:[/cyan]      {cfg['prefix'] or '[dim](none)[/dim]'}\n"
        f"[cyan]Suffix:[/cyan]      {cfg['suffix'] or '[dim](none)[/dim]'}\n"
        f"[cyan]Max tokens:[/cyan]  {cfg['max_new_tokens']}\n"
        f"[cyan]Min words:[/cyan]  {cfg['min_words']} (0 = no minimum)\n"
        f"[cyan]Skip existing:[/cyan] {cfg['skip_existing']}",
        title="[bold]Current Settings[/bold]",
        expand=False,
    ))

    # Only offer to change settings when running interactively (no --path).
    if not args.path and Confirm.ask("Change settings?", default=False):
        cfg["model"] = select_model(cfg["model"])
        cfg["prefix"] = Prompt.ask("Caption prefix (trigger word, etc.)", default=cfg["prefix"])
        cfg["suffix"] = Prompt.ask("Caption suffix", default=cfg["suffix"])
        cfg["max_new_tokens"] = int(Prompt.ask("Max new tokens", default=str(cfg["max_new_tokens"])))
        cfg["min_words"] = int(Prompt.ask("Min words (re-caption if shorter)", default=str(cfg["min_words"])))
        cfg["skip_existing"] = Confirm.ask("Skip already-captioned images?", default=cfg["skip_existing"])
        save_config(cfg)
        console.print("[dim]Settings saved.[/dim]")

    proceed = args.path or Confirm.ask(f"\nCaption [bold]{len(pending)}[/bold] image(s)?", default=True)
    if not proceed:
        console.print("[dim]Aborted.[/dim]")
        sys.exit(0)

    # --- Load model ---
    console.print()
    try:
        model_bundle, device_label = load_model(cfg["model"])
    except Exception as e:
        console.print(f"[red]✗ Failed to load model:[/red] {e}")
        sys.exit(1)

    console.print(f"[green]✓[/green] Model ready on [bold]{device_label}[/bold]\n")

    # --- Caption loop ---
    errors = []
    done = 0

    with Progress(
        SpinnerColumn(),
        TextColumn("[bold cyan]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        TimeElapsedColumn(),
        console=console,
        transient=False,
    ) as progress:
        task = progress.add_task("Captioning...", total=len(pending))

        for img_path in pending:
            progress.update(task, description=f"[bold cyan]{img_path.name}")
            try:
                caption = caption_image(
                    model_bundle, img_path,
                    cfg["prefix"], cfg["suffix"], cfg["max_new_tokens"], cfg["min_words"]
                )
                txt_path = img_path.with_suffix(".txt")
                txt_path.write_text(caption, encoding="utf-8")
                done += 1
            except Exception as e:
                errors.append((img_path.name, str(e)))
            progress.advance(task)

    # --- Summary ---
    console.print()
    if errors:
        console.print(f"[green]✓ {done} captioned[/green]  [red]✗ {len(errors)} failed[/red]")
        for name, err in errors:
            console.print(f"  [red]{name}:[/red] {err}")
    else:
        console.print(Panel(
            f"[bold green]✓ {done} image(s) captioned successfully[/bold green]\n"
            f"[dim]Output: .txt files alongside each image[/dim]",
            style="green",
            expand=False,
        ))


if __name__ == "__main__":
    main()