#!/usr/bin/env python3
"""
generate_cover.py — Generate Imagen 4.0 blog cover images for HoffDesk blog.

Usage:
  python3 generate_cover.py <slug> "<title>" "<excerpt>"

Dependencies: requests, Pillow (python3 -m pip install --break-system-packages requests Pillow)
Auth: Requires VERTEX_AI_KEY and VERTEX_AI_PROJECT env vars.
"""

import os
import sys
import json
import base64
import requests
from PIL import Image
from io import BytesIO

# ── Config ───────────────────────────────────────────────────────────────────
API_KEY = os.environ.get("VERTEX_AI_KEY")
PROJECT = os.environ.get("VERTEX_AI_PROJECT", "gen-lang-client-0081729505")
LOCATION = "us-central1"
PUBLISHER = "google"
MODEL = "imagen-4.0-generate-001"
OUTPUT_DIR = os.environ.get(
    "COVER_OUTPUT_DIR",
    "/home/hoffmann_admin/hoffdesk/blog/static/images/posts/"
)
MAX_WIDTH = 800
JPEG_QUALITY = 85
SAMPLE_COUNT = 1

# ── Templates ────────────────────────────────────────────────────────────────

BASE_PROMPT = (
    "Editorial magazine illustration style, cinematic composition, "
    "high contrast, premium tech magazine aesthetic, clean digital art. "
    "No text, no logos, no watermarks, no brand names, no written words, "
    "no typography, no magazine mastheads, no labels anywhere in the image."
)

def build_prompt(title: str, excerpt: str) -> str:
    """Build a visual prompt from blog metadata."""
    concept = f"{title}: {excerpt}" if excerpt else title
    # Strip very long excerpts
    if len(concept) > 200:
        concept = concept[:197] + "..."
    return (
        f"{concept}. "
        f"Dark background, neon accent lighting, dramatic atmosphere. "
        f"{BASE_PROMPT}"
    )


# ── Generation ───────────────────────────────────────────────────────────────

def generate_image(prompt: str, slug: str) -> str | None:
    """Generate a cover image via Imagen 4.0, save as WebP + JPEG, return path."""
    if not API_KEY:
        print("❌ VERTEX_AI_KEY not set. Export it or set in .bashrc", file=sys.stderr)
        return None

    url = (
        f"https://{LOCATION}-aiplatform.googleapis.com/v1/"
        f"projects/{PROJECT}/locations/{LOCATION}/"
        f"publishers/{PUBLISHER}/models/{MODEL}:predict"
    )

    headers = {
        "x-goog-api-key": API_KEY,
        "Content-Type": "application/json",
    }

    payload = {
        "instances": [{"prompt": prompt}],
        "parameters": {
            "sampleCount": SAMPLE_COUNT,
            "aspectRatio": "4:3",
        },
    }

    print(f"  Generating image for: {slug}")
    resp = requests.post(url, headers=headers, json=payload, timeout=60)

    if resp.status_code != 200:
        err = resp.json().get("error", {}).get("message", resp.text[:200])
        print(f"❌ API error ({resp.status_code}): {err}", file=sys.stderr)
        return None

    data = resp.json()
    if "predictions" not in data or not data["predictions"]:
        print("❌ No predictions returned", file=sys.stderr)
        return None

    img_b64 = data["predictions"][0].get("bytesBase64Encoded")
    if not img_b64:
        print("❌ No image data in response", file=sys.stderr)
        return None

    # Decode
    img_bytes = base64.b64decode(img_b64)
    img = Image.open(BytesIO(img_bytes))

    # Resize
    new_w = min(MAX_WIDTH, img.width)
    new_h = int(img.height * (new_w / img.width))
    img_resized = img.resize((new_w, new_h), Image.LANCZOS)

    # Ensure output dir exists
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Save JPEG
    jpg_path = os.path.join(OUTPUT_DIR, f"{slug}-cover.jpg")
    img_resized.convert("RGB").save(jpg_path, "JPEG", quality=JPEG_QUALITY, optimize=True)

    # Save WebP (smaller for browsers that support it)
    webp_path = os.path.join(OUTPUT_DIR, f"{slug}-cover.webp")
    img_resized.convert("RGB").save(webp_path, "WEBP", quality=JPEG_QUALITY)

    # Save full-res PNG backup
    png_path = os.path.join(OUTPUT_DIR, f"{slug}-cover.png")
    with open(png_path, "wb") as f:
        f.write(img_bytes)

    jpg_size = os.path.getsize(jpg_path)
    print(f"  ✅ JPEG:  {jpg_path} ({jpg_size/1024:.0f} KB)")
    print(f"  ✅ WebP:  {webp_path}")
    print(f"  📦 PNG backup: {png_path}")

    # Return the relative path for use in frontmatter
    return f"/blog/static/images/posts/{slug}-cover.jpg"


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) < 3:
        print("Usage: python3 generate_cover.py <slug> \"<title>\" [\"<excerpt>\"]", file=sys.stderr)
        sys.exit(1)

    slug = sys.argv[1]
    title = sys.argv[2]
    excerpt = sys.argv[3] if len(sys.argv) > 3 else ""

    prompt = build_prompt(title, excerpt)
    print(f"📝 Prompt ({len(prompt)} chars):")
    print(f"   {prompt[:200]}...")
    print()

    cover_path = generate_image(prompt, slug)

    if cover_path:
        print(f"\n✅ Cover generated!")
        print(f"   Frontmatter: cover_image: {cover_path}")
    else:
        print("\n❌ Failed to generate cover", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()