#!/usr/bin/env python3
"""Generate images via the Grok Build CLI /imagine command.

Requires an authenticated grok CLI (SuperGrok/X Premium+ subscription):
    grok login
"""

import argparse
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional


GROK = shutil.which("grok") or str(Path.home() / ".grok" / "bin" / "grok")
SESSION_ROOT = Path.home() / ".grok" / "sessions"
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}


def die(msg: str) -> None:
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(1)


def newest_image(since_ts: float) -> Optional[Path]:
    best = None
    if not SESSION_ROOT.is_dir():
        return None
    for p in SESSION_ROOT.rglob("*"):
        if not p.is_file() or p.suffix.lower() not in IMAGE_EXTS:
            continue
        try:
            mtime = p.stat().st_mtime
        except OSError:
            continue
        if mtime >= since_ts and (best is None or mtime > best.stat().st_mtime):
            best = p
    return best


def path_from_output(text: str) -> Optional[Path]:
    m = re.search(r"`([^`]+\.(?:jpg|jpeg|png|webp))`", text, re.I)
    if not m:
        m = re.search(r"([\w./\\-]+\.(?:jpg|jpeg|png|webp))", text, re.I)
    if not m:
        return None
    candidate = Path(m.group(1))
    if candidate.is_file():
        return candidate.resolve()
    # Output paths are usually relative to the session dir; resolve via search below.
    return None


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate images with Grok Imagine through the grok CLI"
    )
    parser.add_argument("prompt", help="text description of the image")
    parser.add_argument("-a", "--aspect-ratio", help="e.g. 1:1, 16:9, 9:16, 4:3")
    parser.add_argument("-o", "--output", help="copy the result into this directory")
    args = parser.parse_args()

    if not shutil.which(GROK):
        die("grok CLI not found. Install with: curl -fsSL https://x.ai/cli/install.sh | bash")

    full_prompt = f"/imagine {args.prompt}"
    if args.aspect_ratio:
        full_prompt += f", aspect ratio {args.aspect_ratio}"

    before = time.time()
    proc = subprocess.run(
        [GROK, "-p", full_prompt, "--no-auto-update"],
        capture_output=True,
        text=True,
        timeout=600,
    )
    output = proc.stdout + proc.stderr
    print(output)
    if proc.returncode != 0:
        die(f"grok exited with code {proc.returncode}")

    image = newest_image(before - 2) or path_from_output(output)
    if image is None:
        die("could not locate the generated image (check the output above)")

    image = image.resolve()
    if args.output:
        outdir = Path(args.output)
        outdir.mkdir(parents=True, exist_ok=True)
        dest = outdir / image.name
        if dest.exists():
            dest = outdir / f"{image.stem}-{int(time.time())}{image.suffix}"
        shutil.copy2(image, dest)
        image = dest.resolve()
    print(image)


if __name__ == "__main__":
    main()