HexaWetter/api/app/database.py

from __future__ import annotations

import json
import logging
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any

import asyncpg

from app.config import DATABASE_URL

logger = logging.getLogger(__name__)
_pool: asyncpg.Pool | None = None


async def init_db() -> None:
    global _pool
    try:
        _pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=8, command_timeout=30)
        async with _pool.acquire() as conn:
            await conn.execute("SELECT 1")
        logger.info("PostgreSQL connected")
    except Exception as exc:
        logger.warning("PostgreSQL unavailable, running without DB cache: %s", exc)
        _pool = None


async def close_db() -> None:
    global _pool
    if _pool is not None:
        await _pool.close()
        _pool = None


def pool_available() -> bool:
    return _pool is not None


@asynccontextmanager
async def acquire():
    if _pool is None:
        yield None
        return
    async with _pool.acquire() as conn:
        yield conn


async def get_cached_json(table: str, key_cols: dict[str, Any], max_age_seconds: int) -> Any | None:
    async with acquire() as conn:
        if conn is None:
            return None
        if table == "forecast_cache":
            row = await conn.fetchrow(
                """
                SELECT data, fetched_at FROM forecast_cache
                WHERE lat = $1 AND lon = $2
                  AND fetched_at > NOW() - ($3 || ' seconds')::interval
                """,
                key_cols["lat"],
                key_cols["lon"],
                str(max_age_seconds),
            )
        elif table == "observations_cache":
            row = await conn.fetchrow(
                """
                SELECT data, fetched_at FROM observations_cache
                WHERE lat = $1 AND lon = $2
                  AND fetched_at > NOW() - ($3 || ' seconds')::interval
                """,
                key_cols["lat"],
                key_cols["lon"],
                str(max_age_seconds),
            )
        elif table == "geocode_cache":
            row = await conn.fetchrow(
                """
                SELECT results, fetched_at FROM geocode_cache
                WHERE query_hash = $1
                  AND fetched_at > NOW() - ($2 || ' seconds')::interval
                """,
                key_cols["query_hash"],
                str(max_age_seconds),
            )
        else:
            return None
        if not row:
            return None
        data = row["data"] if "data" in row else row["results"]
        return json.loads(data) if isinstance(data, str) else data


async def set_cached_json(table: str, key_cols: dict[str, Any], payload: Any) -> None:
    async with acquire() as conn:
        if conn is None:
            return
        payload_json = json.dumps(payload)
        if table == "forecast_cache":
            await conn.execute(
                """
                INSERT INTO forecast_cache (lat, lon, data, fetched_at)
                VALUES ($1, $2, $3::jsonb, NOW())
                ON CONFLICT (lat, lon) DO UPDATE
                SET data = EXCLUDED.data, fetched_at = NOW()
                """,
                key_cols["lat"],
                key_cols["lon"],
                payload_json,
            )
        elif table == "observations_cache":
            await conn.execute(
                """
                INSERT INTO observations_cache (lat, lon, data, fetched_at)
                VALUES ($1, $2, $3::jsonb, NOW())
                ON CONFLICT (lat, lon) DO UPDATE
                SET data = EXCLUDED.data, fetched_at = NOW()
                """,
                key_cols["lat"],
                key_cols["lon"],
                payload_json,
            )
        elif table == "geocode_cache":
            await conn.execute(
                """
                INSERT INTO geocode_cache (query_hash, query_text, results, fetched_at)
                VALUES ($1, $2, $3::jsonb, NOW())
                ON CONFLICT (query_hash) DO UPDATE
                SET results = EXCLUDED.results, fetched_at = NOW(), query_text = EXCLUDED.query_text
                """,
                key_cols["query_hash"],
                key_cols.get("query_text", ""),
                payload_json,
            )
        elif table == "warnings_cache":
            await conn.execute(
                "INSERT INTO warnings_cache (data, fetched_at) VALUES ($1::jsonb, NOW())",
                payload_json,
            )


async def upsert_radar_meta(
    product: str,
    filename: str,
    file_size: int | None,
    modified_at: datetime | None,
    render_status: str = "pending",
    render_path: str | None = None,
    render_bounds: dict | None = None,
    render_error: str | None = None,
) -> None:
    async with acquire() as conn:
        if conn is None:
            return
        await conn.execute(
            """
            INSERT INTO radar_file_meta (
                product, filename, file_size, modified_at, render_status,
                render_path, render_bounds, render_error, updated_at
            ) VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, NOW())
            ON CONFLICT (filename) DO UPDATE SET
                file_size = EXCLUDED.file_size,
                modified_at = EXCLUDED.modified_at,
                render_status = EXCLUDED.render_status,
                render_path = EXCLUDED.render_path,
                render_bounds = EXCLUDED.render_bounds,
                render_error = EXCLUDED.render_error,
                updated_at = NOW()
            """,
            product,
            filename,
            file_size,
            modified_at,
            render_status,
            render_path,
            json.dumps(render_bounds) if render_bounds else None,
            render_error,
        )


async def list_rendered_radar(product: str, limit: int = 24) -> list[dict[str, Any]]:
    async with acquire() as conn:
        if conn is None:
            return []
        rows = await conn.fetch(
            """
            SELECT filename, render_path, render_bounds, modified_at, updated_at
            FROM radar_file_meta
            WHERE product = $1 AND render_status = 'done' AND render_path IS NOT NULL
            ORDER BY modified_at DESC NULLS LAST
            LIMIT $2
            """,
            product,
            limit,
        )
        result = []
        for row in rows:
            bounds = row["render_bounds"]
            if isinstance(bounds, str):
                bounds = json.loads(bounds)
            result.append(
                {
                    "filename": row["filename"],
                    "image_url": f"/api/admin/radar/rendered/{product}/{row['filename']}.png",
                    "bounds": bounds,
                    "modified_at": row["modified_at"].isoformat() if row["modified_at"] else None,
                    "rendered_at": row["updated_at"].isoformat() if row["updated_at"] else None,
                }
            )
        return result


def utc_now() -> datetime:
    return datetime.now(timezone.utc)