#!/usr/bin/env python3
"""Provision official WALOUS GeoTIFF source rasters for bounded runtime analysis."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import shutil
import sys
from urllib.request import Request, urlopen
from zipfile import BadZipFile, ZipFile


SOURCES = {
    2020: {
        "url": (
            "https://geoservices.wallonie.be/geotraitement/spwdatadownload/results/"
            "47b348f1-6e7a-4baa-963c-0232a43c0cff/WAL_OCS_IA__2020_GEOTIFF_3812.zip"
        ),
        "expected_archive_bytes": 728_244_755,
        "target": "walous_land_cover_2020_3812.tif",
    },
    2023: {
        "url": (
            "https://geoservices.wallonie.be/geotraitement/spwdatadownload/results/"
            "4e780ba1-463c-478e-95df-d2f1963a150d/WAL_OCS_IA__2023_GEOTIFF_3812.zip"
        ),
        "expected_archive_bytes": 876_014_572,
        "target": "walous_land_cover_2023_3812.tif",
    },
}
MAX_ARCHIVE_BYTES = 1_000_000_000
MAX_EXTRACTED_BYTES = 50_000_000_000


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(8 * 1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def download(url: str, destination: Path, expected_bytes: int) -> str:
    temporary = destination.with_suffix(destination.suffix + ".part")
    temporary.unlink(missing_ok=True)
    digest = hashlib.sha256()
    received = 0
    request = Request(url, headers={"User-Agent": "GeoIntel/1.0 WALOUS-source-provisioner"})
    try:
        with urlopen(request, timeout=300) as response, temporary.open("wb") as output:
            content_length = int(response.headers.get("Content-Length") or 0)
            if content_length and content_length != expected_bytes:
                raise RuntimeError(f"official archive size changed: expected {expected_bytes}, advertised {content_length}")
            while chunk := response.read(8 * 1024 * 1024):
                received += len(chunk)
                if received > MAX_ARCHIVE_BYTES:
                    raise RuntimeError("official archive exceeds the governed 1 GB transfer limit")
                digest.update(chunk)
                output.write(chunk)
                if received % (128 * 1024 * 1024) < len(chunk):
                    print(f"  downloaded {received / 1024 / 1024:.0f} MiB", flush=True)
        if received != expected_bytes:
            raise RuntimeError(f"archive is incomplete: expected {expected_bytes} bytes, received {received}")
        temporary.replace(destination)
        return digest.hexdigest()
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def extract_single_geotiff(archive: Path, target: Path) -> None:
    try:
        with ZipFile(archive) as bundle:
            candidates = [item for item in bundle.infolist() if not item.is_dir() and item.filename.lower().endswith((".tif", ".tiff"))]
            if len(candidates) != 1:
                raise RuntimeError(f"archive must contain exactly one GeoTIFF, found {len(candidates)}")
            member = candidates[0]
            if member.file_size <= 0 or member.file_size > MAX_EXTRACTED_BYTES:
                raise RuntimeError(f"GeoTIFF uncompressed size is outside the governed limit: {member.file_size}")
            if Path(member.filename).name != member.filename.replace("\\", "/").split("/")[-1]:
                # Nested paths are accepted only by basename; extraction never trusts archive paths.
                pass
            temporary = target.with_suffix(target.suffix + ".part")
            temporary.unlink(missing_ok=True)
            with bundle.open(member) as source, temporary.open("wb") as output:
                shutil.copyfileobj(source, output, length=8 * 1024 * 1024)
            temporary.replace(target)
    except BadZipFile as exc:
        raise RuntimeError("official WALOUS archive is not a valid ZIP file") from exc


def validate_raster(path: Path) -> dict:
    try:
        import numpy as np
        import rasterio
        from rasterio.enums import Resampling
    except ImportError as exc:
        raise RuntimeError("rasterio and numpy are required to validate WALOUS sources") from exc
    with rasterio.open(path) as source:
        if source.crs is None or source.crs.to_epsg() != 3812:
            raise RuntimeError(f"WALOUS raster must use EPSG:3812, found {source.crs}")
        if source.count != 1:
            raise RuntimeError(f"WALOUS raster must have one band, found {source.count}")
        if not all(abs(abs(float(value)) - 1.0) <= 0.05 for value in source.res):
            raise RuntimeError(f"WALOUS raster must retain 1 m cells, found {source.res}")
        sample_height = min(2048, source.height)
        sample_width = min(2048, source.width)
        sample = source.read(1, out_shape=(sample_height, sample_width), masked=True, resampling=Resampling.nearest)
        values = np.unique(sample.compressed()).astype(int).tolist()
        unexpected = sorted(set(values) - set(range(1, 12)))
        if unexpected:
            raise RuntimeError(f"WALOUS sample contains classes outside 1-11: {unexpected}")
        return {
            "path": str(path),
            "crs": str(source.crs),
            "width": int(source.width),
            "height": int(source.height),
            "resolution": [float(value) for value in source.res],
            "bounds": [float(value) for value in source.bounds],
            "nodata": None if source.nodata is None else float(source.nodata),
            "sample_classes": values,
        }


def provision(year: int, destination: Path, force: bool) -> dict:
    source = SOURCES[year]
    target = destination / source["target"]
    archive = destination / f"{Path(source['target']).stem}.zip"
    if target.is_file() and not force:
        print(f"WALOUS {year}: validating existing source {target}")
        validation = validate_raster(target)
        digest = sha256_file(target)
    else:
        print(f"WALOUS {year}: downloading official archive")
        archive_digest = download(source["url"], archive, source["expected_archive_bytes"])
        print(f"WALOUS {year}: archive sha256 {archive_digest}")
        extract_single_geotiff(archive, target)
        validation = validate_raster(target)
        digest = sha256_file(target)
        archive.unlink(missing_ok=True)
    checksum_path = target.with_suffix(".sha256")
    checksum_path.write_text(f"{digest}  {target.name}\n", encoding="ascii")
    validation.update({"year": year, "sha256": digest, "download_url": source["url"]})
    print(f"WALOUS {year}: ready ({target.stat().st_size / 1024 / 1024:.0f} MiB)")
    return validation


def main() -> int:
    parser = argparse.ArgumentParser(description="Provision official WALOUS 2020/2023 GeoTIFF sources.")
    parser.add_argument("--years", nargs="+", type=int, choices=sorted(SOURCES), default=sorted(SOURCES))
    parser.add_argument("--destination", type=Path, default=Path("storage/source-cache/walous"))
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()
    args.destination.mkdir(parents=True, exist_ok=True)
    report = [provision(year, args.destination.resolve(), args.force) for year in args.years]
    report_path = args.destination / "provisioning-report.json"
    report_path.write_text(json.dumps({"sources": report}, indent=2) + "\n", encoding="utf-8")
    print(f"Provisioning report: {report_path}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"WALOUS_PROVISIONING_FAILED: {exc}", file=sys.stderr)
        raise SystemExit(1) from exc