"""The ERA5 lake: layout, grid math, and read clients for the bucket-hosted archive.

The lake is a set of parquet files in an S3-compatible bucket (Contabo Object
Storage, ``era5-thermograph``), extracted once from the public Earthmover ERA5
Icechunk archive by ``gen_era5_lake.py`` and read three ways:

- the **lake service** (``lake_app.py``, prod-only Swarm service) serves point
  histories and SQL-style queries with a local disk cache in front;
- the **web/worker backend** asks the lake service first
  (``THERMOGRAPH_LAKE_URL``), or reads the bucket directly when only S3
  credentials are configured (beta/LAN, no lake service);
- ad-hoc analysis hits the hive table straight from the bucket.

Two layouts, one manifest (all under ``era5/``):

- ``daily/tile=<ti>_<tj>/year=<Y>/month=<M>/part.parquet`` — the analytical
  table, hive-partitioned by location/year/month. The location grain is a
  12x12-point tile (~3°x3°), aligned to the Earthmover store's chunk grid; a
  per-point x per-month grain would be ~170M objects, which no object store
  enjoys.
- ``points/lat=<i>/lon=<j>.parquet`` — the serving projection: one file per
  ERA5 grid point holding its full daily history, so the hot path (one cell's
  history) is a single GET.
- ``manifest.parquet`` — Iceberg-manifest-style index of every point file:
  grid indices, coordinates, row count, date span. Readers prune with it and
  the extractor resumes from it.

Columns everywhere: the finalized history-store schema (date, tmax, tmin,
precip, wind, gust, humid, fmax, fmin, feels, doy — °F/in/mph/%, as
``climate._finalize_approximated`` emits it), so a lake frame drops straight
into ``climate._write_history_backed`` exactly like a NASA fetch does.
"""
import io
import os

import polars as pl

# --- ERA5 native grid (0.25°, 721 x 1440, lat 90..-90, lon 0..359.75) ---------
LAT_N = 721
LON_N = 1440
STEP = 0.25
TILE = 12                      # matches the Earthmover chunk grid (8736h x 12 x 12)

PREFIX = "era5"
MANIFEST_KEY = f"{PREFIX}/manifest.parquet"

MIN_LAKE_DAYS = 10_000         # a plausibly-full 1980+ record, same spirit as
                               # climate.MIN_ARCHIVE_DAYS (rejects partial files)


def to_idx(lat: float, lon: float) -> tuple[int, int]:
    """Nearest ERA5 grid indices for a lat/lon. lat_idx runs north->south."""
    i = round((90.0 - lat) / STEP)
    j = round((lon % 360.0) / STEP) % LON_N
    return min(max(i, 0), LAT_N - 1), j


def to_coords(i: int, j: int) -> tuple[float, float]:
    """Center lat/lon ([-180, 180) longitude) of grid indices."""
    lon = j * STEP
    return 90.0 - i * STEP, lon - 360.0 if lon >= 180.0 else lon


def tile_of(i: int, j: int) -> tuple[int, int]:
    return i // TILE, j // TILE


def point_key(i: int, j: int) -> str:
    return f"{PREFIX}/points/lat={i}/lon={j}.parquet"


def daily_part_key(ti: int, tj: int, year: int, month: int) -> str:
    return f"{PREFIX}/daily/tile={ti}_{tj}/year={year}/month={month:02d}/part.parquet"


# --- Config (all optional: unset means "no lake configured") ------------------

def lake_url() -> str:
    """Lake service base URL (prod: the Swarm VIP), '' when not deployed."""
    return os.environ.get("THERMOGRAPH_LAKE_URL", "").rstrip("/")


def s3_config() -> "dict | None":
    """Bucket access for direct reads/writes, None unless fully configured.
    Contabo is path-style; region is a signing formality there ('default')."""
    key = os.environ.get("THERMOGRAPH_LAKE_S3_ACCESS_KEY", "")
    secret = os.environ.get("THERMOGRAPH_LAKE_S3_SECRET_KEY", "")
    if not (key and secret):
        return None
    return {
        "endpoint": os.environ.get("THERMOGRAPH_LAKE_S3_ENDPOINT",
                                   "https://eu2.contabostorage.com"),
        "bucket": os.environ.get("THERMOGRAPH_LAKE_S3_BUCKET", "era5-thermograph"),
        "region": os.environ.get("THERMOGRAPH_LAKE_S3_REGION", "default"),
        "access_key": key,
        "secret_key": secret,
    }


def local_dir() -> str:
    """Filesystem mirror of the lake layout — tests and offline dev reads."""
    return os.environ.get("THERMOGRAPH_LAKE_LOCAL_DIR", "")


def storage_options(cfg: dict) -> dict:
    """polars/object_store options for scan_parquet against the bucket."""
    return {
        "aws_access_key_id": cfg["access_key"],
        "aws_secret_access_key": cfg["secret_key"],
        "aws_region": cfg["region"],
        "endpoint_url": cfg["endpoint"],
        # Contabo serves buckets on the path, not as a subdomain.
        "aws_virtual_hosted_style_request": "false",
    }


# --- Point-history reads (the app's hot path) ----------------------------------

class LakeUnavailable(Exception):
    """No lake configured, point not in the lake, or the read failed."""


def _read_local(i: int, j: int) -> pl.DataFrame:
    path = os.path.join(local_dir(), point_key(i, j))
    if not os.path.exists(path):
        raise LakeUnavailable(f"no local lake file for point ({i}, {j})")
    return pl.read_parquet(path)


def _read_service(i: int, j: int) -> pl.DataFrame:
    import httpx  # noqa: PLC0415 - keep module import light for the extractor venv

    lat, lon = to_coords(i, j)
    r = httpx.get(f"{lake_url()}/history", params={"lat": lat, "lon": lon},
                  timeout=httpx.Timeout(20.0, connect=3.0))
    if r.status_code == 404:
        raise LakeUnavailable(f"point ({i}, {j}) not in the lake")
    r.raise_for_status()
    return pl.read_parquet(io.BytesIO(r.content))


def _read_bucket(i: int, j: int, cfg: dict) -> pl.DataFrame:
    try:
        return pl.read_parquet(f"s3://{cfg['bucket']}/{point_key(i, j)}",
                               storage_options=storage_options(cfg))
    except Exception as e:  # noqa: BLE001 - object_store errors are not one type
        raise LakeUnavailable(f"bucket read failed for point ({i}, {j}): {e}") from e


def fetch_history_lake(cell: dict, start: "str | None" = None) -> pl.DataFrame:
    """Finalized daily history for the ERA5 point nearest a grid cell, from the
    first configured source: local dir -> lake service -> bucket. The lake
    stores each point's whole 1940+ record; pass ``start`` (the caller's
    grading window, e.g. climate.START_DATE) to slice — grading semantics stay
    the caller's, the deep record stays in the lake. Raises LakeUnavailable
    when unconfigured, missing, or implausibly short."""
    import datetime  # noqa: PLC0415

    i, j = to_idx(cell["center_lat"], cell["center_lon"])
    if local_dir():
        df = _read_local(i, j)
    elif lake_url():
        df = _read_service(i, j)
    else:
        cfg = s3_config()
        if cfg is None:
            raise LakeUnavailable("no lake configured")
        df = _read_bucket(i, j, cfg)
    if start is not None:
        df = df.filter(pl.col("date") >= datetime.date.fromisoformat(start))
    if df.height < MIN_LAKE_DAYS:
        raise LakeUnavailable(f"lake point ({i}, {j}) is short: {df.height} rows")
    return df.drop([c for c in ("lat_idx", "lon_idx") if c in df.columns])
