Northern-Thailand-Ping-River-Monitor/src/rid_reservoir.py

"""Collector for RID large-dam daily status (app.rid.go.th/reservoir).

The Royal Irrigation Department reservoir app exposes an unauthenticated
JSON API: ``POST https://app.rid.go.th/reservoir/api/dams`` with form field
``date=YYYY-MM-DD`` (empty = today) returns a daily snapshot of every large
dam in Thailand — storage, inflow and outflow in MCM — with archive depth
back to at least 2009. GET returns 404 ("Unknown method."); the POST body
may be empty but must carry a Content-Length.

A sibling endpoint transposes that axis: ``GET .../api/dam`` with
``dam_id``/``date_start``/``date_end`` returns ONE dam over a whole date
range. It is GET-only (POST answers 404 "Unknown method.") and served the
full 2009-01-01..today archive — 6,362 rows, 4.7 MB — in a single ~6 s
response, so backfilling one dam costs one request instead of one per
calendar day. Field names differ from api/dams; see parse_dam_range_records.

The reservoir that matters for P.1 flood forecasting is Mae Ngat Somboon
Chon (DAM_ID 200103), the only large dam upstream of Chiang Mai: during the
Oct 2024 record flood it reached 113% of usable capacity with inflow spikes
of ~19 MCM/day. All dams in the payload are stored (same request cost);
filtering happens at feature-build time.

See docs/DATA_SOURCES.md for the endpoint catalog.
"""

import datetime
import logging
import time
from typing import Any, Dict, List, Optional

import requests

logger = logging.getLogger(__name__)

RID_DAMS_URL = "https://app.rid.go.th/reservoir/api/dams"
RID_DAM_RANGE_URL = "https://app.rid.go.th/reservoir/api/dam"
MAE_NGAT_DAM_ID = "200103"
RID_ARCHIVE_START = datetime.date(2009, 1, 1)  # earliest date api/dam serves

# Per-column NUMERIC capacity; source junk beyond these becomes NULL instead
# of overflowing the insert and discarding the whole daily batch.
_MEASURE_BOUNDS = {
    "storage_mcm": 1e8,
    "storage_pct": 1e6,
    "inflow_mcm": 1e8,
    "outflow_mcm": 1e8,
    "level_msl": 1e6,
}


def _bounded(value: Optional[float], limit: float) -> Optional[float]:
    if value is not None and abs(value) >= limit:
        return None
    return value


def _to_float(value: Any) -> Optional[float]:
    """API numerics arrive as strings ('222.01'), ' - ' placeholders, or None."""
    if value is None:
        return None
    if isinstance(value, str):
        value = value.replace(",", "").strip()
        if value in ("", "-"):
            return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def parse_dam_records(payload: Dict) -> List[Dict]:
    """Flatten the regions/dams payload into per-dam daily rows."""
    records = []
    for region in payload.get("regions") or []:
        for dam in region.get("dams") or []:
            dam_id = dam.get("DAM_ID")
            try:
                date = datetime.date.fromisoformat(dam.get("DMD_Date") or "")
            except ValueError:
                continue
            if not dam_id:
                continue
            records.append(
                {
                    "dam_id": dam_id,
                    "region": region.get("region_name"),
                    "name_th": dam.get("DAM_Name"),
                    "latitude": _to_float(dam.get("DAM_Lat")),
                    "longitude": _to_float(dam.get("DAM_Lon")),
                    "capacity_max_mcm": _to_float(dam.get("DAM_QMax")),
                    "capacity_normal_mcm": _to_float(dam.get("DAM_QStore")),
                    "date": date,
                    "storage_mcm": _to_float(dam.get("DMD_QUse")),
                    "storage_pct": _to_float(dam.get("PERCENT_DMD_QUse")),
                    "inflow_mcm": _to_float(dam.get("DMD_Inflow")),
                    "outflow_mcm": _to_float(dam.get("DMD_Outflow")),
                    "level_msl": _to_float(dam.get("DMD_Q")),
                }
            )
    return records


def parse_dam_range_records(payload: Dict) -> List[Dict]:
    """Flatten one api/dam single-dam range payload into the same rows as
    parse_dam_records, so both endpoints feed one store.

    api/dam names its columns differently and suffixes each measurement
    ``_curr`` / ``_prev``; ``_prev`` is the SAME calendar date one year
    earlier (confirmed against api/dams' DMD_Date_prev) and is dropped —
    those days are rows of their own. Mapping, verified equal to api/dams
    on 2019-01-05, 2024-09-24, 2024-10-05 and 2026-08-08:

        DMD_QUse_curr          -> storage_mcm   (identical)
        PERCENT_DMD_QUse_curr  -> storage_pct   (2 dp; api/dams rounds to
                                                 whole percent, 112.62 vs 113)
        DMD_Inflow_curr        -> inflow_mcm    (identical)
        DMD_Outflow_curr       -> outflow_mcm   (identical)
        DMD_ULevel_curr        -> level_msl     (only source of the level:
                                                 api/dams' DMD_Q is ' - ' for
                                                 all 35 dams, and api/dam's
                                                 DMD_Q_curr is a constant 0.00)

    Dam metadata (name, region, capacities) and coordinates are carried too,
    so the range path never has to blank rid_dams.
    """
    coords = payload.get("dam_coordinates") or {}
    payload_dam_id = payload.get("dam_id")
    records = []
    for row in payload.get("dam_data") or []:
        dam_id = row.get("DAM_ID") or payload_dam_id
        try:
            date = datetime.date.fromisoformat(row.get("DATE_curr") or "")
        except ValueError:
            continue
        if not dam_id:
            continue
        records.append(
            {
                "dam_id": dam_id,
                "region": row.get("DAM_Region") or payload.get("dam_region"),
                "name_th": row.get("DAM_Name") or payload.get("dam_name"),
                "latitude": _to_float(coords.get("lat")),
                "longitude": _to_float(coords.get("lng")),
                "capacity_max_mcm": _to_float(row.get("DAM_QMax")),
                "capacity_normal_mcm": _to_float(row.get("DAM_QStore")),
                "date": date,
                "storage_mcm": _to_float(row.get("DMD_QUse_curr")),
                "storage_pct": _to_float(row.get("PERCENT_DMD_QUse_curr")),
                "inflow_mcm": _to_float(row.get("DMD_Inflow_curr")),
                "outflow_mcm": _to_float(row.get("DMD_Outflow_curr")),
                # An MSL elevation of exactly 0 is "not published", not a
                # reading — these dams sit between 45 and 400 m.
                "level_msl": _to_float(row.get("DMD_ULevel_curr")) or None,
            }
        )
    return records


class RidReservoirClient:
    """HTTP client for the RID reservoir daily-status API."""

    def __init__(
        self,
        url: str = RID_DAMS_URL,
        session: Optional[requests.Session] = None,
        timeout: int = 120,
        range_url: str = RID_DAM_RANGE_URL,
    ):
        self.url = url
        self.range_url = range_url
        self.session = session or requests.Session()
        self.timeout = timeout

    def fetch_day(self, date: Optional[datetime.date] = None) -> List[Dict]:
        """Daily rows for every large dam on `date` (None = today)."""
        data = {"date": date.isoformat()} if date else {"date": ""}
        response = self.session.post(
            self.url,
            data=data,
            headers={"X-Requested-With": "XMLHttpRequest"},
            timeout=self.timeout,
        )
        response.raise_for_status()
        return parse_dam_records(response.json())

    def fetch_dam_range(
        self, dam_id: str, start: datetime.date, end: datetime.date
    ) -> List[Dict]:
        """Every published day in [start, end] for one dam, in one request.

        GET only — api/dam answers POST with 404 "Unknown method.", the exact
        opposite of api/dams. An unrecognised dam_id still returns HTTP 200
        but with a PHP notice page instead of JSON, so a decode failure is
        reported as a bad request rather than a transport error.
        """
        response = self.session.get(
            self.range_url,
            params={
                "dam_id": dam_id,
                "date_start": start.isoformat(),
                "date_end": end.isoformat(),
                "percent": "",
            },
            timeout=self.timeout,
        )
        response.raise_for_status()
        try:
            payload = response.json()
        except ValueError as e:
            raise ValueError(
                f"api/dam returned non-JSON for dam_id '{dam_id}' "
                f"({start}..{end}) — unknown dam_id?"
            ) from e
        return parse_dam_range_records(payload)


class RidReservoirStore:
    """SQL persistence for dam metadata + daily measurements.

    Shares the app's relational database; writes rid_dams (metadata) and
    rid_reservoir_daily keyed (dam_id, date) — the composite natural PK keeps
    the table TimescaleDB-hypertable compatible.
    """

    def __init__(self, connection_string: str, db_type: str):
        self.db_type = db_type.lower()
        if self.db_type not in ("sqlite", "postgresql", "mysql"):
            raise ValueError(
                f"Reservoir collection requires a SQL database, got '{db_type}'"
            )
        self.connection_string = connection_string
        self.engine = None

    def connect(self) -> bool:
        try:
            from sqlalchemy import create_engine

            self.engine = create_engine(self.connection_string, pool_pre_ping=True)
            self._create_tables()
            return True
        except Exception as e:
            logger.error(f"RidReservoirStore failed to connect: {e}")
            self.engine = None
            return False

    def _create_tables(self):
        from sqlalchemy import text

        ddl = [
            """
            CREATE TABLE IF NOT EXISTS rid_dams (
                dam_id VARCHAR(10) PRIMARY KEY,
                region VARCHAR(40),
                name_th VARCHAR(255),
                latitude NUMERIC(10,6),
                longitude NUMERIC(10,6),
                capacity_max_mcm NUMERIC(10,2),
                capacity_normal_mcm NUMERIC(10,2),
                updated_at TIMESTAMP
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS rid_reservoir_daily (
                dam_id VARCHAR(10) NOT NULL,
                date DATE NOT NULL,
                storage_mcm NUMERIC(10,2),
                storage_pct NUMERIC(8,2),
                inflow_mcm NUMERIC(10,2),
                outflow_mcm NUMERIC(10,2),
                level_msl NUMERIC(8,2),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (dam_id, date)
            )
            """,
        ]
        if self.db_type != "mysql":
            ddl.append(
                "CREATE INDEX IF NOT EXISTS idx_rid_reservoir_date "
                "ON rid_reservoir_daily(date)"
            )
        with self.engine.begin() as conn:
            for statement in ddl:
                conn.execute(text(statement))
        # Widen storage_pct on tables created before 2026-08-13: the source
        # publishes junk percents (dam 100602 reports 87798%) that overflowed
        # NUMERIC(6,2) and discarded whole daily batches.
        if self.db_type == "postgresql":
            migrations = (
                "ALTER TABLE rid_reservoir_daily "
                "ALTER COLUMN storage_pct TYPE NUMERIC(8,2)",
            )
        elif self.db_type == "mysql":
            migrations = (
                "ALTER TABLE rid_reservoir_daily MODIFY storage_pct NUMERIC(8,2)",
            )
        else:  # sqlite: NUMERIC is affinity only, nothing to widen
            migrations = ()
        for statement in migrations:
            try:
                with self.engine.begin() as conn:
                    conn.execute(text(statement))
            except Exception as e:
                logger.warning(f"rid_reservoir_daily migration skipped: {e}")

    def _upsert(
        self,
        table: str,
        key_cols: List[str],
        value_cols: List[str],
        preserve_cols: "tuple[str, ...]" = (),
    ) -> str:
        """Build an upsert; columns in `preserve_cols` keep their stored value
        when the incoming one is NULL.

        Dam metadata needs that: any payload that omits a name or coordinate
        would otherwise blank a good rid_dams row on every later write.
        """
        cols = key_cols + value_cols
        col_list = ", ".join(cols)
        params = ", ".join(f":{c}" for c in cols)
        conflict = ", ".join(key_cols)
        if self.db_type == "sqlite":
            if not preserve_cols:
                return f"INSERT OR REPLACE INTO {table} ({col_list}) VALUES ({params})"
            updates = ", ".join(
                f"{c} = COALESCE(excluded.{c}, {table}.{c})"
                if c in preserve_cols
                else f"{c} = excluded.{c}"
                for c in value_cols
            )
            return (
                f"INSERT INTO {table} ({col_list}) VALUES ({params}) "
                f"ON CONFLICT ({conflict}) DO UPDATE SET {updates}"
            )
        if self.db_type == "postgresql":
            updates = ", ".join(
                f"{c} = COALESCE(EXCLUDED.{c}, {table}.{c})"
                if c in preserve_cols
                else f"{c} = EXCLUDED.{c}"
                for c in value_cols
            )
            return (
                f"INSERT INTO {table} ({col_list}) VALUES ({params}) "
                f"ON CONFLICT ({conflict}) DO UPDATE SET {updates}"
            )
        updates = ", ".join(
            f"{c} = COALESCE(VALUES({c}), {c})"
            if c in preserve_cols
            else f"{c} = VALUES({c})"
            for c in value_cols
        )
        return (
            f"INSERT INTO {table} ({col_list}) VALUES ({params}) "
            f"ON DUPLICATE KEY UPDATE {updates}"
        )

    def save(self, records: List[Dict]) -> int:
        """Upsert one day's dam rows (metadata + measurements); idempotent."""
        if not records:
            return 0
        if not self.engine and not self.connect():
            return 0
        from sqlalchemy import text

        dam_cols = [
            "region",
            "name_th",
            "latitude",
            "longitude",
            "capacity_max_mcm",
            "capacity_normal_mcm",
        ]
        measure_cols = [
            "storage_mcm",
            "storage_pct",
            "inflow_mcm",
            "outflow_mcm",
            "level_msl",
        ]
        dam_sql = self._upsert(
            "rid_dams",
            ["dam_id"],
            dam_cols + ["updated_at"],
            preserve_cols=tuple(dam_cols),  # never blank metadata we already have
        )
        measure_sql = self._upsert(
            "rid_reservoir_daily",
            ["dam_id", "date"],
            measure_cols,
            # Only api/dam carries a level (api/dams' DMD_Q is ' - ' for every
            # dam), so the hourly collector would blank the backfilled level
            # of today and yesterday on every cycle.
            preserve_cols=("level_msl",),
        )
        now = datetime.datetime.now()
        dams = {}
        measurements = []
        for record in records:
            dam_row = {c: record.get(c) for c in dam_cols}
            dam_row.update({"dam_id": record["dam_id"], "updated_at": now})
            dams[record["dam_id"]] = dam_row
            measure_row = {
                c: _bounded(record.get(c), _MEASURE_BOUNDS[c]) for c in measure_cols
            }
            measure_row.update(
                {"dam_id": record["dam_id"], "date": record["date"]}
            )
            measurements.append(measure_row)
        try:
            with self.engine.begin() as conn:
                conn.execute(text(dam_sql), list(dams.values()))
                conn.execute(text(measure_sql), measurements)
            return len(measurements)
        except Exception as e:
            logger.error(f"RidReservoirStore save failed: {e}")
            return 0

    def present_dates(
        self,
        start: datetime.date,
        end: datetime.date,
        dam_id: Optional[str] = None,
    ) -> "set[datetime.date]":
        """Dates in [start, end] that already have rows, for backfill skipping.

        `dam_id` narrows the answer to one dam: the per-day fleet backfill can
        treat any stored date as done, but a per-dam backfill must not skip a
        date merely because some other dam published it.
        """
        if not self.engine and not self.connect():
            return set()
        from sqlalchemy import text

        sql = (
            "SELECT DISTINCT date FROM rid_reservoir_daily "
            "WHERE date >= :start AND date <= :end"
        )
        params = {"start": start, "end": end}
        if dam_id is not None:
            sql += " AND dam_id = :dam_id"
            params["dam_id"] = dam_id
        with self.engine.begin() as conn:
            values = conn.execute(text(sql), params).fetchall()
        dates = set()
        for (value,) in values:
            if isinstance(value, str):  # sqlite returns ISO strings
                value = datetime.date.fromisoformat(value[:10])
            if isinstance(value, datetime.datetime):
                value = value.date()
            dates.add(value)
        return dates


class RidReservoirCollector:
    """Fetch + persist the daily dam snapshot (today and yesterday)."""

    def __init__(self, db_config: Dict, client: Optional[RidReservoirClient] = None):
        self.client = client or RidReservoirClient()
        self.store = RidReservoirStore(
            connection_string=db_config["connection_string"],
            db_type=db_config["type"],
        )

    def run_cycle(self) -> int:
        """Collect today's snapshot plus yesterday's (late daily revisions)."""
        saved = 0
        today = datetime.date.today()
        for date in (None, today - datetime.timedelta(days=1)):
            try:
                saved += self.store.save(self.client.fetch_day(date))
            except Exception as e:
                logger.error(f"RID reservoir collection failed for {date}: {e}")
        logger.info(f"RID reservoir collection: {saved} dam-day rows saved")
        return saved


def backfill(
    store: RidReservoirStore,
    start: datetime.date,
    end: Optional[datetime.date] = None,
    client: Optional[RidReservoirClient] = None,
    throttle_seconds: float = 0.4,
) -> int:
    """Fetch every MISSING day in [start, end]; one polite request per day.

    Only dates absent from rid_reservoir_daily are requested, so a rerun
    repairs holes left by transient failures instead of resuming past them
    (the hourly collector writes today's rows immediately, which makes any
    newest-row cursor useless as a resume point). A save that persists
    nothing counts as a failure too — a broken DB must not burn thousands
    of requests against the RID API.
    """
    client = client or RidReservoirClient()
    end = end or datetime.date.today()
    if not store.engine and not store.connect():
        logger.error("backfill aborted: database connection failed")
        return 0
    span = [
        start + datetime.timedelta(days=i) for i in range((end - start).days + 1)
    ]
    present = store.present_dates(start, end)
    targets = [d for d in span if d not in present]
    logger.info(
        f"backfill: {len(targets)} of {len(span)} days missing in [{start}, {end}]"
    )
    total = 0
    failures = 0
    for i, date in enumerate(targets):
        try:
            records = client.fetch_day(date)
            saved = store.save(records)
            if records and not saved:
                raise RuntimeError("database save persisted 0 rows")
            total += saved
            failures = 0
        except Exception as e:
            failures += 1
            logger.warning(f"backfill {date} failed ({failures} in a row): {e}")
            if failures >= 5:
                logger.error("5 consecutive failures — aborting backfill")
                break
        if i % 100 == 0:
            logger.info(f"backfill progress: {date} ({total} rows)")
        time.sleep(throttle_seconds)
    return total


def backfill_dam(
    store: RidReservoirStore,
    dam_id: str = MAE_NGAT_DAM_ID,
    start: datetime.date = RID_ARCHIVE_START,
    end: Optional[datetime.date] = None,
    client: Optional[RidReservoirClient] = None,
    chunk_days: int = 1830,
    throttle_seconds: float = 1.0,
    skip_present: bool = True,
    stats: Optional[Dict] = None,
) -> int:
    """Backfill ONE dam over [start, end] using the range endpoint.

    Costs one request per chunk instead of one per calendar day: Mae Ngat's
    whole 2009-today archive is ~4 requests here versus ~6,400 with
    `backfill`. No server-side range limit was observed (2009-01-01..today
    answered in full), so `chunk_days` exists only to bound the response size
    and the time a single request can hang, not to satisfy the API.

    Chunks whose dates are already stored are skipped without a request, and
    returned rows are filtered to the missing dates, so a rerun repairs holes
    rather than rewriting the archive. `skip_present=False` re-fetches
    everything, which is how rows first written by the api/dams collector
    gain a level_msl and two-decimal storage_pct.

    Returns rows saved. Some dates stay missing however often this runs —
    the source simply never published them (72 days of Mae Ngat's archive,
    absent from api/dams too) — so a 0-row rerun is normal and callers must
    not read it as failure; pass `stats` to get the request/failure counts
    that actually distinguish an outage.
    """
    client = client or RidReservoirClient()
    end = end or datetime.date.today()
    counters = {"requests": 0, "failures": 0, "aborted": False}
    if stats is not None:
        stats.update(counters)
        counters = stats
    if not store.engine and not store.connect():
        logger.error("backfill_dam aborted: database connection failed")
        counters["aborted"] = True
        return 0
    span_days = (end - start).days + 1
    if span_days <= 0:
        return 0
    present = store.present_dates(start, end, dam_id=dam_id) if skip_present else set()
    logger.info(
        f"backfill_dam {dam_id}: {span_days - len(present)} of {span_days} "
        f"days missing in [{start}, {end}]"
    )
    total = 0
    failures = 0
    chunk_start = start
    while chunk_start <= end:
        chunk_end = min(chunk_start + datetime.timedelta(days=chunk_days - 1), end)
        wanted = {
            chunk_start + datetime.timedelta(days=i)
            for i in range((chunk_end - chunk_start).days + 1)
        } - present
        if not wanted:
            chunk_start = chunk_end + datetime.timedelta(days=1)
            continue
        try:
            counters["requests"] += 1
            records = client.fetch_dam_range(dam_id, chunk_start, chunk_end)
            records = [r for r in records if r["date"] in wanted]
            saved = store.save(records)
            if records and not saved:
                raise RuntimeError("database save persisted 0 rows")
            total += saved
            failures = 0
            logger.info(
                f"backfill_dam {dam_id} [{chunk_start}, {chunk_end}]: "
                f"{saved} rows ({total} total)"
            )
        except Exception as e:
            failures += 1
            counters["failures"] += 1
            logger.warning(
                f"backfill_dam {dam_id} [{chunk_start}, {chunk_end}] failed "
                f"({failures} in a row): {e}"
            )
            if failures >= 5:
                logger.error("5 consecutive failures — aborting backfill_dam")
                counters["aborted"] = True
                break
        chunk_start = chunk_end + datetime.timedelta(days=1)
        if chunk_start <= end:
            time.sleep(throttle_seconds)
    return total


def create_collector_from_config() -> Optional[RidReservoirCollector]:
    """Build a collector from app Config; None when disabled or non-SQL DB."""
    from .config import Config

    if not Config.ENABLE_RESERVOIR_COLLECTION:
        return None
    db_config = Config.get_database_config()
    if db_config["type"] not in ("sqlite", "postgresql", "mysql"):
        logger.warning(
            f"Reservoir collection skipped: DB_TYPE '{db_config['type']}' is not SQL"
        )
        return None
    return RidReservoirCollector(db_config)