#!/usr/bin/env python3
"""Download and spatially compress local solar-system texture assets.

The runtime uses 4096x2048 optimized images under the WinUI Assets folder.
Original 8K sources are archived outside the app assets so they are not copied
into the build output.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from shutil import copyfile

from PIL import Image, ImageEnhance, ImageFile, ImageFilter, ImageOps

ImageFile.MAXBLOCK = 128 * 1024 * 1024


ROOT = Path(__file__).resolve().parents[1]
SOURCE_DIR = ROOT / "assets" / "source-textures" / "solar"
OUTPUT_DIR = ROOT / "src" / "box-winUI" / "Assets" / "home-globe" / "textures" / "solar"
MANIFEST_PATH = OUTPUT_DIR / "solar-texture-manifest.json"
TARGET_SIZE = (4096, 2048)
USER_AGENT = "YMhutBoxTexturePipeline/1.0 (+local asset preparation)"


@dataclass(frozen=True)
class TextureJob:
    body_id: str
    source_name: str
    output_name: str
    url: str | None
    credit: str
    mode: str = "rgb"
    quality: int = 94
    contrast: float = 1.02
    sharpness: float = 1.08


JOBS = [
    TextureJob(
        "mercury",
        "8k_mercury.jpg",
        "mercury.jpg",
        "https://www.solarsystemscope.com/textures/download/8k_mercury.jpg",
        "Solar System Scope 8K Mercury texture, NASA/JPL-derived public imagery.",
        contrast=1.06,
        sharpness=1.14,
    ),
    TextureJob(
        "venus",
        "8k_venus_surface.jpg",
        "venus.jpg",
        "https://www.solarsystemscope.com/textures/download/8k_venus_surface.jpg",
        "Solar System Scope 8K Venus surface texture, NASA/JPL-derived public imagery.",
        contrast=1.04,
        sharpness=1.1,
    ),
    TextureJob(
        "earth",
        "8k_earth_daymap.jpg",
        "earth.jpg",
        "https://www.solarsystemscope.com/textures/download/8k_earth_daymap.jpg",
        "Solar System Scope 8K Earth day map based on NASA Blue Marble imagery.",
        contrast=1.03,
        sharpness=1.1,
    ),
    TextureJob(
        "earth-clouds",
        "8k_earth_clouds.jpg",
        "earth-clouds.png",
        "https://www.solarsystemscope.com/textures/download/8k_earth_clouds.jpg",
        "Solar System Scope 8K Earth cloud map based on NASA imagery.",
        mode="cloud-alpha",
        quality=100,
        contrast=1.18,
        sharpness=1.08,
    ),
    TextureJob(
        "mars",
        "8k_mars.jpg",
        "mars.jpg",
        "https://www.solarsystemscope.com/textures/download/8k_mars.jpg",
        "Solar System Scope 8K Mars texture, NASA/JPL-derived public imagery.",
        contrast=1.06,
        sharpness=1.15,
    ),
    TextureJob(
        "jupiter",
        "8k_jupiter.jpg",
        "jupiter.jpg",
        "https://www.solarsystemscope.com/textures/download/8k_jupiter.jpg",
        "Solar System Scope high-resolution Jupiter texture, NASA/JPL-derived public imagery.",
        contrast=1.04,
        sharpness=1.12,
    ),
    TextureJob(
        "saturn",
        "8k_saturn.jpg",
        "saturn.jpg",
        "https://www.solarsystemscope.com/textures/download/8k_saturn.jpg",
        "Solar System Scope high-resolution Saturn texture, NASA/JPL-derived public imagery.",
        contrast=1.04,
        sharpness=1.12,
    ),
    TextureJob(
        "uranus",
        "local_uranus_4k.png",
        "uranus.jpg",
        "local-png:src/box-winUI/Assets/home-globe/textures/solar/uranus.jpg",
        "Local lossless 4K Uranus fallback source, retained because the public 8K direct link is unavailable.",
        contrast=1.08,
        sharpness=1.1,
    ),
    TextureJob(
        "neptune",
        "local_neptune_4k.png",
        "neptune.jpg",
        "local-png:src/box-winUI/Assets/home-globe/textures/solar/neptune.jpg",
        "Local lossless 4K Neptune fallback source, retained because the public 8K direct link is unavailable.",
        contrast=1.08,
        sharpness=1.1,
    ),
]


def download(url: str, destination: Path, force: bool) -> None:
    if destination.exists() and destination.stat().st_size > 0 and not force:
        return

    destination.parent.mkdir(parents=True, exist_ok=True)
    if url.startswith("local:"):
        seed_path = ROOT / url.removeprefix("local:")
        if not seed_path.exists():
            raise FileNotFoundError(seed_path)
        copyfile(seed_path, destination)
        mb = destination.stat().st_size / 1024 / 1024
        print(f"Archived local fallback {destination.name}: {mb:.2f} MB")
        return

    if url.startswith("local-png:"):
        seed_path = ROOT / url.removeprefix("local-png:")
        if not seed_path.exists():
            raise FileNotFoundError(seed_path)
        with Image.open(seed_path) as image:
            image.convert("RGB").save(destination, "PNG", optimize=True, compress_level=9)
        mb = destination.stat().st_size / 1024 / 1024
        print(f"Archived lossless local fallback {destination.name}: {mb:.2f} MB")
        return

    tmp = destination.with_suffix(destination.suffix + ".part")
    last_error: Exception | None = None
    for attempt in range(1, 5):
        request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        try:
            with urllib.request.urlopen(request, timeout=180) as response:
                total = int(response.headers.get("Content-Length") or 0)
                loaded = 0
                started = time.monotonic()
                with tmp.open("wb") as out:
                    while True:
                        chunk = response.read(1024 * 512)
                        if not chunk:
                            break
                        out.write(chunk)
                        loaded += len(chunk)
                        if total:
                            percent = loaded / total * 100
                            sys.stdout.write(f"\rDownloading {destination.name}: {percent:5.1f}%")
                            sys.stdout.flush()
                elapsed = max(time.monotonic() - started, 0.001)
            if tmp.stat().st_size == 0:
                raise RuntimeError("downloaded file is empty")
            tmp.replace(destination)
            mb = destination.stat().st_size / 1024 / 1024
            print(f"\rDownloaded {destination.name}: {mb:.2f} MB in {elapsed:.1f}s")
            return
        except Exception as exc:
            last_error = exc
            wait_seconds = min(18, attempt * 4)
            print(f"\nRetrying {destination.name} after {type(exc).__name__}: attempt {attempt}/4")
            time.sleep(wait_seconds)

    raise RuntimeError(f"download failed for {destination.name}") from last_error


def resize_source(image: Image.Image, target_size: tuple[int, int]) -> Image.Image:
    image = ImageOps.exif_transpose(image)
    if image.size != target_size:
        image = image.resize(target_size, Image.Resampling.LANCZOS)
    return image


def enhance_rgb(image: Image.Image, job: TextureJob) -> Image.Image:
    image = image.convert("RGB")
    if job.contrast != 1:
        image = ImageEnhance.Contrast(image).enhance(job.contrast)
    if job.sharpness != 1:
        image = ImageEnhance.Sharpness(image).enhance(job.sharpness)
    return image.filter(ImageFilter.UnsharpMask(radius=0.9, percent=80, threshold=2))


def build_cloud_alpha(image: Image.Image, job: TextureJob) -> Image.Image:
    image = enhance_rgb(image, job)
    luminance = image.convert("L")
    alpha = ImageEnhance.Contrast(luminance).enhance(1.35)
    white = Image.new("RGB", image.size, (255, 255, 255))
    rgba = Image.merge("RGBA", (*white.split(), alpha))
    return rgba


def source_path_for(job: TextureJob) -> Path:
    return SOURCE_DIR / job.source_name


def output_path_for(job: TextureJob) -> Path:
    return OUTPUT_DIR / job.output_name


def optimize(job: TextureJob) -> dict[str, object]:
    source_path = SOURCE_DIR / job.source_name
    output_path = OUTPUT_DIR / job.output_name
    print(f"Optimizing {job.body_id}: {source_path.name} -> {output_path.name}")

    with Image.open(source_path) as original:
        source_size = original.size
        source_mode = original.mode
        source_raw_bytes = source_size[0] * source_size[1] * max(len(original.getbands()), 3)
        image = resize_source(original, TARGET_SIZE)

    if job.mode == "cloud-alpha":
        image = build_cloud_alpha(image, job)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        image.save(output_path, "PNG", optimize=True, compress_level=9)
    else:
        image = enhance_rgb(image, job)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        image.save(
            output_path,
            "JPEG",
            quality=job.quality,
            optimize=True,
            progressive=True,
            subsampling=0,
        )

    return {
        "id": job.body_id,
        "sourceUrl": job.url,
        "sourceCredit": job.credit,
        "sourceArchivePath": str(source_path.relative_to(ROOT)).replace("\\", "/"),
        "sourceBytes": source_path.stat().st_size,
        "sourceImageSize": {"width": source_size[0], "height": source_size[1], "mode": source_mode},
        "sourceRawBytes": source_raw_bytes,
        "optimizedPath": str(output_path.relative_to(ROOT)).replace("\\", "/"),
        "optimizedBytes": output_path.stat().st_size,
        "optimizedImageSize": {"width": TARGET_SIZE[0], "height": TARGET_SIZE[1]},
        "compression": "spatial Lanczos resize to 4096x2048 plus high-quality JPEG/PNG optimization",
    }


def write_manifest(entries: list[dict[str, object]]) -> None:
    total_source_bytes = sum(int(entry["sourceBytes"]) for entry in entries)
    total_source_raw_bytes = sum(int(entry["sourceRawBytes"]) for entry in entries)
    total_optimized_bytes = sum(int(entry["optimizedBytes"]) for entry in entries)
    manifest = {
        "version": 1,
        "generatedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "policy": {
            "sourceArchive": "Original high-resolution sources are local but outside WinUI copied assets.",
            "runtimeAssets": "Runtime textures are local 4096x2048 spatially compressed images.",
            "minimumSourceRawBytes": 60 * 1024 * 1024,
        },
        "totals": {
            "sourceBytes": total_source_bytes,
            "sourceRawBytes": total_source_raw_bytes,
            "optimizedBytes": total_optimized_bytes,
            "sourceBytesMB": round(total_source_bytes / 1024 / 1024, 2),
            "sourceRawBytesMB": round(total_source_raw_bytes / 1024 / 1024, 2),
            "optimizedBytesMB": round(total_optimized_bytes / 1024 / 1024, 2),
        },
        "sources": [
            "https://www.solarsystemscope.com/textures/",
            "https://visibleearth.nasa.gov/collection/1484/blue-marble",
            "https://maps.jpl.nasa.gov/tmaps/",
        ],
        "textures": entries,
    }
    MANIFEST_PATH.write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")

    raw_mb = manifest["totals"]["sourceRawBytesMB"]
    source_mb = manifest["totals"]["sourceBytesMB"]
    optimized_mb = manifest["totals"]["optimizedBytesMB"]
    print(f"Source archive: {source_mb:.2f} MB downloaded")
    print(f"Source raw pixels: {raw_mb:.2f} MB before spatial compression")
    print(f"Optimized runtime: {optimized_mb:.2f} MB")
    if total_source_raw_bytes < manifest["policy"]["minimumSourceRawBytes"]:
        raise RuntimeError("Source raw texture budget is below 60 MB")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--force", action="store_true", help="redownload and rebuild all textures")
    parser.add_argument("--download-workers", type=int, default=4, help="parallel download workers")
    args = parser.parse_args()

    workers = max(1, min(args.download_workers, len(JOBS)))
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(download, job.url or "", source_path_for(job), args.force): job
            for job in JOBS
        }
        for future in as_completed(futures):
            job = futures[future]
            try:
                future.result()
            except Exception as exc:
                raise RuntimeError(f"Failed to download {job.source_name}") from exc

    entries = [optimize(job) for job in JOBS]
    write_manifest(entries)


if __name__ == "__main__":
    main()
