TJWaterServerBinary/app/infra/db/timescaledb/repositories/realtime.py

from typing import List, Any, Dict
from datetime import datetime, timedelta
from collections import defaultdict
from psycopg import AsyncConnection, Connection, sql
from app.domain.time import parse_utc_time


class RealtimeRepository:
    LINK_FIELDS = frozenset(
        {
            "flow",
            "friction",
            "headloss",
            "quality",
            "reaction",
            "setting",
            "status",
            "velocity",
        }
    )
    NODE_FIELDS = frozenset({"actual_demand", "total_head", "pressure", "quality"})
    LINK_RESULT_COLUMNS = (
        "time, link_id, flow, friction, headloss, quality, reaction, "
        "setting, status, velocity"
    )
    NODE_RESULT_COLUMNS = "time, node_id, actual_demand, total_head, pressure, quality"

    @staticmethod
    def _batch_time(data: List[dict]) -> datetime:
        """Return one normalized timestamp shared by every row in a snapshot."""
        if not data:
            raise ValueError("Realtime batch must not be empty")
        times = {
            parse_utc_time(item["time"], field_name="time")
            for item in data
        }
        if len(times) != 1:
            raise ValueError("Realtime batch must contain exactly one timestamp")
        return times.pop()

    # --- Link Simulation ---

    @staticmethod
    async def _copy_links(cur, data: List[dict], target_time: datetime) -> None:
        async with cur.copy(
            "COPY realtime.link_results (time, link_id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
        ) as copy:
            for item in data:
                await copy.write_row(
                    (
                        target_time,
                        item["id"],
                        item.get("flow"),
                        item.get("friction"),
                        item.get("headloss"),
                        item.get("quality"),
                        item.get("reaction"),
                        item.get("setting"),
                        item.get("status"),
                        item.get("velocity"),
                    )
                )

    @staticmethod
    def _copy_links_sync(cur, data: List[dict], target_time: datetime) -> None:
        with cur.copy(
            "COPY realtime.link_results (time, link_id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
        ) as copy:
            for item in data:
                copy.write_row(
                    (
                        target_time,
                        item["id"],
                        item.get("flow"),
                        item.get("friction"),
                        item.get("headloss"),
                        item.get("quality"),
                        item.get("reaction"),
                        item.get("setting"),
                        item.get("status"),
                        item.get("velocity"),
                    )
                )

    @staticmethod
    async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
        """Batch insert for realtime.link_results using DELETE then COPY."""
        if not data:
            return

        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        async with conn.transaction():
            async with conn.cursor() as cur:
                await cur.execute(
                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
                    (target_time,),
                )
                # 1. 先删除该时间点的旧数据
                await cur.execute(
                    "DELETE FROM realtime.link_results WHERE time = %s",
                    (target_time,),
                )

                await RealtimeRepository._copy_links(cur, data, target_time)

    @staticmethod
    def insert_links_batch_sync(conn: Connection, data: List[dict]):
        """Synchronous batch insert for realtime.link_results."""
        if not data:
            return

        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
                    (target_time,),
                )
                # 1. 先删除该时间点的旧数据
                cur.execute(
                    "DELETE FROM realtime.link_results WHERE time = %s",
                    (target_time,),
                )

                RealtimeRepository._copy_links_sync(cur, data, target_time)

    @staticmethod
    async def get_link_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime, link_id: str
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT {RealtimeRepository.LINK_RESULT_COLUMNS} "
                "FROM realtime.link_results WHERE time >= %s AND time <= %s "
                "AND link_id = %s ORDER BY time",
                (start_time, end_time, link_id),
            )
            return await cur.fetchall()

    @staticmethod
    async def get_links_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime
    ) -> List[dict]:
        normalized_start_time = parse_utc_time(start_time, field_name="start_time")
        normalized_end_time = parse_utc_time(end_time, field_name="end_time")
        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT {RealtimeRepository.LINK_RESULT_COLUMNS} "
                "FROM realtime.link_results WHERE time >= %s AND time <= %s "
                "ORDER BY time, link_id",
                (normalized_start_time, normalized_end_time),
            )
            return await cur.fetchall()

    @staticmethod
    async def get_link_fields_by_ids_time_range(
        conn: AsyncConnection,
        start_time: datetime,
        end_time: datetime,
        link_ids: list[str],
        field: str,
    ) -> dict[str, list[dict[str, Any]]]:
        if field not in RealtimeRepository.LINK_FIELDS:
            raise ValueError(f"Invalid field: {field}")
        result = {link_id: [] for link_id in link_ids}
        if not link_ids:
            return result
        query = sql.SQL(
            "SELECT link_id, time, {} FROM realtime.link_results "
            "WHERE time BETWEEN %s AND %s AND link_id = ANY(%s) "
            "ORDER BY link_id, time"
        ).format(sql.Identifier(field))
        async with conn.cursor() as cur:
            await cur.execute(query, (start_time, end_time, link_ids))
            for row in await cur.fetchall():
                result.setdefault(str(row["link_id"]), []).append(
                    {"time": row["time"].isoformat(), "value": row[field]}
                )
        return result

    @staticmethod
    async def get_links_field_by_time_range(
        conn: AsyncConnection,
        start_time: datetime,
        end_time: datetime,
        field: str,
    ) -> dict:
        # Validate field name to prevent SQL injection
        if field not in RealtimeRepository.LINK_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
            "SELECT link_id, time, {} FROM realtime.link_results "
            "WHERE time >= %s AND time <= %s ORDER BY link_id, time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(query, (start_time, end_time))
            rows = await cur.fetchall()
            result = defaultdict(list)
            for row in rows:
                result[row["link_id"]].append(
                    {"time": row["time"].isoformat(), "value": row[field]}
                )
            return dict(result)

    @staticmethod
    async def update_link_field(
        conn: AsyncConnection,
        time: datetime,
        link_id: str,
        field: str,
        value: Any,
    ):
        if field not in RealtimeRepository.LINK_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
            "UPDATE realtime.link_results SET {} = %s WHERE time = %s AND link_id = %s"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(query, (value, time, link_id))

    @staticmethod
    async def delete_links_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime
    ):
        async with conn.cursor() as cur:
            await cur.execute(
                "DELETE FROM realtime.link_results WHERE time >= %s AND time <= %s",
                (start_time, end_time),
            )

    # --- Node Simulation ---

    @staticmethod
    async def _copy_nodes(cur, data: List[dict], target_time: datetime) -> None:
        async with cur.copy(
            "COPY realtime.node_results (time, node_id, actual_demand, total_head, pressure, quality) FROM STDIN"
        ) as copy:
            for item in data:
                await copy.write_row(
                    (
                        target_time,
                        item["id"],
                        item.get("actual_demand"),
                        item.get("total_head"),
                        item.get("pressure"),
                        item.get("quality"),
                    )
                )

    @staticmethod
    def _copy_nodes_sync(cur, data: List[dict], target_time: datetime) -> None:
        with cur.copy(
            "COPY realtime.node_results (time, node_id, actual_demand, total_head, pressure, quality) FROM STDIN"
        ) as copy:
            for item in data:
                copy.write_row(
                    (
                        target_time,
                        item["id"],
                        item.get("actual_demand"),
                        item.get("total_head"),
                        item.get("pressure"),
                        item.get("quality"),
                    )
                )

    @staticmethod
    async def insert_nodes_batch(conn: AsyncConnection, data: List[dict]):
        if not data:
            return

        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        async with conn.transaction():
            async with conn.cursor() as cur:
                await cur.execute(
                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
                    (target_time,),
                )
                # 1. 先删除该时间点的旧数据
                await cur.execute(
                    "DELETE FROM realtime.node_results WHERE time = %s",
                    (target_time,),
                )

                await RealtimeRepository._copy_nodes(cur, data, target_time)

    @staticmethod
    def insert_nodes_batch_sync(conn: Connection, data: List[dict]):
        if not data:
            return

        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
                    (target_time,),
                )
                # 1. 先删除该时间点的旧数据
                cur.execute(
                    "DELETE FROM realtime.node_results WHERE time = %s",
                    (target_time,),
                )

                RealtimeRepository._copy_nodes_sync(cur, data, target_time)

    @staticmethod
    async def get_node_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime, node_id: str
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT {RealtimeRepository.NODE_RESULT_COLUMNS} "
                "FROM realtime.node_results WHERE time >= %s AND time <= %s "
                "AND node_id = %s ORDER BY time",
                (start_time, end_time, node_id),
            )
            return await cur.fetchall()

    @staticmethod
    async def get_nodes_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime
    ) -> List[dict]:
        normalized_start_time = parse_utc_time(start_time, field_name="start_time")
        normalized_end_time = parse_utc_time(end_time, field_name="end_time")
        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT {RealtimeRepository.NODE_RESULT_COLUMNS} "
                "FROM realtime.node_results WHERE time >= %s AND time <= %s "
                "ORDER BY time, node_id",
                (normalized_start_time, normalized_end_time),
            )
            return await cur.fetchall()

    @staticmethod
    async def get_node_fields_by_ids_time_range(
        conn: AsyncConnection,
        start_time: datetime,
        end_time: datetime,
        node_ids: list[str],
        field: str,
    ) -> dict[str, list[dict[str, Any]]]:
        if field not in RealtimeRepository.NODE_FIELDS:
            raise ValueError(f"Invalid field: {field}")
        result = {node_id: [] for node_id in node_ids}
        if not node_ids:
            return result
        query = sql.SQL(
            "SELECT node_id, time, {} FROM realtime.node_results "
            "WHERE time BETWEEN %s AND %s AND node_id = ANY(%s) "
            "ORDER BY node_id, time"
        ).format(sql.Identifier(field))
        async with conn.cursor() as cur:
            await cur.execute(query, (start_time, end_time, node_ids))
            for row in await cur.fetchall():
                result.setdefault(str(row["node_id"]), []).append(
                    {"time": row["time"].isoformat(), "value": row[field]}
                )
        return result

    @staticmethod
    async def get_nodes_field_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime, field: str
    ) -> dict:
        if field not in RealtimeRepository.NODE_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
            "SELECT node_id, time, {} FROM realtime.node_results "
            "WHERE time >= %s AND time <= %s ORDER BY node_id, time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(query, (start_time, end_time))
            rows = await cur.fetchall()
            result = defaultdict(list)
            for row in rows:
                result[row["node_id"]].append(
                    {"time": row["time"].isoformat(), "value": row[field]}
                )
            return dict(result)

    @staticmethod
    async def update_node_field(
        conn: AsyncConnection,
        time: datetime,
        node_id: str,
        field: str,
        value: Any,
    ):
        if field not in RealtimeRepository.NODE_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
            "UPDATE realtime.node_results SET {} = %s WHERE time = %s AND node_id = %s"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(query, (value, time, node_id))

    @staticmethod
    async def delete_nodes_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime
    ):
        async with conn.cursor() as cur:
            await cur.execute(
                "DELETE FROM realtime.node_results WHERE time >= %s AND time <= %s",
                (start_time, end_time),
            )

    # --- 复合查询 ---

    @staticmethod
    async def store_realtime_simulation_result(
        conn: AsyncConnection,
        node_result_list: List[Dict[str, any]],
        link_result_list: List[Dict[str, any]],
        result_start_time: str,
    ):
        """
        Store realtime simulation results to TimescaleDB.

        Args:
            conn: Database connection
            node_result_list: List of node simulation results
            link_result_list: List of link simulation results
            result_start_time: Start time for the results (ISO format string)
        """
        simulation_time = parse_utc_time(
            result_start_time, field_name="result_start_time"
        )

        # Prepare node data for batch insert
        node_data = []
        for node_result in node_result_list:
            node_id = node_result.get("node")
            data = node_result.get("result", [])[0]  # 实时模拟只有一个周期
            node_data.append(
                {
                    "time": simulation_time,
                    "id": node_id,
                    "actual_demand": data.get("demand"),
                    "total_head": data.get("head"),
                    "pressure": data.get("pressure"),
                    "quality": data.get("quality"),
                }
            )

        # Prepare link data for batch insert
        link_data = []
        for link_result in link_result_list:
            link_id = link_result.get("link")
            data = link_result.get("result", [])[0]
            link_data.append(
                {
                    "time": simulation_time,
                    "id": link_id,
                    "flow": data.get("flow"),
                    "friction": data.get("friction"),
                    "headloss": data.get("headloss"),
                    "quality": data.get("quality"),
                    "reaction": data.get("reaction"),
                    "setting": data.get("setting"),
                    "status": data.get("status"),
                    "velocity": data.get("velocity"),
                }
            )

        # Keep node and link replacement atomic with one lock and one delete per
        # table. Copy-only helpers avoid repeating replacement SQL.
        async with conn.transaction():
            async with conn.cursor() as cur:
                await cur.execute(
                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
                    (simulation_time,),
                )
                await cur.execute(
                    "DELETE FROM realtime.node_results WHERE time = %s",
                    (simulation_time,),
                )
                await cur.execute(
                    "DELETE FROM realtime.link_results WHERE time = %s",
                    (simulation_time,),
                )
                if node_data:
                    await RealtimeRepository._copy_nodes(cur, node_data, simulation_time)
                if link_data:
                    await RealtimeRepository._copy_links(cur, link_data, simulation_time)

    @staticmethod
    def store_realtime_simulation_result_sync(
        conn: Connection,
        node_result_list: List[Dict[str, any]],
        link_result_list: List[Dict[str, any]],
        result_start_time: str,
    ):
        """
        Store realtime simulation results to TimescaleDB (sync version).

        Args:
            conn: Database connection
            node_result_list: List of node simulation results
            link_result_list: List of link simulation results
            result_start_time: Start time for the results (ISO format string)
        """
        simulation_time = parse_utc_time(
            result_start_time, field_name="result_start_time"
        )

        # Prepare node data for batch insert
        node_data = []
        for node_result in node_result_list:
            node_id = node_result.get("node")
            data = node_result.get("result", [])[0]  # 实时模拟只有一个周期
            node_data.append(
                {
                    "time": simulation_time,
                    "id": node_id,
                    "actual_demand": data.get("demand"),
                    "total_head": data.get("head"),
                    "pressure": data.get("pressure"),
                    "quality": data.get("quality"),
                }
            )

        # Prepare link data for batch insert
        link_data = []
        for link_result in link_result_list:
            link_id = link_result.get("link")
            data = link_result.get("result", [])[0]
            link_data.append(
                {
                    "time": simulation_time,
                    "id": link_id,
                    "flow": data.get("flow"),
                    "friction": data.get("friction"),
                    "headloss": data.get("headloss"),
                    "quality": data.get("quality"),
                    "reaction": data.get("reaction"),
                    "setting": data.get("setting"),
                    "status": data.get("status"),
                    "velocity": data.get("velocity"),
                }
            )

        # Keep node and link replacement atomic with one lock and one delete per
        # table. Copy-only helpers avoid repeating replacement SQL.
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
                    (simulation_time,),
                )
                cur.execute(
                    "DELETE FROM realtime.node_results WHERE time = %s",
                    (simulation_time,),
                )
                cur.execute(
                    "DELETE FROM realtime.link_results WHERE time = %s",
                    (simulation_time,),
                )
                if node_data:
                    RealtimeRepository._copy_nodes_sync(cur, node_data, simulation_time)
                if link_data:
                    RealtimeRepository._copy_links_sync(cur, link_data, simulation_time)

    @staticmethod
    async def query_all_record_by_time_property(
        conn: AsyncConnection,
        query_time: str,
        type: str,
        property: str,
    ) -> list:
        """
        Query all records by time and property from TimescaleDB.

        Args:
            conn: Database connection
            query_time: Time to query (ISO format string)
            type: Type of data ("node" or "link")
            property: Property/field to query

        Returns:
            List of records matching the criteria
        """
        target_time = parse_utc_time(query_time, field_name="query_time")

        # Create time range: query_time ± 1 second
        start_time = target_time - timedelta(seconds=1)
        end_time = target_time + timedelta(seconds=1)

        # Query based on type
        if type.lower() == "node":
            data = await RealtimeRepository.get_nodes_field_by_time_range(
                conn, start_time, end_time, property
            )
        elif type.lower() == "link":
            data = await RealtimeRepository.get_links_field_by_time_range(
                conn, start_time, end_time, property
            )
        else:
            raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")

        # Format the results
        result = []
        for id, items in data.items():
            for item in items:
                result.append({"ID": id, "value": item["value"]})
        return result

    @staticmethod
    async def query_simulation_result_by_id_time(
        conn: AsyncConnection,
        id: str,
        type: str,
        query_time: str,
    ) -> list[dict]:
        """
        Query simulation results by id and time from TimescaleDB.

        Args:
            conn: Database connection
            id: The id of the node or link
            type: Type of data ("node" or "link")
            query_time: Time to query (ISO format string)

        Returns:
            List of records matching the criteria
        """
        target_time = parse_utc_time(query_time, field_name="query_time")

        # Create time range: query_time ± 1 second
        start_time = target_time - timedelta(seconds=1)
        end_time = target_time + timedelta(seconds=1)

        # Query based on type
        if type.lower() == "node":
            return await RealtimeRepository.get_node_by_time_range(
                conn, start_time, end_time, id
            )
        elif type.lower() == "link":
            return await RealtimeRepository.get_link_by_time_range(
                conn, start_time, end_time, id
            )
        else:
            raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")