"""Short-transaction storage for replay control and resumable batches."""
from contextlib import asynccontextmanager
from datetime import datetime, timezone

from fastapi import HTTPException
from psycopg.types.json import Jsonb


class ScadaReplayRepository:
    def __init__(self, connection, project_id):
        self.connection, self.project_id = connection, str(project_id)
        self.lock_key = f"scada-replay:{self.project_id}"

    async def available(self) -> bool:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT to_regclass('peripheral_sample.scada_replays') IS NOT NULL "
                "AND to_regclass('peripheral_sample.scada_replay_batches') IS NOT NULL AS ready"
            )
            return (await cur.fetchone())["ready"]

    @asynccontextmanager
    async def control(self):
        async with self.connection.transaction():
            await self.connection.execute("SELECT pg_advisory_xact_lock(hashtextextended(%s,0))", (self.lock_key,))
            yield

    @asynccontextmanager
    async def processing(self):
        # Session lock outlives short transactions, never an open SQL transaction
        # across Timescale IO. A dead connection automatically releases ownership.
        async with self.connection.transaction():
            cur = await self.connection.execute("SELECT pg_try_advisory_lock(hashtextextended(%s,0)) AS owned", (self.lock_key,))
            owned = (await cur.fetchone())["owned"]
        try:
            yield owned
        finally:
            if owned:
                async with self.connection.transaction():
                    await self.connection.execute("SELECT pg_advisory_unlock(hashtextextended(%s,0))", (self.lock_key,))

    async def get(self, replay_id) -> dict:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT document FROM peripheral_sample.scada_replays WHERE project_id=%s AND replay_id=%s",
                (self.project_id, str(replay_id)),
            )
            row = await cur.fetchone()
        if not row:
            raise HTTPException(404, "当前项目不存在此 SCADA 回放")
        return row["document"]

    async def current(self) -> dict | None:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT document FROM peripheral_sample.scada_replays WHERE project_id=%s"
                " ORDER BY created_at DESC LIMIT 1", (self.project_id,),
            )
            row = await cur.fetchone()
        return row["document"] if row else None

    async def by_request(self, request_id) -> dict | None:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT document FROM peripheral_sample.scada_replays WHERE project_id=%s AND request_id=%s",
                (self.project_id, str(request_id)),
            )
            row = await cur.fetchone()
        return row["document"] if row else None

    async def active_other(self, replay_id) -> bool:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT 1 FROM peripheral_sample.scada_replays WHERE project_id=%s AND replay_id<>%s "
                "AND status IN ('ready','paused') LIMIT 1",
                (self.project_id, str(replay_id)),
            )
            return await cur.fetchone() is not None

    async def insert(self, document: dict):
        await self.connection.execute(
            "INSERT INTO peripheral_sample.scada_replays(project_id,replay_id,request_id,status,document) "
            "VALUES(%s,%s,%s,%s,%s)",
            (self.project_id, document["replay_id"], document["config"]["request_id"], document["status"], Jsonb(document)),
        )

    async def save(self, document: dict):
        document["updated_at"] = datetime.now(timezone.utc).isoformat()
        await self.connection.execute(
            "UPDATE peripheral_sample.scada_replays SET status=%s,document=%s,updated_at=now() "
            "WHERE project_id=%s AND replay_id=%s",
            (document["status"], Jsonb(document), self.project_id, document["replay_id"]),
        )

    async def pending_batch(self, replay_id) -> dict | None:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT document FROM peripheral_sample.scada_replay_batches "
                "WHERE project_id=%s AND replay_id=%s AND stage<>'detected' ORDER BY sequence LIMIT 1",
                (self.project_id, str(replay_id)),
            )
            row = await cur.fetchone()
        return row["document"] if row else None

    async def save_batch(self, replay_id, batch: dict):
        await self.connection.execute(
            "INSERT INTO peripheral_sample.scada_replay_batches(project_id,replay_id,sequence,stage,document) "
            "VALUES(%s,%s,%s,%s,%s) ON CONFLICT(project_id,replay_id,sequence) "
            "DO UPDATE SET stage=excluded.stage,document=excluded.document",
            (self.project_id, str(replay_id), batch["sequence"], batch["stage"], Jsonb(batch)),
        )

    async def events(self, replay_id, offset=0, limit=100) -> list[dict]:
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT e.value AS event FROM peripheral_sample.scada_replay_batches b "
                "CROSS JOIN LATERAL jsonb_array_elements(COALESCE(b.document->'events','[]'::jsonb)) e "
                "WHERE b.project_id=%s AND b.replay_id=%s ORDER BY b.sequence DESC LIMIT %s OFFSET %s",
                (self.project_id, str(replay_id), limit, offset),
            )
            return [r["event"] for r in await cur.fetchall()]

    async def monitoring_outputs(
        self,
        replay_id,
        device_ids: list[str],
        start: datetime,
        end: datetime,
        before_sequence: int,
    ) -> list[dict]:
        """Read only committed replay outputs available before the current batch."""
        if not device_ids:
            return []
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT r.value AS reading FROM peripheral_sample.scada_replay_batches b "
                "CROSS JOIN LATERAL jsonb_array_elements(COALESCE(b.document->'outputs','[]'::jsonb)) r "
                "WHERE b.project_id=%s AND b.replay_id=%s AND b.stage='detected' AND b.sequence<%s "
                "AND r.value->>'device_id'=ANY(%s) "
                "AND (r.value->>'time')::timestamptz>=%s AND (r.value->>'time')::timestamptz<=%s "
                "ORDER BY (r.value->>'time')::timestamptz,r.value->>'device_id'",
                (self.project_id, str(replay_id), before_sequence, device_ids, start, end),
            )
            return [row["reading"] for row in await cur.fetchall()]

    async def detected_outputs(
        self,
        replay_id,
        device_ids: list[str],
        start: datetime,
        end: datetime,
        before_sequence: int,
        *,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict]:
        """Read immutable detector outputs available to a bounded replay investigation."""
        statement = (
            "SELECT r.value AS reading FROM peripheral_sample.scada_replay_batches b "
            "CROSS JOIN LATERAL jsonb_array_elements(COALESCE(b.document->'outputs','[]'::jsonb)) r "
            "WHERE b.project_id=%s AND b.replay_id=%s AND b.sequence<=%s AND b.stage='detected' "
            "AND r.value->>'device_id'=ANY(%s) "
            "AND (r.value->>'time')::timestamptz>=%s AND (r.value->>'time')::timestamptz<=%s "
            "ORDER BY (r.value->>'time')::timestamptz,r.value->>'device_id'"
        )
        params: tuple = (
            self.project_id, str(replay_id), before_sequence, device_ids, start, end,
        )
        if limit is not None:
            statement += " LIMIT %s OFFSET %s"
            params += (limit, offset)
        async with self.connection.transaction():
            cur = await self.connection.execute(statement, params)
            return [row["reading"] for row in await cur.fetchall()]

    async def latest_outputs(self, replay_id, through: datetime, limit: int = 500) -> list[dict]:
        """Return the latest committed cleaned output for each replay device."""
        async with self.connection.transaction():
            cur = await self.connection.execute(
                "SELECT DISTINCT ON(r.value->>'device_id') r.value AS reading "
                "FROM peripheral_sample.scada_replay_batches b "
                "CROSS JOIN LATERAL jsonb_array_elements(COALESCE(b.document->'outputs','[]'::jsonb)) r "
                "WHERE b.project_id=%s AND b.replay_id=%s AND b.stage IN ('cleaned','detected') "
                "AND (r.value->>'time')::timestamptz<=%s "
                "ORDER BY r.value->>'device_id',(r.value->>'time')::timestamptz DESC LIMIT %s",
                (self.project_id, str(replay_id), through, limit),
            )
            return [row["reading"] for row in await cur.fetchall()]