from datetime import datetime, timedelta
from typing import Any
from uuid import UUID

from psycopg import AsyncConnection, Connection, sql

from app.domain.time import parse_utc_time


class AnalysisResultsRepository:
    NODE_FIELDS = {"actual_demand", "total_head", "pressure", "quality"}
    LINK_FIELDS = {
        "flow",
        "friction",
        "headloss",
        "quality",
        "reaction",
        "setting",
        "status",
        "velocity",
    }

    @staticmethod
    def prepare_simulation_rows(
        node_results: list[dict[str, Any]],
        link_results: list[dict[str, Any]],
        result_start_time: str,
        num_periods: int,
        result_timestep_seconds: int,
    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
        start_time = parse_utc_time(
            result_start_time, field_name="result_start_time"
        )
        timestep = timedelta(seconds=result_timestep_seconds)
        node_rows: list[dict[str, Any]] = []
        for node_result in node_results:
            for period_index, values in enumerate(
                node_result.get("result", [])[:num_periods]
            ):
                node_rows.append(
                    {
                        "time": start_time + timestep * period_index,
                        "node_id": node_result["node"],
                        "actual_demand": values.get("demand"),
                        "total_head": values.get("head"),
                        "pressure": values.get("pressure"),
                        "quality": values.get("quality"),
                    }
                )
        link_rows: list[dict[str, Any]] = []
        for link_result in link_results:
            for period_index, values in enumerate(
                link_result.get("result", [])[:num_periods]
            ):
                link_rows.append(
                    {
                        "time": start_time + timestep * period_index,
                        "link_id": link_result["link"],
                        **{field: values.get(field) for field in AnalysisResultsRepository.LINK_FIELDS},
                    }
                )
        return node_rows, link_rows

    @staticmethod
    async def store_results(
        conn: AsyncConnection,
        run_id: UUID,
        node_rows: list[dict[str, Any]],
        link_rows: list[dict[str, Any]],
    ) -> None:
        async with conn.transaction(), conn.cursor() as cur:
            await AnalysisResultsRepository._lock_run(cur, run_id)
            await AnalysisResultsRepository._assert_run_is_empty(cur, run_id)
            if node_rows:
                async with cur.copy(
                    "COPY analysis.node_results "
                    "(time, run_id, node_id, actual_demand, total_head, pressure, quality) "
                    "FROM STDIN"
                ) as copy:
                    for row in node_rows:
                        await copy.write_row(
                            (
                                row["time"],
                                run_id,
                                row["node_id"],
                                row.get("actual_demand"),
                                row.get("total_head"),
                                row.get("pressure"),
                                row.get("quality"),
                            )
                        )
            if link_rows:
                async with cur.copy(
                    "COPY analysis.link_results "
                    "(time, run_id, link_id, flow, friction, headloss, quality, "
                    "reaction, setting, status, velocity) FROM STDIN"
                ) as copy:
                    for row in link_rows:
                        await copy.write_row(
                            (
                                row["time"],
                                run_id,
                                row["link_id"],
                                row.get("flow"),
                                row.get("friction"),
                                row.get("headloss"),
                                row.get("quality"),
                                row.get("reaction"),
                                row.get("setting"),
                                row.get("status"),
                                row.get("velocity"),
                            )
                        )

    @staticmethod
    async def _assert_run_is_empty(cur, run_id: UUID) -> None:
        await cur.execute(
            """
            SELECT EXISTS (
                SELECT 1 FROM analysis.node_results WHERE run_id = %s
                UNION ALL
                SELECT 1 FROM analysis.link_results WHERE run_id = %s
            ) AS exists
            """,
            (run_id, run_id),
        )
        row = await cur.fetchone()
        if row and row["exists"]:
            raise ValueError(f"analysis results already exist for run {run_id}")

    @staticmethod
    async def _lock_run(cur, run_id: UUID) -> None:
        await cur.execute(
            "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 0))",
            (run_id,),
        )

    @staticmethod
    async def get_node_series(
        conn: AsyncConnection,
        run_id: UUID,
        node_id: str,
        start_time: datetime,
        end_time: datetime,
        field: str,
    ) -> list[dict[str, Any]]:
        if field not in AnalysisResultsRepository.NODE_FIELDS:
            raise ValueError(f"invalid node result field: {field}")
        query = sql.SQL(
            "SELECT time, {} AS value FROM analysis.node_results "
            "WHERE run_id = %s AND node_id = %s AND time BETWEEN %s AND %s "
            "ORDER BY time"
        ).format(sql.Identifier(field))
        async with conn.cursor() as cur:
            await cur.execute(query, (run_id, node_id, start_time, end_time))
            return await cur.fetchall()

    @staticmethod
    async def get_link_series(
        conn: AsyncConnection,
        run_id: UUID,
        link_id: str,
        start_time: datetime,
        end_time: datetime,
        field: str,
    ) -> list[dict[str, Any]]:
        if field not in AnalysisResultsRepository.LINK_FIELDS:
            raise ValueError(f"invalid link result field: {field}")
        query = sql.SQL(
            "SELECT time, {} AS value FROM analysis.link_results "
            "WHERE run_id = %s AND link_id = %s AND time BETWEEN %s AND %s "
            "ORDER BY time"
        ).format(sql.Identifier(field))
        async with conn.cursor() as cur:
            await cur.execute(query, (run_id, link_id, start_time, end_time))
            return await cur.fetchall()

    @staticmethod
    async def get_series_by_ids(
        conn: AsyncConnection,
        run_id: UUID,
        element_type: str,
        element_ids: list[str],
        start_time: datetime,
        end_time: datetime,
        field: str,
    ) -> dict[str, list[dict[str, Any]]]:
        if element_type == "node":
            table_name, id_column, valid_fields = (
                "node_results", "node_id", AnalysisResultsRepository.NODE_FIELDS
            )
        elif element_type == "link":
            table_name, id_column, valid_fields = (
                "link_results", "link_id", AnalysisResultsRepository.LINK_FIELDS
            )
        else:
            raise ValueError(f"invalid analysis element type: {element_type}")
        if field not in valid_fields:
            raise ValueError(f"invalid {element_type} result field: {field}")

        result: dict[str, list[dict[str, Any]]] = {
            element_id: [] for element_id in element_ids
        }
        if not element_ids:
            return result
        query = sql.SQL(
            "SELECT {} AS element_id, time, {} AS value FROM analysis.{} "
            "WHERE run_id = %s AND {} = ANY(%s) AND time BETWEEN %s AND %s "
            "ORDER BY {}, time"
        ).format(
            sql.Identifier(id_column),
            sql.Identifier(field),
            sql.Identifier(table_name),
            sql.Identifier(id_column),
            sql.Identifier(id_column),
        )
        async with conn.cursor() as cur:
            await cur.execute(query, (run_id, element_ids, start_time, end_time))
            for row in await cur.fetchall():
                result.setdefault(str(row["element_id"]), []).append(
                    {"time": row["time"], "value": row["value"]}
                )
        return result

    @staticmethod
    async def get_values_at_time(
        conn: AsyncConnection,
        run_id: UUID,
        element_type: str,
        result_time: datetime,
        field: str,
    ) -> dict[str, Any]:
        if element_type == "node":
            table, id_column, fields = (
                "node_results",
                "node_id",
                AnalysisResultsRepository.NODE_FIELDS,
            )
        elif element_type == "link":
            table, id_column, fields = (
                "link_results",
                "link_id",
                AnalysisResultsRepository.LINK_FIELDS,
            )
        else:
            raise ValueError("element_type must be node or link")
        if field not in fields:
            raise ValueError(f"invalid {element_type} result field: {field}")
        query = sql.SQL(
            "SELECT {id_column}, {field} AS value FROM analysis.{table} "
            "WHERE run_id = %s AND time = %s ORDER BY {id_column}"
        ).format(
            id_column=sql.Identifier(id_column),
            field=sql.Identifier(field),
            table=sql.Identifier(table),
        )
        async with conn.cursor() as cur:
            await cur.execute(query, (run_id, result_time))
            return {row[id_column]: row["value"] for row in await cur.fetchall()}

    @staticmethod
    def store_results_sync(
        conn: Connection,
        run_id: UUID,
        node_rows: list[dict[str, Any]],
        link_rows: list[dict[str, Any]],
    ) -> None:
        with conn.transaction(), conn.cursor() as cur:
            cur.execute(
                "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 0))",
                (run_id,),
            )
            cur.execute(
                """
                SELECT EXISTS (
                    SELECT 1 FROM analysis.node_results WHERE run_id = %s
                    UNION ALL
                    SELECT 1 FROM analysis.link_results WHERE run_id = %s
                ) AS exists
                """,
                (run_id, run_id),
            )
            row = cur.fetchone()
            if row and row["exists"]:
                raise ValueError(f"analysis results already exist for run {run_id}")
            if node_rows:
                with cur.copy(
                    "COPY analysis.node_results "
                    "(time, run_id, node_id, actual_demand, total_head, pressure, quality) "
                    "FROM STDIN"
                ) as copy:
                    for item in node_rows:
                        copy.write_row(
                            (
                                item["time"], run_id, item["node_id"],
                                item.get("actual_demand"), item.get("total_head"),
                                item.get("pressure"), item.get("quality"),
                            )
                        )
            if link_rows:
                with cur.copy(
                    "COPY analysis.link_results "
                    "(time, run_id, link_id, flow, friction, headloss, quality, "
                    "reaction, setting, status, velocity) FROM STDIN"
                ) as copy:
                    for item in link_rows:
                        copy.write_row(
                            (
                                item["time"], run_id, item["link_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"),
                            )
                        )