#!/usr/bin/env python3

# Download NSIDC Antarctic ice shelf imagery for the selected summer window.
#
# Processing:
#   - Traverse ice shelf and sensor folders in the NSIDC-0102 dataset
#   - Download January 10–31 and all February TIFF imagery
#   - Optionally restrict downloads to an inclusive overall start/end date
#   - Exclude filenames containing "_therm"
#   - Save imagery to cwd/imagery/orig/<ice_shelf_name>/
#   - Skip files that already exist and are non-empty
#
# Earthdata credentials:
#   cwd/earthdata_credentials.txt
#
# Required format:
#   username=...
#   password=...
#
# Install:
#   pip install requests beautifulsoup4

from __future__ import annotations

import os
import re
import sys
import time
from datetime import date
from http.cookiejar import MozillaCookieJar
from pathlib import Path
from urllib.parse import urljoin, urlparse

import requests
from bs4 import BeautifulSoup


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

BASE_DIR = Path.cwd()
IMAGERY_DIR = BASE_DIR / "imagery"

BASE_URL = (
    "https://daacdata.apps.nsidc.org/pub/DATASETS/nsidc0102_antarctic_iceshelf_images_v2/"
)

CRED_FILE = BASE_DIR / "earthdata_credentials.txt"
COOKIE_JAR_FILE = BASE_DIR / "urs_cookies.txt"

# Summer window.
JAN_START_DAY = 10
JAN_END_DAY = 31

# Optional overall date limits, inclusive.
# Set either or both to None for no limit.
START_DATE: date | None = None
END_DATE: date | None = None

# Examples:
# START_DATE = date(2000, 1, 1)
# END_DATE = date(2020, 12, 31)

# Exclude matching filenames, case-insensitive.
EXCLUDE_SUBSTRINGS = ["_therm"]

TIMEOUT = 60
RETRIES = 5
SLEEP_SEC = 0.2

DATE_RE = re.compile(r"_(\d{8})_")


# =============================================================================
# AUTHENTICATION
# =============================================================================

def load_credentials(path: Path) -> tuple[str, str]:
    user = pwd = None

    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()

            if not line or line.startswith("#"):
                continue

            if line.lower().startswith("username="):
                user = line.split("=", 1)[1].strip()
            elif line.lower().startswith("password="):
                pwd = line.split("=", 1)[1].strip()

    if not user or not pwd:
        raise ValueError(
            f"{path.name} must contain username=... and password=..."
        )

    return user, pwd


def make_session() -> requests.Session:
    if not CRED_FILE.exists():
        raise FileNotFoundError(
            f"Missing {CRED_FILE.name} in {BASE_DIR}"
        )

    session = requests.Session()
    session.headers.update({
        "User-Agent": "nsidc0102-tif-downloader/5.0"
    })

    jar = MozillaCookieJar(str(COOKIE_JAR_FILE))

    if COOKIE_JAR_FILE.exists():
        try:
            jar.load(ignore_discard=True, ignore_expires=True)
        except Exception:
            pass

    session.cookies = jar
    return session


def save_cookies(session: requests.Session) -> None:
    if isinstance(session.cookies, MozillaCookieJar):
        session.cookies.save(
            ignore_discard=True,
            ignore_expires=True,
        )


# =============================================================================
# URL HELPERS
# =============================================================================

DATASET_HOST = urlparse(BASE_URL).netloc.lower()
DATASET_BASE_PATH = urlparse(BASE_URL).path.rstrip("/") + "/"


def in_dataset_tree(url: str) -> bool:
    parsed = urlparse(url)

    return (
        parsed.netloc.lower() == DATASET_HOST
        and parsed.path.startswith(DATASET_BASE_PATH)
    )


def is_dir_url(url: str) -> bool:
    path = urlparse(url).path

    if path.endswith("/"):
        return True

    return "." not in path.rsplit("/", 1)[-1]


def looks_like_login_html(html: str) -> bool:
    html = html.lower()

    return (
        "urs.earthdata.nasa.gov" in html
        and "login" in html
    )


def fetch_manual_redirect(
    url: str,
    session: requests.Session,
    auth: tuple[str, str],
) -> requests.Response:

    current = url

    for _ in range(12):
        host = urlparse(current).netloc.lower()
        use_auth = (
            auth if host == "urs.earthdata.nasa.gov" else None
        )

        response = session.get(
            current,
            timeout=TIMEOUT,
            allow_redirects=False,
            auth=use_auth,
        )

        if response.is_redirect or response.is_permanent_redirect:
            location = response.headers.get("Location")

            if not location:
                response.raise_for_status()
                save_cookies(session)
                return response

            current = urljoin(current, location)
            continue

        response.raise_for_status()
        save_cookies(session)
        return response

    raise RuntimeError(f"Too many redirects for {url}")


def get_text(
    url: str,
    session: requests.Session,
    auth: tuple[str, str],
) -> str:

    for attempt in range(1, RETRIES + 1):
        try:
            response = fetch_manual_redirect(
                url,
                session,
                auth,
            )

            text = response.text

            if looks_like_login_html(text):
                raise RuntimeError(
                    "Got URS login HTML "
                    "(auth not established / terms not accepted)."
                )

            return text

        except (requests.RequestException, RuntimeError):
            if attempt == RETRIES:
                raise

            time.sleep(1.2 * attempt)


def list_dir(
    url: str,
    session: requests.Session,
    auth: tuple[str, str],
) -> list[str]:

    html = get_text(url, session, auth)
    soup = BeautifulSoup(html, "html.parser")

    urls = []

    for link in soup.find_all("a", href=True):
        href = link["href"].strip()

        if href in ("../", "./") or href.startswith("?"):
            continue

        absolute_url = urljoin(url, href)

        if in_dataset_tree(absolute_url):
            urls.append(absolute_url)

    return urls


def one_level_children(
    entries: list[str],
    parent_url: str,
) -> list[str]:

    base_path = urlparse(parent_url).path.rstrip("/") + "/"
    out = []

    for url in entries:
        path = urlparse(url).path

        if not path.startswith(base_path):
            continue

        relative = path[len(base_path):].strip("/")

        if relative and "/" not in relative:
            out.append(url)

    return sorted(set(out))


# =============================================================================
# FILTERING
# =============================================================================

def is_tif_url(url: str) -> bool:
    path = urlparse(url).path.lower()
    return path.endswith(".tif") or path.endswith(".tiff")


def parse_yyyymmdd_from_filename(filename: str) -> date | None:
    match = DATE_RE.search(filename)

    if not match:
        return None

    value = match.group(1)

    return date(
        int(value[0:4]),
        int(value[4:6]),
        int(value[6:8]),
    )


def in_summer_window(dt: date) -> bool:
    if dt.month == 1:
        return JAN_START_DAY <= dt.day <= JAN_END_DAY

    if dt.month == 2:
        return True

    return False


def should_download_url(url: str) -> bool:
    if not is_tif_url(url):
        return False

    filename = os.path.basename(urlparse(url).path)
    lower_name = filename.lower()

    if any(
        substring in lower_name
        for substring in EXCLUDE_SUBSTRINGS
    ):
        return False

    dt = parse_yyyymmdd_from_filename(filename)

    if dt is None:
        return False

    # Existing seasonal filter.
    if not in_summer_window(dt):
        return False

    # Optional overall date limits.
    if START_DATE is not None and dt < START_DATE:
        return False

    if END_DATE is not None and dt > END_DATE:
        return False

    return True


# =============================================================================
# DOWNLOAD
# =============================================================================

def download(
    url: str,
    dest: Path,
    session: requests.Session,
    auth: tuple[str, str],
) -> None:

    dest.parent.mkdir(parents=True, exist_ok=True)

    if dest.exists() and dest.stat().st_size > 0:
        print(f"SKIP (exists): {dest}")
        return

    tmp = dest.with_suffix(dest.suffix + ".part")

    for attempt in range(1, RETRIES + 1):
        try:
            current = url

            for _ in range(12):
                host = urlparse(current).netloc.lower()
                use_auth = (
                    auth if host == "urs.earthdata.nasa.gov" else None
                )

                response = session.get(
                    current,
                    timeout=TIMEOUT,
                    allow_redirects=False,
                    auth=use_auth,
                    stream=True,
                )

                if response.is_redirect or response.is_permanent_redirect:
                    location = response.headers.get("Location")
                    response.close()

                    if not location:
                        break

                    current = urljoin(current, location)
                    continue

                response.raise_for_status()
                save_cookies(session)

                content_type = (
                    response.headers.get("content-type") or ""
                ).lower()

                if "text/html" in content_type:
                    snippet = (
                        response.text or ""
                    )[:200].replace("\n", " ")

                    response.close()

                    raise RuntimeError(
                        f"Got HTML not TIFF. "
                        f"Snippet: {snippet!r}"
                    )

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

                response.close()
                tmp.replace(dest)

                print(f"OK: {dest}")
                return

            raise RuntimeError(
                "Too many redirects while downloading"
            )

        except Exception as error:
            print(
                f"Retry {attempt}/{RETRIES} failed: "
                f"{url} -> {error}"
            )

            try:
                if tmp.exists():
                    tmp.unlink()
            except Exception:
                pass

            time.sleep(1.2 * attempt)

    raise RuntimeError(
        f"Failed after {RETRIES} attempts: {url}"
    )


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

def main() -> None:
    if (
        START_DATE is not None
        and END_DATE is not None
        and START_DATE > END_DATE
    ):
        raise ValueError(
            "START_DATE must be earlier than or equal to END_DATE"
        )

    user, pwd = load_credentials(CRED_FILE)
    auth = (user, pwd)

    session = make_session()

    print(f"Listing ice shelves under: {BASE_URL}")

    top = list_dir(
        BASE_URL,
        session,
        auth,
    )

    shelves = one_level_children(
        top,
        BASE_URL,
    )

    # Skip dataset grid-overlay folder.
    shelves = [
        url
        for url in shelves
        if Path(
            urlparse(url).path.rstrip("/")
        ).name != "0102_grid_overlays"
    ]

    print(f"Found {len(shelves)} ice shelf folders.")

    print(
        f"Seasonal filter: Jan "
        f"{JAN_START_DAY:02d}–{JAN_END_DAY:02d} "
        f"and all February"
    )

    if START_DATE is not None:
        print(f"Overall start date: {START_DATE}")

    if END_DATE is not None:
        print(f"Overall end date:   {END_DATE}")

    for shelf_url in shelves:
        shelf = Path(
            urlparse(shelf_url).path.rstrip("/")
        ).name

        print(f"\n=== {shelf} ===")

        time.sleep(SLEEP_SEC)

        shelf_entries = list_dir(
            shelf_url,
            session,
            auth,
        )

        sensors = [
            url
            for url in one_level_children(
                shelf_entries,
                shelf_url,
            )
            if is_dir_url(url)
        ]

        for sensor_url in sensors:
            sensor = Path(
                urlparse(sensor_url).path.rstrip("/")
            ).name

            images_url = urljoin(
                sensor_url.rstrip("/") + "/",
                "images/",
            )

            try:
                time.sleep(SLEEP_SEC)

                image_entries = list_dir(
                    images_url,
                    session,
                    auth,
                )

            except requests.HTTPError as error:
                if getattr(
                    error.response,
                    "status_code",
                    None,
                ) == 404:
                    continue

                print(
                    f"  WARN: {images_url} -> {error}"
                )
                continue

            except Exception as error:
                print(
                    f"  WARN: {images_url} -> {error}"
                )
                continue

            tif_urls = sorted([
                url
                for url in image_entries
                if should_download_url(url)
            ])

            if not tif_urls:
                continue

            print(
                f"  {sensor}: "
                f"{len(tif_urls)} tif(s) matching filters"
            )

            for tif_url in tif_urls:
                filename = os.path.basename(
                    urlparse(tif_url).path
                )

                dest = (
                    IMAGERY_DIR
                    / shelf
                    / "orig"
                    / filename
                )

                time.sleep(SLEEP_SEC)

                download(
                    tif_url,
                    dest,
                    session,
                    auth,
                )

    print("\nDone.")


if __name__ == "__main__":
    try:
        main()

    except Exception as error:
        print("\nERROR:", error)
        sys.exit(1)