#!/usr/bin/env python3
"""
Upscale MAP.png en 4 passes (4 quadrants) via Vertex AI Imagen upscale.

La map étant trop lourde pour l'API en un seul appel (> 10 Mo), on :
  1. Découpe en 4 quadrants de 2048×2048
  2. Upscale chaque quadrant ×2  →  4096×4096 chacun
  3. Recolle en une image finale 8192×8192

Usage :
  cd backend && python -m scripts.upscale_map
  cd backend && python -m scripts.upscale_map --factor x2 --input static/MAP.png --output static/MAP.png
"""
from __future__ import annotations

import argparse
import io
import sys
from pathlib import Path

from PIL import Image

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

from scripts.upscale_cover import upscale_image


def upscale_map(
    input_path: Path,
    output_path: Path,
    factor: str = "x2",
) -> None:
    src = Image.open(input_path).convert("RGB")
    w, h = src.size
    print(f"Image source : {w}×{h}  ({input_path.stat().st_size / 1024 / 1024:.1f} Mo)")

    half_w, half_h = w // 2, h // 2
    quadrants = [
        (0,      0,      half_w, half_h),   # top-left
        (half_w, 0,      w,      half_h),   # top-right
        (0,      half_h, half_w, h     ),   # bottom-left
        (half_w, half_h, w,      h     ),   # bottom-right
    ]
    labels = ["top-left", "top-right", "bottom-left", "bottom-right"]

    upscaled_quadrants: list[Image.Image] = []

    for i, (box, label) in enumerate(zip(quadrants, labels)):
        print(f"\n[{i+1}/4] Quadrant {label} {box}")

        quad_img = src.crop(box)
        tmp_in = Path(f"/tmp/map_quad_{i}.png")
        tmp_out = Path(f"/tmp/map_quad_{i}_up.png")

        buf = io.BytesIO()
        quad_img.save(buf, format="PNG", optimize=False)
        size_mb = len(buf.getvalue()) / 1024 / 1024
        print(f"       Taille quadrant : {size_mb:.1f} Mo", end="")
        if size_mb > 10:
            print("  ⚠ > 10 Mo, passage en JPEG lossy pour respecter la limite API")
            buf = io.BytesIO()
            quad_img.save(buf, format="JPEG", quality=92)
            tmp_in = Path(f"/tmp/map_quad_{i}.jpg")
        else:
            print()

        tmp_in.write_bytes(buf.getvalue())

        upscale_image(tmp_in, tmp_out, factor=factor)
        upscaled_quadrants.append(Image.open(tmp_out).convert("RGB"))

    uw, uh = upscaled_quadrants[0].size
    print(f"\nRecomposition : quadrants upscalés {uw}×{uh} → image finale {uw*2}×{uh*2}")
    final = Image.new("RGB", (uw * 2, uh * 2))
    final.paste(upscaled_quadrants[0], (0,   0))
    final.paste(upscaled_quadrants[1], (uw,  0))
    final.paste(upscaled_quadrants[2], (0,   uh))
    final.paste(upscaled_quadrants[3], (uw,  uh))

    output_path.parent.mkdir(parents=True, exist_ok=True)
    final.save(output_path, format="PNG", optimize=False)
    print(f"\nSauvegardé : {output_path} ({output_path.stat().st_size / 1024 / 1024:.1f} Mo)")


def main() -> None:
    default_map = _backend / "static" / "MAP.png"
    parser = argparse.ArgumentParser(description="Upscale MAP.png en 4 quadrants via Vertex AI")
    parser.add_argument("--input",  type=Path, default=default_map)
    parser.add_argument("--output", type=Path, default=None)
    parser.add_argument("--factor", choices=("x2", "x4"), default="x2")
    args = parser.parse_args()

    out = args.output or args.input
    if not args.input.is_file():
        raise SystemExit(f"Fichier introuvable : {args.input}")

    upscale_map(args.input, out, factor=args.factor)


if __name__ == "__main__":
    main()