# %%
# Download time-varying forcing data
# ----------------------------------
# This is the third step of the Great Barrier Reef workflow. It downloads the
# time-dependent ocean and atmospheric fields used by prepro.py: 3D currents,
# temperature, salinity, sea surface elevation, wind, heat fluxes, precipitation,
# and sea-surface temperature.
#
# Most products are downloaded through NCSS/THREDDS requests in monthly or
# shorter chunks. The SST product is downloaded through the Copernicus Marine
# client. Existing files are checked before downloading so interrupted runs can
# be resumed safely.

import os
import sys
import time
import threading
from datetime import datetime, timedelta
from typing import List, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
import copernicusmarine

import requests
import xarray as xr
import numpy as np


def _get_next_date(date: datetime) -> datetime:
    return datetime(date.year + date.month // 12, date.month % 12 + 1, 1)


def generate_months_from_timestamps(
    start_ts: float, end_ts: float, chunk: str | int = "month"
):
    """Split a time interval into monthly or shorter download chunks.

    Parameters
    ----------
    start_ts, end_ts : float
        Start and end times as UTC timestamps.
    chunk : {"month", "week"} or int, optional
        Chunking strategy. Integer values are interpreted as a number of days.

    Returns
    -------
    months : list of str
        Month tags in ``YYYYMM`` format, one per chunk.
    intervals : list of tuple of str
        Start and end date tags in ``YYYYMMDD`` format, one per chunk.
    """
    start = datetime(*time.gmtime(start_ts)[:6])
    end = datetime(*time.gmtime(end_ts)[:6])

    months = []
    intervals = []

    # normalize start to midnight
    start = datetime(start.year, start.month, start.day)
    end = datetime(end.year, end.month, end.day)

    if chunk == "month":
        cur = start
        while cur <= end:
            months.append(cur.strftime("%Y%m"))
            last = min(end, _get_next_date(cur) - timedelta(days=1))
            intervals.append((cur.strftime("%Y%m%d"), last.strftime("%Y%m%d")))
            cur = _get_next_date(cur)

    else:
        if chunk == "week":
            step = 7
        elif isinstance(chunk, int) and chunk > 0:
            step = chunk
        else:
            raise ValueError(f"Invalid chunk: {chunk}")
        cur = start
        while cur <= end:
            month_end = _get_next_date(cur) - timedelta(days=1)
            last = min(end, month_end, cur + timedelta(days=step - 1))
            months.append(cur.strftime("%Y%m"))
            intervals.append((cur.strftime("%Y%m%d"), last.strftime("%Y%m%d")))
            cur = last + timedelta(days=1)

    return months, intervals


def file_exists_and_probably_ok(path: str) -> bool:
    """Check whether a downloaded netCDF file can be reused.

    Parameters
    ----------
    path : str
        Path to the candidate netCDF file.

    Returns
    -------
    bool
        ``True`` if the file exists, can be opened by xarray, and does not end
        with an all-NaN time slice; otherwise ``False``. Invalid files are
        removed so they can be downloaded again.
    """
    if not os.path.exists(path):
        return False
    try:
        with xr.open_dataset(path) as dsdisk:
            for variable in list(dsdisk.keys()):
                diskshape = dsdisk[variable].shape
                if (
                    len(dsdisk[variable].dims) > 0
                    and dsdisk[variable].dims[0] == "time"
                ):
                    if np.all(np.isnan(dsdisk[variable][-1])):
                        print(
                            f"{path} : All NANS in variable {variable}, re-downloading"
                        )
                        os.remove(path)
                        return False
            print(f"{path} already present")
            return True
    except Exception:
        print(f"{path} corrupted, re-downloading")
        os.remove(path)
        return False


def _month_dataset_url(variable: str, month: str, dataset="barra2r2") -> str:
    """Build the THREDDS NCSS URL for a monthly source file.

    Parameters
    ----------
    variable : str
        Name of the variable to download.
    month : str
        Month tag in ``YYYYMM`` format.
    dataset : {"barra2r2", "barra2c2", "bran"}, optional
        Source dataset identifier.

    Returns
    -------
    str
        THREDDS NCSS endpoint for the requested source file.
    """
    THREDDS_BASE = "https://thredds.nci.org.au/thredds/ncss/grid"
    if dataset == "barra2r2":
        dspath = (
            "ob53/output/reanalysis/AUS-11/BOM/ERA5/historical/hres/BARRA-R2/v1/1hr"
        )
        return f"{THREDDS_BASE}/{dspath}/{variable}/latest/{variable}_AUS-11_ERA5_historical_hres_BOM_BARRA-R2_v1_1hr_{month}-{month}.nc"
    if dataset == "barra2c2":
        dspath = (
            "ob53/output/reanalysis/AUST-04/BOM/ERA5/historical/hres/BARRA-C2/v1/1hr"
        )
        return f"{THREDDS_BASE}/{dspath}/{variable}/latest/{variable}_AUST-04_ERA5_historical_hres_BOM_BARRA-C2_v1_1hr_{month}-{month}.nc"
    if dataset == "bran":
        dspath = "gb6/BRAN/BRAN2023/daily"
        return f"{THREDDS_BASE}/{dspath}/ocean_{variable}_{month[0:4]}_{month[4:6]}.nc"
    else:
        raise ValueError(f"Unknown dataset: {dataset}")


def _interval_to_iso_utc(interval: Tuple[str, str]) -> Tuple[str, str]:
    start_tag, end_tag = interval
    start_iso = datetime.strptime(start_tag, "%Y%m%d").strftime("%Y-%m-%dT00:00:00Z")
    end_iso = datetime.strptime(end_tag, "%Y%m%d").strftime("%Y-%m-%dT23:59:59Z")
    return start_iso, end_iso


class MultiProgress:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._lines: List[str] = []
        self._initialized = False
        self._interactive = sys.stdout.isatty()

    def add(self, text: str) -> int:
        with self._lock:
            idx = len(self._lines)
            self._lines.append(text)
            if self._interactive:
                sys.stdout.write(text + "\n")
                sys.stdout.flush()
                self._initialized = True
            else:
                print(text, flush=True)

            return idx

    def update(self, idx: int, text: str) -> None:
        with self._lock:
            self._lines[idx] = text
            if not self._interactive:
                return
            n = len(self._lines)
            # Move cursor to the first progress line
            sys.stdout.write(f"\x1b[{n}F")
            for line in self._lines:
                sys.stdout.write("\x1b[2K")  # clear entire line
                sys.stdout.write(line + "\n")
            sys.stdout.flush()

    def finish(self) -> None:
        if self._interactive:
            sys.stdout.flush()


def _format_progress_line(
    filename: str,
    downloaded: int,
    total_size: int,
    elapsed: float,
    status: str,
) -> str:
    """Format one human-readable download-progress line.

    Parameters
    ----------
    filename : str
        Destination file name.
    downloaded : int
        Number of bytes already written.
    total_size : int
        Expected total size in bytes, or zero if unknown.
    elapsed : float
        Elapsed time in seconds.
    status : str
        Short status label.

    Returns
    -------
    str
        Progress line suitable for terminal output.
    """
    name = os.path.basename(filename)
    elapsed = max(elapsed, 1e-12)
    speed_mbs = downloaded / elapsed / 2**20
    if total_size > 0:
        percent = 100.0 * downloaded / total_size
        remaining = max(total_size - downloaded, 0)
        eta = remaining / max(downloaded / elapsed, 1e-12) if downloaded > 0 else 0.0
        return f"{name}: {status:11s} {percent:5.1f}% {downloaded / 2**20:8.1f}/{total_size / 2**20:8.1f} MB {speed_mbs:6.1f} MB/s ETA {eta:6.1f}s"
    return f"{name}: {status:11s} {downloaded / 2**20:8.1f} MB {speed_mbs:6.1f} MB/s"


def _download_with_progress(
    response: requests.Response,
    dst_file: str,
    progress: MultiProgress,
    progress_idx: int,
    chunk_size: int = 1024 * 1024,
) -> None:
    """Stream a response to disk while updating progress output.

    Parameters
    ----------
    response : requests.Response
        Open streaming HTTP response.
    dst_file : str
        Final output file path.
    progress : MultiProgress
        Progress manager.
    progress_idx : int
        Index of the progress line to update.
    chunk_size : int, optional
        Size of chunks read from the HTTP stream, in bytes.
    """
    tmp_file = dst_file + ".part"
    total_size = int(response.headers.get("Content-Length", 0))
    downloaded = 0
    t0 = time.time()
    last_update = 0.0

    with open(tmp_file, "wb") as f:
        for chunk in response.iter_content(chunk_size=chunk_size):
            if not chunk:
                continue
            f.write(chunk)
            downloaded += len(chunk)
            now = time.time()
            if now - last_update >= 0.2:
                progress.update(
                    progress_idx,
                    _format_progress_line(
                        dst_file, downloaded, total_size, now - t0, "downloading"
                    ),
                )
                last_update = now

    os.replace(tmp_file, dst_file)
    progress.update(
        progress_idx,
        _format_progress_line(
            dst_file, downloaded, total_size, time.time() - t0, "done"
        ),
    )


def _download_ncss_subset(
    variable: str,
    dataset: str,
    month: str,
    interval: Tuple[str, str],
    bounds,
    dst_file: str,
    progress: MultiProgress,
    progress_idx: int,
    timeout: int = 300,
) -> str:
    minlon, maxlon, minlat, maxlat = bounds
    time_start, time_end = _interval_to_iso_utc(interval)

    url = _month_dataset_url(variable, month, dataset)
    params = {
        "var": variable,
        "west": float(minlon),
        "east": float(maxlon),
        "south": float(minlat),
        "north": float(maxlat),
        "time_start": time_start,
        "time_end": time_end,
        "accept": "netcdf4",
    }

    progress.update(progress_idx, f"{os.path.basename(dst_file)}: connecting")
    with requests.get(url, params=params, stream=True, timeout=timeout) as response:
        response.raise_for_status()
        content_type = response.headers.get("Content-Type", "")
        if "html" in content_type.lower():
            text_preview = response.text[:1000]
            raise RuntimeError(
                f"Server returned HTML instead of netCDF for {response.url}\nContent-Type: {content_type}\nResponse preview:\n{text_preview}"
            )
        _download_with_progress(response, dst_file, progress, progress_idx)
    return dst_file


def _download_one_task(task, progress: MultiProgress) -> str:
    dst_file = task["dst_file"]
    progress_idx = progress.add(f"{os.path.basename(dst_file)}: queued")
    try:
        return _download_ncss_subset(
            variable=task["variable"],
            dataset=task["dataset"],
            month=task["month"],
            interval=task["interval"],
            bounds=task["bounds"],
            dst_file=dst_file,
            progress=progress,
            progress_idx=progress_idx,
        )
    except Exception as exc:
        tmp_file = dst_file + ".part"
        if os.path.exists(tmp_file):
            os.remove(tmp_file)
        progress.update(progress_idx, f"{os.path.basename(dst_file)}: FAILED - {exc}")
        return ""


def download_variables(
    dst_dir: str,
    dataset: str,
    variables: List[str],
    bounds,
    time_from: float,
    time_to: float,
    chunk: str | int = "month",
    max_workers: int = 4,
) -> List[str]:
    """Download one or more variables from a gridded THREDDS dataset.

    Parameters
    ----------
    dst_dir : str
        Directory where downloaded netCDF files are stored.
    dataset : {"barra2r2", "barra2c2", "bran"}
        Source dataset identifier.
    variables : list of str
        Variables to download.
    bounds : sequence of float
        Bounding box ``[minlon, maxlon, minlat, maxlat]``.
    time_from, time_to : float
        Download interval as UTC timestamps.
    chunk : {"month", "week"} or int, optional
        Chunking strategy. Integer values are interpreted as a number of days.
    max_workers : int, optional
        Number of concurrent downloads.

    Returns
    -------
    list of str
        Sorted list of downloaded or reused files.
    """
    os.makedirs(dst_dir, exist_ok=True)
    months, intervals = generate_months_from_timestamps(time_from, time_to, chunk)

    all_files: List[str] = []
    tasks = []

    for variable in variables:
        for month, interval in zip(months, intervals):
            start_tag, end_tag = interval
            dst_file = os.path.join(dst_dir, f"{variable}_{start_tag}_{end_tag}.nc")

            if file_exists_and_probably_ok(dst_file):
                all_files.append(dst_file)
                continue

            tasks.append(
                {
                    "variable": variable,
                    "dataset": dataset,
                    "month": month,
                    "interval": interval,
                    "bounds": bounds,
                    "dst_file": dst_file,
                }
            )

    if not tasks:
        return sorted(all_files)

    progress = MultiProgress()
    results: List[str] = []

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(_download_one_task, task, progress) for task in tasks
        ]
        for future in as_completed(futures):
            results.append(future.result())

    progress.finish()
    all_files.extend(results)
    return sorted(all_files)


def download_sst(
    dst_dir: str,
    variables: List[str],
    bounds,
    time_from: float,
    time_to: float,
    chunk: str | int = "month",
    username: str | None = None,
    password: str | None = None,
) -> List[str]:
    """Download sea-surface-temperature products from Copernicus Marine.

    Parameters
    ----------
    dst_dir : str
        Directory where downloaded netCDF files are stored.
    variables : list of str
        SST variables to request.
    bounds : sequence of float
        Bounding box ``[minlon, maxlon, minlat, maxlat]``.
    time_from, time_to : float
        Download interval as UTC timestamps.
    chunk : {"month", "week"} or int, optional
        Chunking strategy. Kept consistent with ``download_variables``.
    username, password : str, optional
        Copernicus Marine credentials.

    Returns
    -------
    list of str
        Sorted list of downloaded or reused files.
    """
    os.makedirs(dst_dir, exist_ok=True)
    months, intervals = generate_months_from_timestamps(time_from, time_to, chunk)
    all_files: List[str] = []
    for variable in variables:
        for month, interval in zip(months, intervals):
            start_tag, end_tag = interval
            fname = f"{variable}_{start_tag}_{end_tag}.nc"
            dst_file = os.path.join(dst_dir, fname)

            if file_exists_and_probably_ok(dst_file):
                all_files.append(dst_file)
                continue

            copernicusmarine.subset(
                dataset_id="METOFFICE-GLO-SST-L4-REP-OBS-SST",
                variables=[f"{variable}"],
                start_datetime=f"{start_tag[0:4]}-{start_tag[4:6]}-{start_tag[6:8]}",
                end_datetime=f"{end_tag[0:4]}-{end_tag[4:6]}-{end_tag[6:8]}",
                minimum_longitude=137.0,
                maximum_longitude=159.0,
                minimum_latitude=-32.0,
                maximum_latitude=-6.0,
                output_directory=dst_dir,
                output_filename=fname,
                username=username,
                password=password,
                overwrite=True,
            )
            all_files.append(dst_file)

    return sorted(all_files)


# %%
# When executed as a script, we download all forcing products required by the
# preprocessing step.

if __name__ == "__main__":
    from slim4.tools import parse_time
    import params

    bounds = params.bounds
    time_from = parse_time(params.initial_time)
    time_to = parse_time(params.final_time)

    files = download_variables(
        "data/netcdf/hydro3d",
        "bran",
        ["u", "v", "temp", "salt"],
        bounds,
        time_from,
        time_to,
        chunk="month",
        max_workers=5,
    )

    files = download_variables(
        "data/netcdf/hydro2d",
        "bran",
        ["eta_t"],
        bounds,
        time_from,
        time_to,
        chunk="month",
        max_workers=5,
    )

    files = download_variables(
        "data/netcdf/wind",
        "barra2c2",
        [
            "uas",  # east wind 10m
            "vas",  # notrh wind 10m
            "psl",  # sea level pressure
        ],
        bounds,
        time_from,
        time_to,
        chunk=16,
        max_workers=5,
    )

    files = download_variables(
        "data/netcdf/atm",
        "barra2c2",
        [
            "hurs",  # surface relative humidity [%]
            "hfls",  # surface upward latent heat flux [W/m²]
            "hfss",  # surface upward sensible heat flux [W/m²]
            "rsds",  # surface downwelling shortwave radiation [W/m²]
            "rsus",  # surface upwelling shortwave radiation [W/m²]
            "rlds",  # surface downwelling longwave radiation [W/m²]
            "rlus",  # surface upwelling longwave radiation [W/m²]
            "pr",  # precipitation [kg/m²/s]
            "tas",  # near-surface air temp [K]
            "ts",  # surface skin temp [K]
        ],
        bounds,
        time_from,
        time_to,
        chunk=16,
        max_workers=5,
    )

    files = download_sst(
        "data/netcdf/sst", ["analysed_sst"], bounds, time_from, time_to
    )