from pathlib import Path
from datetime import datetime, timedelta, timezone
import time

try:
    import requests
except ModuleNotFoundError:  # pragma: no cover - dependency guard for runtime environments
    requests = None

BASE_URL = "https://nomads.ncep.noaa.gov/cgi-bin/filter_gfs_0p25.pl"
WAVE_BASE_URL = "https://nomads.ncep.noaa.gov/cgi-bin/filter_gfswave.pl"
PROJECT_ROOT = Path(__file__).resolve().parent.parent
OUTPUT_DIR = PROJECT_ROOT / "data" / "grib"
REQUEST_TIMEOUT = (10, 120)
RETRY_LIMIT = 3
MIN_FILE_SIZE_BYTES = 1024

FORECAST_HOURS = [
    0, 3, 6, 9, 12, 15, 18, 21, 24,
    27, 30, 33, 36, 39, 42, 45,
    48, 51, 54, 57, 60, 63, 66,
    69, 72,
]

REGION = {
    "leftlon": 120,
    "rightlon": 150,
    "toplat": 50,
    "bottomlat": 20,
}

ATMOS_VARIABLES = [
    "UGRD",
    "VGRD",
    "APCP",
    "PRMSL",
    "TMP",
]

WAVE_VARIABLES = [
    "HTSGW",
    "DIRPW",
    "PERPW",
]

ATMOS_LEVELS = [
    "lev_10_m_above_ground",
    "lev_surface",
    "lev_mean_sea_level",
]

WAVE_LEVELS = [
    "lev_surface",
]

OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


def get_cycle(reference_time=None):
    now = reference_time or datetime.now(timezone.utc)
    # NOMADS availability usually lags the wall clock. Bias one cycle back.
    candidate = now - timedelta(hours=5)
    hour = (candidate.hour // 6) * 6
    cycle_time = candidate.replace(hour=hour, minute=0, second=0, microsecond=0)
    return cycle_time.strftime("%Y%m%d"), f"{cycle_time.hour:02d}"


def build_atmos_request(date, cycle, forecast_hour):
    filename = f"gfs.t{cycle}z.pgrb2.0p25.f{forecast_hour}"
    params = {
        "file": filename,
        "leftlon": REGION["leftlon"],
        "rightlon": REGION["rightlon"],
        "toplat": REGION["toplat"],
        "bottomlat": REGION["bottomlat"],
        "dir": f"/gfs.{date}/{cycle}/atmos",
    }

    for variable in ATMOS_VARIABLES:
        params[f"var_{variable}"] = "on"

    for level in ATMOS_LEVELS:
        params[level] = "on"

    return params


def build_wave_request(date, cycle, forecast_hour):
    filename = f"gfswave.t{cycle}z.global.0p25.f{forecast_hour}.grib2"
    params = {
        "file": filename,
        "leftlon": REGION["leftlon"],
        "rightlon": REGION["rightlon"],
        "toplat": REGION["toplat"],
        "bottomlat": REGION["bottomlat"],
        "dir": f"/gfs.{date}/{cycle}/wave/gridded",
    }

    for variable in WAVE_VARIABLES:
        params[f"var_{variable}"] = "on"

    for level in WAVE_LEVELS:
        params[level] = "on"

    return params


def validate_response(response):
    content_type = response.headers.get("Content-Type", "").lower()
    if "html" in content_type or "text/plain" in content_type:
        preview = response.text[:200].strip().replace("\n", " ")
        raise ValueError(f"unexpected response content type {content_type}: {preview}")


def is_fatal_network_error(error):
    error_text = str(error)
    fatal_markers = (
        "NameResolutionError",
        "Failed to resolve",
        "Temporary failure in name resolution",
    )
    return any(marker in error_text for marker in fatal_markers)


def download_file(session, base_url, params, output_path):
    temp_path = output_path.with_suffix(".grib2.part")

    if output_path.exists() and output_path.stat().st_size >= MIN_FILE_SIZE_BYTES:
        print("skip", output_path)
        return

    for attempt in range(1, RETRY_LIMIT + 1):
        try:
            print(f"downloading {output_path} (attempt {attempt}/{RETRY_LIMIT})")
            with session.get(
                base_url,
                params=params,
                stream=True,
                timeout=REQUEST_TIMEOUT,
            ) as response:
                response.raise_for_status()
                validate_response(response)

                with temp_path.open("wb") as file_handle:
                    for chunk in response.iter_content(1024 * 1024):
                        if chunk:
                            file_handle.write(chunk)

            if temp_path.stat().st_size < MIN_FILE_SIZE_BYTES:
                raise ValueError(f"downloaded file too small: {temp_path.stat().st_size} bytes")

            temp_path.replace(output_path)
            return
        except (requests.RequestException, ValueError) as exc:
            if temp_path.exists():
                temp_path.unlink()
            print("  download failed:", exc)
            if is_fatal_network_error(exc):
                raise RuntimeError("fatal network error while reaching NOAA") from exc
            if attempt == RETRY_LIMIT:
                raise
            time.sleep(attempt * 2)


def download_forecast(session, date, cycle, forecast_hour):
    atmos_path = OUTPUT_DIR / f"{date}_{cycle}_f{forecast_hour}.grib2"
    wave_path = OUTPUT_DIR / f"{date}_{cycle}_f{forecast_hour}_wave.grib2"

    download_file(
        session,
        BASE_URL,
        build_atmos_request(date, cycle, forecast_hour),
        atmos_path,
    )
    download_file(
        session,
        WAVE_BASE_URL,
        build_wave_request(date, cycle, forecast_hour),
        wave_path,
    )


def main():
    if requests is None:
        raise RuntimeError("requests is required to download GFS data")

    date, cycle = get_cycle()
    print("cycle:", date, cycle)

    session = requests.Session()
    session.headers["User-Agent"] = "weather-pipeline/1.0"

    failures = []
    for forecast_hour in FORECAST_HOURS:
        forecast_hour_str = f"{forecast_hour:03d}"
        try:
            download_forecast(session, date, cycle, forecast_hour_str)
        except RuntimeError as exc:
            failures.append((forecast_hour_str, str(exc)))
            break
        except Exception as exc:
            failures.append((forecast_hour_str, str(exc)))

    if failures:
        for forecast_hour, error in failures:
            print(f"failed forecast {forecast_hour}: {error}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()
