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

"""Collector for HII/ThaiWater open api-v3 feeds (rainfall + water level).

Polls the unauthenticated api-v3.thaiwater.net public endpoints, filters to the
Ping basin, and persists to dedicated tables alongside the RID data:

- hii_rain_stations / hii_rainfall        (rain_1h / rain_24h gauge telemetry)
- hii_wl_stations   / hii_waterlevel      (independent water-level source, m MSL)

Water levels are kept in a separate table (not a column on water_measurements)
because HII reports in m MSL from a different station universe; the per-station
``offset`` column (gauge zero in m MSL) converts to gauge datum when needed.

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

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

import requests

logger = logging.getLogger(__name__)

HII_API_BASE = "https://api-v3.thaiwater.net/api/v1/thaiwater30/public"
PING_BASIN_CODE = 6

# Matches 'P.1', 'ridhydro_P.67', 'ridtele_TUP.14' -> canonical RID code suffix
_RID_CODE_RE = re.compile(r"(?:^|_)(P\.\d+[A-Z]?)$")


def _to_float(value: Any) -> Optional[float]:
    """API numerics arrive as strings ('335.00'), numbers, or None."""
    if value is None or value == "":
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _parse_datetime(value: Any) -> Optional[datetime.datetime]:
    """Timestamps are Thai local time, e.g. '2026-08-11 13:00'."""
    if not value:
        return None
    for fmt in ("%Y-%m-%d %H:%M", "%Y-%m-%d %H:%M:%S"):
        try:
            return datetime.datetime.strptime(value, fmt)
        except ValueError:
            continue
    return None


def _name(station: Dict, lang: str) -> Optional[str]:
    name = station.get("tele_station_name")
    if isinstance(name, dict):
        return name.get(lang)
    return name if lang == "th" else None


def rid_code_from_oldcode(oldcode: Optional[str]) -> Optional[str]:
    """Normalize a ThaiWater oldcode to the RID P-code it mirrors, if any."""
    if not oldcode:
        return None
    match = _RID_CODE_RE.search(oldcode)
    return match.group(1) if match else None


def parse_rain_records(
    payload: Dict, basin_code: int = PING_BASIN_CODE
) -> List[Dict]:
    """Extract per-station rainfall rows from a rain_24h payload."""
    records = []
    for row in payload.get("data") or []:
        basin = row.get("basin") or {}
        if basin.get("basin_code") != basin_code:
            continue
        station = row.get("station") or {}
        station_id = station.get("id")
        timestamp = _parse_datetime(row.get("rainfall_datetime"))
        if station_id is None or timestamp is None:
            continue
        records.append(
            {
                "station_id": station_id,
                "oldcode": station.get("tele_station_oldcode"),
                "name_th": _name(station, "th"),
                "name_en": _name(station, "en"),
                "latitude": _to_float(station.get("tele_station_lat")),
                "longitude": _to_float(station.get("tele_station_long")),
                "sub_basin_id": str(station.get("sub_basin_id") or "") or None,
                "agency": ((row.get("agency") or {}).get("agency_shortname") or {}).get(
                    "en"
                ),
                "timestamp": timestamp,
                "rain_1h": _to_float(row.get("rain_1h")),
                "rain_24h": _to_float(row.get("rain_24h")),
            }
        )
    return records


def parse_waterlevel_records(
    payload: Dict, basin_code: int = PING_BASIN_CODE
) -> List[Dict]:
    """Extract per-station water-level rows from a waterlevel_load payload."""
    data = (payload.get("waterlevel_data") or {}).get("data") or []
    records = []
    for row in data:
        basin = row.get("basin") or {}
        if basin.get("basin_code") != basin_code:
            continue
        station = row.get("station") or {}
        station_id = station.get("id")
        timestamp = _parse_datetime(row.get("waterlevel_datetime"))
        if station_id is None or timestamp is None:
            continue
        oldcode = station.get("tele_station_oldcode")
        records.append(
            {
                "station_id": station_id,
                "oldcode": oldcode,
                "rid_code": rid_code_from_oldcode(oldcode),
                "name_th": _name(station, "th"),
                "name_en": _name(station, "en"),
                "latitude": _to_float(station.get("tele_station_lat")),
                "longitude": _to_float(station.get("tele_station_long")),
                "agency": ((row.get("agency") or {}).get("agency_shortname") or {}).get(
                    "en"
                ),
                "river_name": row.get("river_name"),
                "offset_msl": _to_float(station.get("offset")),
                "ground_level_msl": _to_float(station.get("ground_level")),
                "min_bank_msl": _to_float(station.get("min_bank")),
                "critical_level_msl": _to_float(station.get("critical_level_msl")),
                "critical_level_m": _to_float(station.get("critical_level_m")),
                "qmax": _to_float(station.get("qmax")),
                "is_key_station": bool(station.get("is_key_station")),
                "timestamp": timestamp,
                "wl_msl": _to_float(row.get("waterlevel_msl")),
                "wl_m": _to_float(row.get("waterlevel_m")),
                "discharge": _to_float(row.get("discharge")),
                "flow_rate": _to_float(row.get("flow_rate")),
                "storage_percent": _to_float(row.get("storage_percent")),
                "situation_level": row.get("situation_level"),
                "diff_wl_bank": _to_float(row.get("diff_wl_bank")),
            }
        )
    return records


class HiiClient:
    """HTTP client for the open api-v3 public endpoints."""

    def __init__(
        self,
        base_url: str = HII_API_BASE,
        session: Optional[requests.Session] = None,
        timeout: int = 90,
    ):
        self.base_url = base_url.rstrip("/")
        self.session = session or requests.Session()
        self.timeout = timeout

    def get(self, endpoint: str, params: Optional[Dict] = None) -> Dict:
        response = self.session.get(
            f"{self.base_url}/{endpoint}", params=params, timeout=self.timeout
        )
        response.raise_for_status()
        return response.json()

    def fetch_rain(self, basin_code: int = PING_BASIN_CODE) -> List[Dict]:
        return parse_rain_records(self.get("rain_24h"), basin_code)

    def fetch_waterlevel(self, basin_code: int = PING_BASIN_CODE) -> List[Dict]:
        return parse_waterlevel_records(self.get("waterlevel_load"), basin_code)


class HiiStore:
    """SQL persistence for HII feeds (sqlite / postgresql / mysql).

    Reuses the app's main relational database (same connection string as
    the RID tables) but writes to its own hii_* tables.
    """

    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"HII 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"HiiStore failed to connect: {e}")
            self.engine = None
            return False

    def _create_tables(self):
        from sqlalchemy import text

        bool_type = "BOOLEAN" if self.db_type != "mysql" else "TINYINT(1)"
        ddl = [
            """
            CREATE TABLE IF NOT EXISTS hii_rain_stations (
                id INTEGER PRIMARY KEY,
                oldcode VARCHAR(60),
                name_th VARCHAR(255),
                name_en VARCHAR(255),
                latitude NUMERIC(10,6),
                longitude NUMERIC(10,6),
                sub_basin_id VARCHAR(10),
                agency VARCHAR(40),
                updated_at TIMESTAMP
            )
            """,
            # Composite natural PK (no surrogate id): TimescaleDB hypertable
            # conversion requires every unique index to include the time column.
            """
            CREATE TABLE IF NOT EXISTS hii_rainfall (
                station_id INTEGER NOT NULL,
                timestamp TIMESTAMP NOT NULL,
                rain_1h NUMERIC(7,2),
                rain_24h NUMERIC(8,2),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (station_id, timestamp)
            )
            """,
            f"""
            CREATE TABLE IF NOT EXISTS hii_wl_stations (
                id INTEGER PRIMARY KEY,
                oldcode VARCHAR(60),
                rid_code VARCHAR(10),
                name_th VARCHAR(255),
                name_en VARCHAR(255),
                latitude NUMERIC(10,6),
                longitude NUMERIC(10,6),
                agency VARCHAR(40),
                river_name VARCHAR(255),
                offset_msl NUMERIC(8,3),
                ground_level_msl NUMERIC(8,3),
                min_bank_msl NUMERIC(8,3),
                critical_level_msl NUMERIC(8,3),
                critical_level_m NUMERIC(8,3),
                qmax NUMERIC(10,2),
                is_key_station {bool_type},
                updated_at TIMESTAMP
            )
            """,
            """
            CREATE TABLE IF NOT EXISTS hii_waterlevel (
                station_id INTEGER NOT NULL,
                timestamp TIMESTAMP NOT NULL,
                wl_msl NUMERIC(8,3),
                wl_m NUMERIC(8,3),
                discharge NUMERIC(10,2),
                flow_rate NUMERIC(10,2),
                storage_percent NUMERIC(6,2),
                situation_level INTEGER,
                diff_wl_bank NUMERIC(8,3),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (station_id, timestamp)
            )
            """,
            "CREATE INDEX IF NOT EXISTS idx_hii_rainfall_ts ON hii_rainfall(timestamp)",
            "CREATE INDEX IF NOT EXISTS idx_hii_waterlevel_ts ON hii_waterlevel(timestamp)",
        ]
        # MySQL (<8.0.13 semantics) lacks CREATE INDEX IF NOT EXISTS; the unique
        # constraints already cover the hot (station_id, timestamp) lookups there.
        if self.db_type == "mysql":
            ddl = ddl[:4]
        with self.engine.begin() as conn:
            for statement in ddl:
                conn.execute(text(statement))

    def _upsert(self, table: str, key_cols: List[str], value_cols: List[str]) -> str:
        cols = key_cols + value_cols
        col_list = ", ".join(cols)
        params = ", ".join(f":{c}" for c in cols)
        if self.db_type == "sqlite":
            return f"INSERT OR REPLACE INTO {table} ({col_list}) VALUES ({params})"
        if self.db_type == "postgresql":
            updates = ", ".join(f"{c} = EXCLUDED.{c}" for c in value_cols)
            conflict = ", ".join(key_cols)
            return (
                f"INSERT INTO {table} ({col_list}) VALUES ({params}) "
                f"ON CONFLICT ({conflict}) DO UPDATE SET {updates}"
            )
        updates = ", ".join(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_rain(self, records: List[Dict]) -> int:
        return self._save(
            records,
            station_table="hii_rain_stations",
            station_cols=[
                "oldcode",
                "name_th",
                "name_en",
                "latitude",
                "longitude",
                "sub_basin_id",
                "agency",
            ],
            measurement_table="hii_rainfall",
            measurement_cols=["rain_1h", "rain_24h"],
        )

    def save_waterlevel(self, records: List[Dict]) -> int:
        return self._save(
            records,
            station_table="hii_wl_stations",
            station_cols=[
                "oldcode",
                "rid_code",
                "name_th",
                "name_en",
                "latitude",
                "longitude",
                "agency",
                "river_name",
                "offset_msl",
                "ground_level_msl",
                "min_bank_msl",
                "critical_level_msl",
                "critical_level_m",
                "qmax",
                "is_key_station",
            ],
            measurement_table="hii_waterlevel",
            measurement_cols=[
                "wl_msl",
                "wl_m",
                "discharge",
                "flow_rate",
                "storage_percent",
                "situation_level",
                "diff_wl_bank",
            ],
        )

    def save_waterlevel_history(self, station_id: int, rows: List[Dict]) -> int:
        """Upsert backfilled history rows, touching only wl_msl and discharge.

        Live-snapshot rows for the same (station, hour) keep their extra
        columns (storage_percent, situation_level, ...) untouched.
        """
        if not rows:
            return 0
        if not self.engine and not self.connect():
            return 0
        from sqlalchemy import text

        cols = "(station_id, timestamp, wl_msl, discharge)"
        values = "(:station_id, :timestamp, :wl_msl, :discharge)"
        if self.db_type == "mysql":
            sql = (
                f"INSERT INTO hii_waterlevel {cols} VALUES {values} "
                "ON DUPLICATE KEY UPDATE wl_msl = VALUES(wl_msl), "
                "discharge = VALUES(discharge)"
            )
        else:  # sqlite (>=3.24) and postgresql share upsert syntax
            sql = (
                f"INSERT INTO hii_waterlevel {cols} VALUES {values} "
                "ON CONFLICT (station_id, timestamp) DO UPDATE SET "
                "wl_msl = EXCLUDED.wl_msl, discharge = EXCLUDED.discharge"
            )
        params = [{**row, "station_id": station_id} for row in rows]
        try:
            with self.engine.begin() as conn:
                conn.execute(text(sql), params)
            return len(params)
        except Exception as e:
            logger.error(f"HiiStore history save failed: {e}")
            return 0

    def _save(
        self,
        records: List[Dict],
        station_table: str,
        station_cols: List[str],
        measurement_table: str,
        measurement_cols: List[str],
    ) -> int:
        if not records:
            return 0
        if not self.engine and not self.connect():
            return 0
        from sqlalchemy import text

        now = datetime.datetime.now()
        station_sql = self._upsert(
            station_table, ["id"], station_cols + ["updated_at"]
        )
        measurement_sql = self._upsert(
            measurement_table, ["station_id", "timestamp"], measurement_cols
        )
        # Dedupe stations (one row per station per snapshot anyway) and build
        # parameter dicts limited to each statement's columns.
        stations = {}
        measurements = []
        for record in records:
            sid = record["station_id"]
            station_row = {c: record.get(c) for c in station_cols}
            station_row.update({"id": sid, "updated_at": now})
            stations[sid] = station_row
            measurement_row = {c: record.get(c) for c in measurement_cols}
            measurement_row.update(
                {"station_id": sid, "timestamp": record["timestamp"]}
            )
            measurements.append(measurement_row)

        try:
            with self.engine.begin() as conn:
                conn.execute(text(station_sql), list(stations.values()))
                conn.execute(text(measurement_sql), measurements)
            return len(measurements)
        except Exception as e:
            logger.error(f"HiiStore save to {measurement_table} failed: {e}")
            return 0


class HiiCollector:
    """Fetch + persist one snapshot of both HII feeds."""

    def __init__(
        self,
        db_config: Dict,
        basin_code: int = PING_BASIN_CODE,
        client: Optional[HiiClient] = None,
    ):
        self.client = client or HiiClient()
        self.basin_code = basin_code
        self.store = HiiStore(
            connection_string=db_config["connection_string"],
            db_type=db_config["type"],
        )

    def run_cycle(self) -> Dict[str, int]:
        """Collect both feeds; each is independent and failure-isolated."""
        counts = {"rainfall": 0, "waterlevel": 0}
        try:
            counts["rainfall"] = self.store.save_rain(
                self.client.fetch_rain(self.basin_code)
            )
        except Exception as e:
            logger.error(f"HII rainfall collection failed: {e}")
        try:
            counts["waterlevel"] = self.store.save_waterlevel(
                self.client.fetch_waterlevel(self.basin_code)
            )
        except Exception as e:
            logger.error(f"HII waterlevel collection failed: {e}")
        logger.info(
            f"HII collection: {counts['rainfall']} rainfall, "
            f"{counts['waterlevel']} waterlevel rows saved"
        )
        return counts


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

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