#!/usr/bin/env python3
"""Gera a textura galáctica de alta resolução usada pelo wallpaper.

Fonte: "Gaia's sky in colour" (ESA/Gaia/DPAC, CC BY-SA 3.0 IGO),
8000x4000 em projeção Hammer e coordenadas galácticas. A saída é plate
carrée, adequada para amostragem direta por longitude/latitude no canvas.

O processamento é feito em blocos para não manter várias grades 6K inteiras
na memória. Requer Pillow e numpy.
"""

import io
import math
import urllib.request
from pathlib import Path

import numpy as np
from PIL import Image


URL = (
    "https://upload.wikimedia.org/wikipedia/commons/e/ea/"
    "Gaia%E2%80%99s_sky_in_colour_ESA393127.png"
)
CACHE = Path.home() / ".cache/skeledance/gaia_allsky.png"
OUT = Path(__file__).resolve().parent.parent / "plugin/contents/ui/milkyway.jpg"
SIZE = (6144, 3072)
BLOCK_ROWS = 96
QUALITY = 92


def fetch() -> bytes:
    if CACHE.exists():
        print(f"Usando cache {CACHE}")
        return CACHE.read_bytes()
    print(f"Baixando {URL} ...")
    req = urllib.request.Request(URL, headers={"User-Agent": "skeledance-build/2.0"})
    raw = urllib.request.urlopen(req, timeout=180).read()
    CACHE.parent.mkdir(parents=True, exist_ok=True)
    CACHE.write_bytes(raw)
    print(f"  {len(raw) / 1e6:.1f} MB (cacheado)")
    return raw


def main() -> None:
    Image.MAX_IMAGE_PIXELS = None
    src = np.asarray(Image.open(io.BytesIO(fetch())).convert("RGB"))
    sh, sw = src.shape[:2]
    tw, th = SIZE
    out = np.empty((th, tw, 3), dtype=np.uint8)

    # Longitude cresce para a esquerda para coincidir com a imagem Gaia.
    lon = np.deg2rad(180 - (np.arange(tw, dtype=np.float64) + 0.5) * 360 / tw)

    for y_start in range(0, th, BLOCK_ROWS):
        y_stop = min(th, y_start + BLOCK_ROWS)
        lat = np.deg2rad(
            90 - (np.arange(y_start, y_stop, dtype=np.float64) + 0.5) * 180 / th
        )
        lon_grid, lat_grid = np.meshgrid(lon, lat)

        denom = np.sqrt(1 + np.cos(lat_grid) * np.cos(lon_grid / 2))
        hx = 2 * math.sqrt(2) * np.cos(lat_grid) * np.sin(lon_grid / 2) / denom
        hy = math.sqrt(2) * np.sin(lat_grid) / denom
        sx = (1 - hx / (2 * math.sqrt(2))) * sw / 2
        sy = (1 - hy / math.sqrt(2)) * sh / 2

        # Amostragem bilinear preserva filamentos sem serrilhar a reprojeção.
        x0 = np.clip(np.floor(sx).astype(np.int32), 0, sw - 1)
        y0 = np.clip(np.floor(sy).astype(np.int32), 0, sh - 1)
        x1 = np.minimum(x0 + 1, sw - 1)
        y1 = np.minimum(y0 + 1, sh - 1)
        fx = (sx - x0)[..., None]
        fy = (sy - y0)[..., None]
        top = src[y0, x0] * (1 - fx) + src[y0, x1] * fx
        bottom = src[y1, x0] * (1 - fx) + src[y1, x1] * fx
        out[y_start:y_stop] = np.clip(top * (1 - fy) + bottom * fy, 0, 255)
        print(f"\rReprojetando: {y_stop * 100 // th:3d}%", end="", flush=True)

    print()
    Image.fromarray(out).save(OUT, "JPEG", quality=QUALITY, optimize=True)
    print(f"{OUT} ({OUT.stat().st_size / 1024 / 1024:.1f} MB, {tw}x{th})")


if __name__ == "__main__":
    main()