"""SQLAlchemy-backed stream-truncation event repository."""

from __future__ import annotations

import logging
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy import delete, desc, func, select
from sqlalchemy.ext.asyncio import async_sessionmaker

from deerflow.persistence.stream_truncation.base import (
    RETENTION_DAYS,
    StreamTruncationEvent,
    StreamTruncationQuery,
    StreamTruncationStore,
)
from deerflow.persistence.stream_truncation.model import StreamTruncationEventRow

logger = logging.getLogger(__name__)


class SqlStreamTruncationStore(StreamTruncationStore):
    """Persist truncation events with a hard one-month retention cap."""

    def __init__(self, session_factory: async_sessionmaker) -> None:
        self._session_factory = session_factory

    async def _ensure_table(self, session) -> None:
        # Table creation follows the repo's create_all-on-first-use convention
        # (see e.g. task_buttons); safe to call repeatedly.
        try:
            await session.run_sync(
                lambda sync_session: StreamTruncationEventRow.__table__.create(
                    sync_session.bind, checkfirst=True
                )
            )
        except Exception:
            logger.debug("stream_truncation table create skipped", exc_info=True)

    async def record(self, event: StreamTruncationEvent) -> None:
        try:
            async with self._session_factory() as session:
                await self._ensure_table(session)
                session.add(
                    StreamTruncationEventRow(
                        id=event.id,
                        created_at=event.created_at,
                        model_name=event.model_name,
                        provider_completion_id=event.provider_completion_id,
                        run_id=event.run_id,
                        event_type=event.event_type,
                        strict_mode=event.strict_mode,
                        finish_reason=event.finish_reason,
                        status=event.status,
                        error_type=event.error_type,
                        chunk_count=event.chunk_count,
                        duration_ms=event.duration_ms,
                        first_chunk_ms=event.first_chunk_ms,
                        detail=event.detail,
                    )
                )
                await session.commit()
                # Same session/transaction scope: enforce the retention cap on
                # every insert. Truncations are rare, so this stays cheap and
                # keeps the one-month promise without any scheduler.
                await self._enforce_retention(session)
        except Exception:
            logger.warning(
                "Failed to persist stream truncation event model=%s type=%s",
                event.model_name,
                event.event_type,
                exc_info=True,
            )

    async def _enforce_retention(self, session) -> None:
        """Delete rows beyond the retention cap. Best-effort, never raises."""
        try:
            cutoff = datetime.now(UTC) - timedelta(days=RETENTION_DAYS)
            await session.execute(
                delete(StreamTruncationEventRow).where(
                    StreamTruncationEventRow.created_at < cutoff
                )
            )
            await session.commit()
        except Exception:
            logger.debug("stream truncation retention sweep failed", exc_info=True)

    def _apply_filters(self, query: StreamTruncationQuery):
        stmt = select(StreamTruncationEventRow)
        if query.model_name:
            stmt = stmt.where(StreamTruncationEventRow.model_name == query.model_name)
        if query.event_type:
            stmt = stmt.where(StreamTruncationEventRow.event_type == query.event_type)
        if query.since is not None:
            stmt = stmt.where(StreamTruncationEventRow.created_at >= query.since)
        if query.until is not None:
            stmt = stmt.where(StreamTruncationEventRow.created_at <= query.until)
        return stmt

    async def list(self, query: StreamTruncationQuery) -> list[dict[str, Any]]:
        try:
            async with self._session_factory() as session:
                await self._ensure_table(session)
                stmt = self._apply_filters(query).order_by(
                    desc(StreamTruncationEventRow.created_at)
                ).limit(query.limit).offset(query.offset)
                rows = (await session.execute(stmt)).scalars().all()
                return [row.to_dict() for row in rows]
        except Exception:
            logger.warning("Failed to list stream truncation events", exc_info=True)
            return []

    async def count(self, query: StreamTruncationQuery) -> int:
        try:
            async with self._session_factory() as session:
                await self._ensure_table(session)
                stmt = self._apply_filters(query).with_only_columns(
                    func.count(StreamTruncationEventRow.id)
                )
                return int((await session.execute(stmt)).scalar() or 0)
        except Exception:
            logger.warning("Failed to count stream truncation events", exc_info=True)
            return 0

    async def model_breakdown(self, query: StreamTruncationQuery) -> list[dict[str, Any]]:
        try:
            async with self._session_factory() as session:
                await self._ensure_table(session)
                stmt = self._apply_filters(query).with_only_columns(
                    StreamTruncationEventRow.model_name,
                    StreamTruncationEventRow.event_type,
                    func.count(StreamTruncationEventRow.id).label("total"),
                    func.max(StreamTruncationEventRow.created_at).label("last_at"),
                ).group_by(StreamTruncationEventRow.model_name, StreamTruncationEventRow.event_type)
                rows = (await session.execute(stmt)).all()
                merged: dict[str, dict[str, Any]] = {}
                for row in rows:
                    name = row.model_name or "(unknown)"
                    bucket = merged.setdefault(
                        name, {"model_name": name, "total": 0, "last_at": None}
                    )
                    bucket[row.event_type] = int(row.total or 0)
                    bucket["total"] += int(row.total or 0)
                    row_last = row.last_at
                    if row_last is not None and (
                        bucket["last_at"] is None or row_last > bucket["last_at"]
                    ):
                        bucket["last_at"] = row_last
                result = sorted(
                    merged.values(), key=lambda item: item["last_at"] or datetime.min.replace(tzinfo=UTC), reverse=True
                )
                for item in result:
                    if isinstance(item["last_at"], datetime):
                        item["last_at"] = item["last_at"].isoformat()
                return result
        except Exception:
            logger.warning("Failed to aggregate stream truncation events", exc_info=True)
            return []

    async def purge_before(self, cutoff: datetime) -> int:
        try:
            async with self._session_factory() as session:
                await self._ensure_table(session)
                result = await session.execute(
                    delete(StreamTruncationEventRow).where(
                        StreamTruncationEventRow.created_at < cutoff
                    )
                )
                await session.commit()
                rowcount = result.rowcount
                return 0 if rowcount is None else int(rowcount)
        except Exception:
            logger.warning("Failed to purge stream truncation events", exc_info=True)
            return 0