from typing import List, Any, Optional
from datetime import datetime
from psycopg import AsyncConnection, sql

class RealtimeRepository:
    
    # --- Link Simulation ---
    
    @staticmethod
    async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
        """Batch insert for realtime.link_simulation using COPY for performance."""
        if not data:
            return

        async with conn.cursor() as cur:
            async with cur.copy(
                "COPY realtime.link_simulation (time, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
            ) as copy:
                for item in data:
                    await copy.write_row((
                        item['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 get_links_by_time(conn: AsyncConnection, start_time: datetime, end_time: datetime) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
                (start_time, end_time)
            )
            return await cur.fetchall()

    @staticmethod
    async def get_link_field(conn: AsyncConnection, time: datetime, link_id: str, field: str) -> Any:
        # Validate field name to prevent SQL injection
        valid_fields = {"flow", "friction", "headloss", "quality", "reaction", "setting", "status", "velocity"}
        if field not in valid_fields:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL("SELECT {} FROM realtime.link_simulation WHERE time = %s AND id = %s").format(sql.Identifier(field))
        
        async with conn.cursor() as cur:
            await cur.execute(query, (time, link_id))
            row = await cur.fetchone()
            return row[field] if row else None

    @staticmethod
    async def update_link_field(conn: AsyncConnection, time: datetime, link_id: str, field: str, value: Any):
        valid_fields = {"flow", "friction", "headloss", "quality", "reaction", "setting", "status", "velocity"}
        if field not in valid_fields:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL("UPDATE realtime.link_simulation SET {} = %s WHERE time = %s AND 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(conn: AsyncConnection, start_time: datetime, end_time: datetime):
        async with conn.cursor() as cur:
            await cur.execute(
                "DELETE FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
                (start_time, end_time)
            )

    # --- Node Simulation ---

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

        async with conn.cursor() as cur:
            async with cur.copy(
                "COPY realtime.node_simulation (time, id, actual_demand, total_head, pressure, quality) FROM STDIN"
            ) as copy:
                for item in data:
                    await copy.write_row((
                        item['time'], item['id'], item.get('actual_demand'), item.get('total_head'), item.get('pressure'), item.get('quality')
                    ))

    @staticmethod
    async def get_nodes_by_time(conn: AsyncConnection, start_time: datetime, end_time: datetime) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT * FROM realtime.node_simulation WHERE time >= %s AND time <= %s",
                (start_time, end_time)
            )
            return await cur.fetchall()
            
    @staticmethod
    async def get_node_field(conn: AsyncConnection, time: datetime, node_id: str, field: str) -> Any:
        valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
        if field not in valid_fields:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL("SELECT {} FROM realtime.node_simulation WHERE time = %s AND id = %s").format(sql.Identifier(field))
        
        async with conn.cursor() as cur:
            await cur.execute(query, (time, node_id))
            row = await cur.fetchone()
            return row[field] if row else None

    @staticmethod
    async def update_node_field(conn: AsyncConnection, time: datetime, node_id: str, field: str, value: Any):
        valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
        if field not in valid_fields:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL("UPDATE realtime.node_simulation SET {} = %s WHERE time = %s AND id = %s").format(sql.Identifier(field))
        
        async with conn.cursor() as cur:
            await cur.execute(query, (value, time, node_id))