"""Read historical station measurements from PostgreSQL."""

import datetime
import os
from typing import Dict, List, Optional, Tuple

from sqlalchemy import create_engine, text

# Stage-discharge rating curves: Q = a * (H - b)^c
# Key: station_code, Value: (a, b, c)
# Use linear fallback Q = slope * H if a curve is not defined.
_RATING_CURVES: Dict[str, Tuple[float, float, float]] = {}
_DEFAULT_LINEAR_SLOPE = 20.0  # m^3/s per meter


def _calculate_discharge(
    water_level: Optional[float], station_code: str = None
) -> Optional[float]:
    """Estimate discharge from water level using a rating curve or linear fallback."""
    if water_level is None:
        return None

    curve = _RATING_CURVES.get(station_code)
    if curve:
        a, b, c = curve
        h_excess = water_level - b
        if h_excess <= 0:
            return 0.0
        return round(a * (h_excess**c), 2)

    # Linear fallback: Q = slope * H
    return round(_DEFAULT_LINEAR_SLOPE * water_level, 2)


class PostgresHistory:
    def __init__(self, connection_string: Optional[str] = None, engine=None):
        connection_string = connection_string or os.getenv("POSTGRES_CONNECTION_STRING")
        if engine is None and not connection_string:
            raise RuntimeError("POSTGRES_CONNECTION_STRING is not configured")
        self.engine = engine or create_engine(connection_string, pool_pre_ping=True)

    def station_history(
        self,
        station_code: str,
        start: datetime.datetime,
        end: datetime.datetime,
        limit: int = 2000,
    ) -> List[Dict]:
        if not 1 <= limit <= 100000:
            raise ValueError("limit must be between 1 and 100000")
        if start >= end:
            raise ValueError("start must be before end")

        query = text(
            """
            SELECT m.timestamp, s.station_code, m.water_level,
                   m.discharge, m.discharge_percent
            FROM water_measurements m
            JOIN stations s ON m.station_id = s.id
            WHERE s.station_code = :station_code
              AND m.timestamp >= :start_time
              AND m.timestamp <= :end_time
            ORDER BY m.timestamp ASC
            LIMIT :limit
            """
        )
        with self.engine.connect() as connection:
            rows = connection.execute(
                query,
                {
                    "station_code": station_code,
                    "start_time": start,
                    "end_time": end,
                    "limit": limit,
                },
            )
            result = []
            for row in rows:
                timestamp = row[0]
                if isinstance(timestamp, str):
                    timestamp = datetime.datetime.fromisoformat(timestamp)
                station_code = row[1]
                water_level = float(row[2]) if row[2] is not None else None
                discharge = float(row[3]) if row[3] is not None else None
                # Estimate discharge from water level if DB value is missing
                if discharge is None and water_level is not None:
                    discharge = _calculate_discharge(water_level, station_code)
                result.append(
                    {
                        "timestamp": timestamp,
                        "station_code": station_code,
                        "water_level": water_level,
                        "discharge": discharge,
                        "discharge_percent": float(row[4])
                        if row[4] is not None
                        else None,
                    }
                )
            return result