#!/usr/bin/env python3
"""Generate an image through Pollinations.ai (free/no key) and save locally.

Usage:
  python3 ~/.hermes/scripts/pollinations_image.py --prompt "..." --aspect portrait --out /tmp/image.jpg

Notes:
- Pollinations is a free public service, so quality/latency may vary.
- The service often adds a bottom-right watermark. By default this script crops a small bottom strip.
"""
import argparse, pathlib, urllib.parse, urllib.request, time, sys

ASPECTS = {
    'portrait': (768, 1024),
    'square': (1024, 1024),
    'landscape': (1280, 720),
}

def crop_watermark(path: pathlib.Path):
    try:
        from PIL import Image
    except Exception:
        return False, 'Pillow not installed; kept original image with possible watermark.'
    try:
        im = Image.open(path)
        w, h = im.size
        # Pollinations watermark sits near bottom-right. Crop a modest bottom band.
        crop_h = max(40, min(90, int(h * 0.085)))
        im.crop((0, 0, w, h - crop_h)).save(path, quality=95)
        return True, f'Cropped bottom {crop_h}px to remove possible watermark.'
    except Exception as e:
        return False, f'Could not crop watermark: {e}'

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prompt', required=True)
    ap.add_argument('--aspect', choices=ASPECTS, default='portrait')
    ap.add_argument('--out', default='')
    ap.add_argument('--seed', default='')
    ap.add_argument('--raw', action='store_true', help='Do not crop watermark')
    args = ap.parse_args()
    width, height = ASPECTS[args.aspect]
    seed = args.seed or str(int(time.time()) % 10_000_000)
    prompt = args.prompt.strip()
    params = urllib.parse.urlencode({
        'width': width,
        'height': height,
        'seed': seed,
        'nologo': 'true',
        'private': 'true',
        'enhance': 'true',
        'safe': 'true',
    })
    url = 'https://image.pollinations.ai/prompt/' + urllib.parse.quote(prompt) + '?' + params
    out = pathlib.Path(args.out or f'/home/spazio/.hermes/cache/images/pollinations_{int(time.time())}.jpg')
    out.parent.mkdir(parents=True, exist_ok=True)
    req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0 (Hermes Agent)'})
    with urllib.request.urlopen(req, timeout=180) as r:
        ctype = r.headers.get('content-type','')
        data = r.read()
    if not data or not ctype.startswith('image/'):
        raise RuntimeError(f'Pollinations returned {ctype}: {data[:200]!r}')
    out.write_bytes(data)
    note = 'Original image kept.'
    if not args.raw:
        _, note = crop_watermark(out)
    print(str(out))
    print(note, file=sys.stderr)

if __name__ == '__main__':
    main()