"""Streaming über PostgreSQL (→ Konzept 4.4).

Der Worker schreibt Ereignisse mit fortlaufender Nummer je Run in `run_events`
und weckt per NOTIFY die Web-Prozesse. Die lesen per `seq > letzte` nach.
NOTIFY ist nur ein Signal; die Daten kommen immer aus der Tabelle.
"""

import asyncio
import logging
import time
import uuid
from typing import Any

import psycopg
from sqlalchemy import func, select, text
from sqlalchemy.orm import Session

from app.core.database import Database
from app.tasks.models import RunEvent

CHANNEL = "run_events"
log = logging.getLogger(__name__)


def append(session: Session, run_id: uuid.UUID, type_: str, data: dict[str, Any]) -> RunEvent:
    """Ereignis anhängen; NOTIFY wird mit dem Commit derselben Transaktion zugestellt."""
    last = session.scalar(select(func.max(RunEvent.seq)).where(RunEvent.run_id == run_id)) or 0
    event = RunEvent(run_id=run_id, seq=last + 1, type=type_, data=data)
    session.add(event)
    session.flush()
    session.execute(text("SELECT pg_notify(:channel, :payload)"), {"channel": CHANNEL, "payload": str(run_id)})
    return event


def read(session: Session, run_id: uuid.UUID, after: int, limit: int = 500) -> list[RunEvent]:
    return list(session.scalars(select(RunEvent).where(RunEvent.run_id == run_id, RunEvent.seq > after)
                                .order_by(RunEvent.seq).limit(limit)))


class EventWriter:
    """Für den Worker: Tokens sammeln und gebündelt schreiben, statt ein Ereignis je Token."""

    def __init__(self, db: Database, run_id: uuid.UUID, flush_ms: int):
        self.db, self.run_id, self.flush_seconds = db, run_id, flush_ms / 1000
        self.buffer: list[str] = []
        self.last_flush = time.monotonic()

    def token(self, piece: str) -> None:
        self.buffer.append(piece)
        if time.monotonic() - self.last_flush >= self.flush_seconds:
            self.flush()

    def flush(self) -> None:
        if self.buffer:
            pieces, self.buffer = "".join(self.buffer), []
            with self.db.session() as session:
                append(session, self.run_id, "token", {"text": pieces})
        self.last_flush = time.monotonic()

    def emit(self, type_: str, data: dict[str, Any], session: Session | None = None) -> None:
        """Offene Tokens zuerst, damit die Reihenfolge stimmt."""
        self.flush()
        if session is not None:
            append(session, self.run_id, type_, data)
            return
        with self.db.session() as own:
            append(own, self.run_id, type_, data)


class EventHub:
    """Im Web-Prozess: eine LISTEN-Verbindung je Prozess, die wartende SSE-Streams weckt.

    Fällt die Verbindung aus, verbindet sich der Hub neu; die Streams fragen in der
    Zwischenzeit von selbst in kurzen Abständen nach (Fallback-Polling).
    """

    def __init__(self, url: str):
        self.url = url
        self.waiters: dict[str, set[asyncio.Event]] = {}
        self.connected = False

    def subscribe(self, run_id: uuid.UUID) -> asyncio.Event:
        event = asyncio.Event()
        self.waiters.setdefault(str(run_id), set()).add(event)
        return event

    def unsubscribe(self, run_id: uuid.UUID, event: asyncio.Event) -> None:
        waiting = self.waiters.get(str(run_id))
        if waiting is not None:
            waiting.discard(event)
            if not waiting:
                self.waiters.pop(str(run_id), None)

    def wake(self, run_id: str) -> None:
        for event in self.waiters.get(run_id, ()):
            event.set()

    async def run(self) -> None:
        while True:
            try:
                async with await psycopg.AsyncConnection.connect(self.url, autocommit=True) as connection:
                    await connection.execute(f"LISTEN {CHANNEL}")
                    self.connected = True
                    async for notify in connection.notifies():
                        self.wake(notify.payload)
            except asyncio.CancelledError:
                raise
            except Exception:
                log.warning("LISTEN-Verbindung unterbrochen, neuer Versuch in 2 s", exc_info=True)
            self.connected = False
            await asyncio.sleep(2)