#!/usr/bin/env python3

# Download monthly Antarctic sea-ice extent polygon shapefiles from NSIDC.
#
# - Traverses the 12 monthly folders in the NSIDC G02135 archive
# - Downloads polygon ZIPs directly to cwd/input
# - Skips existing non-empty files
#
# Install:
#   pip install requests beautifulsoup4

from pathlib import Path
import re
from urllib.parse import urljoin, urlparse

import requests
from bs4 import BeautifulSoup


# =============================================================================
# CONFIGURATION
# =============================================================================

BASE_DIR = Path.cwd()
INPUT_DIR = BASE_DIR / "input"

BASE_URL = (
    "https://noaadata.apps.nsidc.org/NOAA/G02135/south/monthly/shapefiles/shp_extent/"
)

TIMEOUT = 60
MONTH_RE = re.compile(r"^\d{2}_[A-Za-z]{3}$")


# =============================================================================
# HELPERS
# =============================================================================

def get_links(session: requests.Session, url: str) -> list[str]:
    response = session.get(url, timeout=TIMEOUT)
    response.raise_for_status()

    soup = BeautifulSoup(response.text, "html.parser")

    return [
        urljoin(url, link["href"])
        for link in soup.find_all("a", href=True)
    ]


def is_month_folder(url: str) -> bool:
    folder = Path(urlparse(url).path.rstrip("/")).name
    return url.endswith("/") and bool(MONTH_RE.fullmatch(folder))


def download_file(session: requests.Session, url: str) -> None:
    filename = Path(urlparse(url).path).name
    output_path = INPUT_DIR / filename

    if output_path.exists() and output_path.stat().st_size > 0:
        print(f"[SKIP] {output_path.relative_to(BASE_DIR)}")
        return

    temp_path = output_path.with_suffix(output_path.suffix + ".part")

    print(f"[DOWNLOAD] {filename}")

    with session.get(url, stream=True, timeout=TIMEOUT) as response:
        response.raise_for_status()

        with open(temp_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    f.write(chunk)

    temp_path.replace(output_path)


# =============================================================================
# MAIN
# =============================================================================

def main() -> None:
    INPUT_DIR.mkdir(parents=True, exist_ok=True)

    with requests.Session() as session:
        month_urls = sorted(
            url
            for url in get_links(session, BASE_URL)
            if is_month_folder(url)
        )

        if not month_urls:
            raise RuntimeError(f"No monthly folders found at:\n{BASE_URL}")

        print(f"Found {len(month_urls)} monthly folders")

        for month_url in month_urls:
            month = Path(urlparse(month_url).path.rstrip("/")).name
            print(f"\n[{month}]")

            for file_url in get_links(session, month_url):
                filename = Path(urlparse(file_url).path).name.lower()

                if filename.endswith(".zip") and "_polygon_" in filename:
                    download_file(session, file_url)

    print("\n[OK] Done")


if __name__ == "__main__":
    main()