from typing import List, Any
from datetime import datetime, timedelta
from collections import defaultdict
from psycopg import AsyncConnection, Connection, sql
from psycopg.rows import dict_row


class ScadaRepository:
    VALUE_FIELDS = frozenset({"monitored_value", "cleaned_value"})
    RESULT_COLUMNS = "time, device_id, monitored_value, cleaned_value"

    @staticmethod
    async def begin_replay_baseline_snapshot(conn: AsyncConnection) -> None:
        """Set the caller's source transaction to a repeatable read-only snapshot."""
        await conn.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")

    @staticmethod
    async def get_warmup_values(
        conn: AsyncConnection,
        device_ids: list[str],
        before: datetime,
        per_device_limit: int,
    ) -> list[dict]:
        """Return the latest raw values before a replay starts for detector warmup."""
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT d.device_id,h.monitored_value FROM unnest(%s::text[]) d(device_id) "
                "CROSS JOIN LATERAL (SELECT monitored_value,time FROM scada.measurements "
                "WHERE device_id=d.device_id AND time<%s ORDER BY time DESC LIMIT %s) h "
                "ORDER BY d.device_id,h.time",
                (device_ids, before, per_device_limit),
            )
            return [dict(row) for row in await cur.fetchall()]

    @staticmethod
    async def get_raw_page(
        conn: AsyncConnection, device_ids: list[str], start: datetime,
        end: datetime, through: datetime, limit: int,
        after: tuple[datetime, str] | None = None, *, frame: bool = False,
    ) -> list[dict]:
        """Read raw only, with a stable composite key and a bounded frame/page."""
        predicate = "device_id=ANY(%s) AND time>=%s AND time<%s AND time<=%s"
        params: tuple = (device_ids, start, end, through)
        if after is not None:
            predicate += " AND (time,device_id)>(%s,%s)"
            params += after
        if frame:
            statement = (
                "WITH next_time AS (SELECT min(time) AS time FROM scada.measurements WHERE "
                + predicate + ") SELECT m.time,m.device_id,m.monitored_value,"
                "md5(COALESCE(m.monitored_value::text,'NULL')) AS source_hash "
                "FROM scada.measurements m WHERE m.device_id=ANY(%s) "
                "AND m.time=(SELECT time FROM next_time) ORDER BY m.time,m.device_id LIMIT %s"
            )
            params += (device_ids, limit)
        else:
            statement = (
                "SELECT time,device_id,monitored_value,"
                "md5(COALESCE(monitored_value::text,'NULL')) AS source_hash "
                "FROM scada.measurements WHERE " + predicate
                + " ORDER BY time,device_id LIMIT %s"
            )
            params += (limit,)
        async with conn.cursor() as cur:
            await cur.execute(statement, params)
            return [dict(row) for row in await cur.fetchall()]

    @staticmethod
    async def apply_replay_cleaning(conn: AsyncConnection, rows: list[dict]) -> None:
        """Guard every raw version; update corrected values without inserting raw."""
        if not rows:
            return
        async with conn.transaction():
            async with conn.cursor() as cur:
                await cur.execute(
                    "SELECT m.time,m.device_id,md5(COALESCE(m.monitored_value::text,'NULL')) AS source_hash "
                    "FROM scada.measurements m JOIN unnest(%s::timestamptz[],%s::text[]) "
                    "b(time,device_id) USING(time,device_id) FOR UPDATE OF m",
                    ([r["time"] for r in rows], [r["device_id"] for r in rows]),
                )
                existing = {(r["time"].isoformat(), r["device_id"]): r["source_hash"]
                            for r in await cur.fetchall()}
                if len(existing) != len(rows) or any(
                    existing.get((datetime.fromisoformat(r["time"]).isoformat(), r["device_id"]))
                    != r["source_hash"] for r in rows
                ):
                    raise ValueError("SOURCE_CHANGED: SCADA 原始值在处理期间发生变化")
                await cur.execute(
                    "UPDATE scada.measurements m SET cleaned_value=b.value "
                    "FROM unnest(%s::timestamptz[],%s::text[],%s::double precision[]) "
                    "b(time,device_id,value) WHERE m.time=b.time AND m.device_id=b.device_id "
                    "AND m.cleaned_value IS DISTINCT FROM b.value",
                    ([r["time"] for r in rows], [r["device_id"] for r in rows],
                     [r["cleaned_value"] for r in rows]),
                )

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

        async with conn.cursor() as cur:
            async with cur.copy(
                "COPY scada.measurements (time, device_id, monitored_value, cleaned_value) FROM STDIN"
            ) as copy:
                for item in data:
                    await copy.write_row(
                        (
                            item["time"],
                            item["device_id"],
                            item.get("monitored_value"),
                            item.get("cleaned_value"),
                        )
                    )

    @staticmethod
    async def get_scada_by_ids_time_range(
        conn: AsyncConnection,
        device_ids: List[str],
        start_time: datetime,
        end_time: datetime,
        *, row_limit: int | None = None,
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT {ScadaRepository.RESULT_COLUMNS} FROM scada.measurements "
                "WHERE device_id = ANY(%s) "
                "AND time >= %s AND time <= %s ORDER BY device_id, time"
                + (" LIMIT %s" if row_limit is not None else ""),
                (device_ids, start_time, end_time) + ((row_limit,) if row_limit is not None else ()),
            )
            return await cur.fetchall()

    @staticmethod
    async def get_recovery_scada_readings(
        conn: AsyncConnection, device_ids: list[str], incident_at: datetime,
        returned_at: datetime, through: datetime, *, baseline_limit: int = 100_001,
    ) -> list[dict]:
        """Read the causal reference and exactly the first stored frame after T."""
        if not device_ids or through <= returned_at:
            return []
        async with conn.cursor() as cur:
            await cur.execute(
                "WITH next_frame AS (SELECT min(time) AS time FROM scada.measurements "
                "WHERE device_id=ANY(%s) AND time>%s AND time<=%s), "
                "reference AS (SELECT time,device_id,monitored_value,cleaned_value "
                "FROM scada.measurements WHERE device_id=ANY(%s) AND time>=%s AND time<%s "
                "ORDER BY time,device_id LIMIT %s) "
                "SELECT time,device_id,monitored_value,cleaned_value FROM reference UNION ALL "
                "SELECT m.time,m.device_id,m.monitored_value,m.cleaned_value FROM scada.measurements m "
                "JOIN next_frame f ON m.time=f.time WHERE m.device_id=ANY(%s) ORDER BY time,device_id",
                (device_ids, returned_at, through, device_ids, incident_at - timedelta(days=15),
                 incident_at, baseline_limit, device_ids),
            )
            return [dict(row) for row in await cur.fetchall()]

    @staticmethod
    async def get_latest_scada_readings(
        conn: AsyncConnection,
        device_ids: List[str],
    ) -> dict[str, dict[str, Any]]:
        """Return at most one latest row for each requested canonical device ID."""
        if not device_ids:
            return {}

        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT DISTINCT ON (device_id) {ScadaRepository.RESULT_COLUMNS} "
                "FROM scada.measurements WHERE device_id = ANY(%s) "
                "ORDER BY device_id, time DESC",
                (device_ids,),
            )
            rows = await cur.fetchall()
        return {str(row["device_id"]): dict(row) for row in rows}

    @staticmethod
    def get_scada_by_ids_time_range_sync(
        conn: Connection,
        device_ids: List[str],
        start_time: datetime,
        end_time: datetime,
    ) -> List[dict]:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(
                f"SELECT {ScadaRepository.RESULT_COLUMNS} FROM scada.measurements "
                "WHERE device_id = ANY(%s) "
                "AND time >= %s AND time <= %s ORDER BY device_id, time",
                (device_ids, start_time, end_time),
            )
            return cur.fetchall()

    @staticmethod
    def get_latest_scada_time_sync(
        conn: Connection,
        device_ids: List[str],
        before_time: datetime | None = None,
    ) -> datetime | None:
        with conn.cursor(row_factory=dict_row) as cur:
            if before_time is None:
                cur.execute(
                    "SELECT max(time) AS time FROM scada.measurements WHERE device_id = ANY(%s)",
                    (device_ids,),
                )
            else:
                cur.execute(
                    "SELECT max(time) AS time FROM scada.measurements "
                    "WHERE device_id = ANY(%s) AND time <= %s",
                    (device_ids, before_time),
                )
            row = cur.fetchone()
            return row["time"] if row else None

    @staticmethod
    async def get_scada_field_by_id_time_range(
        conn: AsyncConnection,
        device_ids: List[str],
        start_time: datetime,
        end_time: datetime,
        field: str,
    ) -> dict:
        if field not in ScadaRepository.VALUE_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
            "SELECT device_id, time, {} FROM scada.measurements "
            "WHERE time >= %s AND time <= %s AND device_id = ANY(%s) "
            "ORDER BY device_id, time"
        ).format(sql.Identifier(field))

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

    @staticmethod
    async def update_scada_field(
        conn: AsyncConnection, time: datetime, device_id: str, field: str, value: Any
    ):
        if field not in ScadaRepository.VALUE_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        update_query = sql.SQL(
            "UPDATE scada.measurements SET {} = %s WHERE time = %s AND device_id = %s"
        ).format(sql.Identifier(field))
        insert_query = sql.SQL(
            "INSERT INTO scada.measurements (time, device_id, {}) VALUES (%s, %s, %s)"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(update_query, (value, time, device_id))
            if cur.rowcount == 0:
                await cur.execute(insert_query, (time, device_id, value))

    @staticmethod
    async def update_scada_field_batch(
        conn: AsyncConnection,
        rows: list[tuple[datetime, str, float | None]],
        field: str,
    ) -> int:
        """Update existing SCADA samples in one set-based statement."""
        if field not in ScadaRepository.VALUE_FIELDS:
            raise ValueError(f"Invalid field: {field}")
        if not rows:
            return 0

        query = sql.SQL(
            "UPDATE scada.measurements AS measurement SET {} = batch.value "
            "FROM unnest(%s::timestamptz[], %s::text[], %s::double precision[]) "
            "AS batch(time, device_id, value) "
            "WHERE measurement.time = batch.time "
            "AND measurement.device_id = batch.device_id"
        ).format(sql.Identifier(field))
        times, device_ids, values = zip(*rows)
        async with conn.cursor() as cur:
            await cur.execute(query, (list(times), list(device_ids), list(values)))
            return cur.rowcount

    @staticmethod
    async def delete_scada_by_id_time_range(
        conn: AsyncConnection, device_id: str, start_time: datetime, end_time: datetime
    ):
        async with conn.cursor() as cur:
            await cur.execute(
                "DELETE FROM scada.measurements WHERE device_id = %s AND time >= %s AND time <= %s",
                (device_id, start_time, end_time),
            )