#!/usr/bin/env python3
"""
Upscale l'image cover (lobby) via Vertex AI Imagen 4.0 upscale.

Prérequis : GCP_PROJECT_ID et GCP_LOCATION dans .env, et `gcloud auth application-default login`
(ou GOOGLE_APPLICATION_CREDENTIALS).

Usage :
  cd backend && python -m scripts.upscale_cover
  cd backend && python -m scripts.upscale_cover --factor x2 --input ../frontend/src/cover.jpg --output ../frontend/src/cover.jpg
"""
from __future__ import annotations

import argparse
import base64
import subprocess
import sys
from pathlib import Path

_backend = Path(__file__).resolve().parent.parent
if str(_backend) not in sys.path:
    sys.path.insert(0, str(_backend))

from config import GCP_PROJECT_ID, GCP_LOCATION


def get_access_token() -> str:
    """Token via gcloud (ADC)."""
    out = subprocess.run(
        ["gcloud", "auth", "print-access-token"],
        capture_output=True,
        text=True,
        check=False,
    )
    if out.returncode != 0:
        raise RuntimeError(
            "Échec gcloud auth. Lancez: gcloud auth application-default login"
        )
    return out.stdout.strip()


def upscale_image(
    image_path: Path,
    out_path: Path,
    factor: str = "x4",
    region: str | None = None,
    project_id: str | None = None,
) -> None:
    region = region or GCP_LOCATION
    project_id = project_id or GCP_PROJECT_ID
    if not project_id:
        raise SystemExit("Définir GCP_PROJECT_ID dans .env")

    data = image_path.read_bytes()
    if len(data) > 10 * 1024 * 1024:
        raise SystemExit("Image trop lourde (max 10 Mo)")

    b64 = base64.standard_b64encode(data).decode("ascii")
    token = get_access_token()
    url = (
        f"https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}"
        f"/locations/{region}/publishers/google/models/imagen-4.0-upscale-preview:predict"
    )
    import httpx

    payload = {
        "instances": [
            {"prompt": "Upscale the image", "image": {"bytesBase64Encoded": b64}}
        ],
        "parameters": {
            "mode": "upscale",
            "upscaleConfig": {"upscaleFactor": factor},
            "outputOptions": {"mimeType": "image/png"},
        },
    }
    resp = httpx.post(
        url,
        json=payload,
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        },
        timeout=120.0,
    )
    resp.raise_for_status()
    body = resp.json()
    preds = body.get("predictions") or []
    if not preds or "bytesBase64Encoded" not in preds[0]:
        raise SystemExit("Réponse Vertex sans image: " + str(body)[:500])
    out_b64 = preds[0]["bytesBase64Encoded"]
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_bytes(base64.standard_b64decode(out_b64))
    print(f"Upscale {factor} OK → {out_path} ({out_path.stat().st_size / 1024:.1f} Ko)")


def main() -> None:
    parser = argparse.ArgumentParser(description="Upscale cover image via Vertex AI Imagen")
    parser.add_argument(
        "--input",
        type=Path,
        default=_backend.parent / "frontend" / "src" / "cover.jpg",
        help="Image source",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Image de sortie (défaut: --input)",
    )
    parser.add_argument(
        "--factor",
        choices=("x2", "x3", "x4"),
        default="x4",
        help="Facteur d'upscale (défaut: x4)",
    )
    args = parser.parse_args()
    out = args.output or args.input
    if not args.input.is_file():
        raise SystemExit(f"Fichier introuvable: {args.input}")
    upscale_image(args.input, out, factor=args.factor)


if __name__ == "__main__":
    main()