"""Prepare Varda-single-1.0 initial conditions in a standalone process.

Fetches the KENDA-CH1 (regional, 1 km) and ECMWF open-data IFS (global) analyses at `init` and
`init - 6h`, preprocesses them exactly like `anemoi-inference run` does, and saves:
  * the model input tensor  (multi_step_input, variables, values)  -> <out>
  * the mesh (lat, lon, lam mask)                                   -> <mesh> (if missing)

Running this out-of-process keeps the native GRIB/regridding libraries (eccodes, MIR, atlas…)
out of the ZeroGPU main process, and lets the web app stay responsive meanwhile.

Usage: python prep.py <config.json> <init ISO> <out.npy> <mesh.npz>
"""

import json
import logging
import os
import sys
import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor, wait
from datetime import datetime, timedelta

from eccodes_cosmo_resources import get_definitions_path

os.environ["ECCODES_DEFINITION_PATH"] = str(get_definitions_path())
os.environ.setdefault("ECCODES_VERSION_CHECK_OFF", "1")
os.environ["CUDA_VISIBLE_DEVICES"] = ""  # CPU only

import earthkit.data as ekd  # noqa: E402
import numpy as np  # noqa: E402
import requests  # noqa: E402
from anemoi_plugins_meteoswiss.inference.inputs import oper_kenda_opendata as kenda  # noqa: E402
from anemoi_plugins_meteoswiss.transform.filters.vertical_interpolation import ModelToPressureLevel  # noqa: E402

from anemoi.inference.config.run import RunConfiguration  # noqa: E402
from anemoi.inference.runners import create_runner  # noqa: E402

logging.basicConfig(level=logging.INFO, format="[prep] %(levelname)s %(name)s: %(message)s")
LOG = logging.getLogger("prep")

# The KENDA-CH1 server is throttled per connection (~1-2 MB/s seen from the Space) but supports HTTP
# Range requests, so every file is fetched as CHUNK-sized byte ranges over many connections at once.
CHUNK = 16 * 1024 * 1024
CONNECTIONS = 32
DOWNLOAD_RETRIES = 4
# KENDA-CH1 fields the Varda input needs at each analysis time (logged by the plugin as
# "fetching ... variable(s)"). Only used to start downloads early: anything missing from this
# list is still fetched by the plugin itself, an extra entry just costs bandwidth.
KENDA_VARS = ["P", "PMSL", "PS", "QV", "T", "TD_2M", "T_2M", "U", "U_10M", "V", "V_10M"]
INPUT_TIMES_H = (-6, 0)  # the forecaster needs the analyses at init - 6h and init

_pool = ThreadPoolExecutor(64)  # one light coordinating thread per file
_chunks = ThreadPoolExecutor(CONNECTIONS)  # the actual transfers; FIFO, so files finish in queue order
_inflight: dict[str, Future] = {}  # destination path -> download, shared by prefetch and plugin
_inflight_lock = threading.Lock()
_session = threading.local()


def _http() -> requests.Session:
    if not hasattr(_session, "s"):  # one keep-alive session per thread
        _session.s = requests.Session()
    return _session.s


def _retrying(what: str, fn):
    for attempt in range(1, DOWNLOAD_RETRIES + 1):
        try:
            return fn()
        except Exception as e:  # noqa: BLE001
            LOG.warning("  %s: attempt %d/%d failed: %s", what, attempt, DOWNLOAD_RETRIES, e)
            if attempt == DOWNLOAD_RETRIES:
                raise
            time.sleep(3 * attempt)


def _get_range(href: str, fd: int, start: int, end: int) -> None:
    """Fetch bytes [start, end] of `href` into `fd` at the same offset."""
    # (connect, read) timeouts: a transfer that sends nothing for 60 s is retried
    with _http().get(href, headers={"Range": f"bytes={start}-{end}"}, stream=True, timeout=(30, 60)) as r:
        r.raise_for_status()
        if r.status_code != 206:
            raise IOError(f"server ignored the Range header (HTTP {r.status_code})")
        pos = start
        for block in r.iter_content(chunk_size=1024 * 1024):
            os.pwrite(fd, block, pos)
            pos += len(block)
    if pos != end + 1:
        raise IOError(f"short range {start}-{end}: got {pos - start} bytes")


def _size(href: str) -> int:
    # a 1-byte GET rather than HEAD: presigned URLs are often signed for GET only
    with _http().get(href, headers={"Range": "bytes=0-0"}, stream=True, timeout=(30, 60)) as r:
        r.raise_for_status()
        if r.status_code == 206:
            return int(r.headers["Content-Range"].rsplit("/", 1)[1])
        return -1  # no Range support


def _download_stream(href: str, part: str) -> None:
    with _http().get(href, stream=True, timeout=(30, 60)) as r, open(part, "wb") as f:
        r.raise_for_status()
        for block in r.iter_content(chunk_size=4 * 1024 * 1024):
            f.write(block)


def _download(href: str, path: str) -> None:
    """Download `href` to `path` (via a .part file) as parallel byte ranges, each retried on failure."""
    name = os.path.relpath(path, os.path.dirname(os.path.dirname(path)))
    part = path + ".part"
    t0 = time.time()
    size = _retrying(name, lambda: _size(href))
    if size < 0:
        _retrying(name, lambda: _download_stream(href, part))
    else:
        fd = os.open(part, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
        try:
            os.ftruncate(fd, size)
            ranges = [(s, min(s + CHUNK, size) - 1) for s in range(0, size, CHUNK)]
            futs = [
                _chunks.submit(_retrying, f"{name}@{s}", lambda s=s, e=e: _get_range(href, fd, s, e))
                for s, e in ranges
            ]
            wait(futs)  # let every range finish before closing fd, even if one failed for good
            for f in futs:
                f.result()
        finally:
            os.close(fd)
    os.replace(part, path)
    dt = time.time() - t0
    mb = os.path.getsize(path) / 1e6
    LOG.info("  %s: %.0f MB in %.0fs (%.1f MB/s)", name, mb, dt, mb / max(dt, 1e-3))


def _fetch(href: str, path: str) -> Future | None:
    """Start downloading `href` to `path` unless it is cached or already being downloaded."""
    with _inflight_lock:
        if path in _inflight:
            return _inflight[path]
        if os.path.exists(path):
            return None
        os.makedirs(os.path.dirname(path), exist_ok=True)
        fut = _inflight[path] = _pool.submit(_download, href, path)
        return fut


def _parallel_cached_fields(cache_dir, subdir, hrefs_by_filename):
    """Drop-in for the plugin's `_cached_fields`: same on-disk cache layout, but missing KENDA-CH1
    assets (~180-275 MB each for the 3-D fields) download concurrently, reusing prefetches in flight."""
    if not cache_dir:
        return _orig_cached_fields(cache_dir, subdir, hrefs_by_filename)
    paths = {fn: os.path.join(cache_dir, subdir, fn) for fn in hrefs_by_filename}
    t0 = time.time()
    futs = [f for f in (_fetch(hrefs_by_filename[fn], p) for fn, p in paths.items()) if f is not None]
    for f in futs:
        f.result()
    if futs:
        LOG.info("KENDA-CH1 %s: waited %.0fs for %d download(s)", subdir, time.time() - t0, len(futs))
    return ekd.from_source("file", list(paths.values()))


_orig_cached_fields = kenda._cached_fields
kenda._cached_fields = _parallel_cached_fields


def _m2p_forward_parallel(self, data):
    """ModelToPressureLevel.forward with the analysis times interpolated in parallel threads (the GRIB
    decoding and numpy work release the GIL): ~1.6x faster than the plugin's serial loop, same output."""
    groups = [ekd.SimpleFieldList(list(g)) for g in data.group_by("valid_datetime")]
    if len(groups) < 2:
        return _orig_m2p_forward(self, data)
    with ThreadPoolExecutor(len(groups)) as pool:
        outs = list(pool.map(lambda g: _orig_m2p_forward(self, g), groups))
    # same order as the serial filter: every interpolated field first, then the surface passthrough
    fields = [f for o in outs for f in o]
    interp = [f for f in fields if f.metadata("typeOfLevel") == "isobaricInhPa"]
    rest = [f for f in fields if f.metadata("typeOfLevel") != "isobaricInhPa"]
    return ekd.SimpleFieldList(interp + rest)


_orig_m2p_forward = ModelToPressureLevel.forward
ModelToPressureLevel.forward = _m2p_forward_parallel


def _prefetch_kenda(cache_dir: str, init: datetime) -> None:
    """Queue every KENDA-CH1 download the run needs (grid constants + both analysis times) at once,
    so the t download overlaps decoding t-6h instead of waiting for it. Best effort."""
    try:
        for name, href in kenda._collection_static_assets().items():
            _fetch(href, os.path.join(cache_dir, "constants", name))
        for dh in INPUT_TIMES_H:
            t = init + timedelta(hours=dh)
            hrefs = kenda._hrefs_by_variable(kenda._stac_items(t))
            for v in KENDA_VARS:
                if v in hrefs:
                    _fetch(hrefs[v], os.path.join(cache_dir, "raw", t.strftime("%Y%m%dT%H%M%S"), f"{v}.grib2"))
        LOG.info("prefetch: %d KENDA-CH1 download(s) queued", sum(not f.done() for f in _inflight.values()))
    except Exception as e:  # noqa: BLE001
        LOG.warning("prefetch failed (the plugin will fetch on demand): %s", e)


def _kenda_cache_dir(cfg: dict) -> str | None:
    for src in cfg["input"]["cutout"]:
        for name, opts in src.get("lam_0", {}).items():
            if name.replace("_", "-") == "oper-kenda-opendata":
                return opts.get("cache_dir")
    return None


def main(cfg_path: str, init: str, out: str, mesh: str) -> None:
    with open(cfg_path) as f:
        cfg = json.load(f)
    cfg["device"] = "cpu"
    date = datetime.fromisoformat(init)
    t0 = time.time()
    cache_dir = _kenda_cache_dir(cfg)
    if cache_dir:
        threading.Thread(target=_prefetch_kenda, args=(cache_dir, date), daemon=True).start()
    runner = create_runner(RunConfiguration.load(cfg))
    ds = runner.dataset_names[0]
    prog = runner.prognostics_inputs[ds].create_input_state(date=date)
    LOG.info("prognostics fetched in %.0fs", time.time() - t0)
    const = runner.constant_forcings_inputs[ds].create_input_state(date=date)
    forc = runner.dynamic_forcings_inputs[ds].create_input_state(date=date)
    LOG.info("forcings ready at %.0fs", time.time() - t0)
    state = runner._combine_states(prog, const, forc)
    tensor = runner.tensor_handlers[ds].prepare_input_tensor(state)
    LOG.info("input tensor %s ready in %.0fs", tensor.shape, time.time() - t0)
    if not os.path.exists(mesh):
        np.savez(
            mesh + ".part.npz",
            lat=np.asarray(state["latitudes"], dtype=np.float64),
            lon=np.asarray(state["longitudes"], dtype=np.float64),
            lam=np.asarray(state["_mask"]["lam_0"], dtype=bool),
        )
        os.replace(mesh + ".part.npz", mesh)
    np.save(out + ".part.npy", tensor.astype(np.float32))
    os.replace(out + ".part.npy", out)
    _pool.shutdown(wait=False, cancel_futures=True)
    _chunks.shutdown(wait=False, cancel_futures=True)


if __name__ == "__main__":
    main(*sys.argv[1:5])