Nexus/server/infrastructure/database/watch_repo.py

"""Repository for watch pins, sessions, and probe records."""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Any

from sqlalchemy import delete, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from server.domain.models import Server, WatchPin, WatchPinSession, WatchProbeRecord
from server.utils.watch_metrics import normalize_watch_ttl_hours


def _utc_naive_now() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


def _pin_still_occupied(pin: WatchPin, now: datetime) -> bool:
    """Pin occupies a slot until manual remove or TTL expiry while monitoring is on."""
    if not pin.monitoring_enabled:
        return True
    return pin.expires_at > now


class WatchRepositoryImpl:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def list_active_pins_for_admin(self, admin_id: int) -> list[dict[str, Any]]:
        now = _utc_naive_now()
        q = (
            select(WatchPin, WatchPinSession, Server.name)
            .join(WatchPinSession, WatchPin.session_id == WatchPinSession.id)
            .join(Server, WatchPin.server_id == Server.id)
            .where(
                WatchPin.admin_id == admin_id,
                or_(WatchPin.monitoring_enabled.is_(False), WatchPin.expires_at > now),
            )
            .order_by(WatchPin.slot_index)
        )
        rows = (await self.session.execute(q)).all()
        out: list[dict[str, Any]] = []
        for pin, sess, server_name in rows:
            if not _pin_still_occupied(pin, now):
                continue
            out.append(
                {
                    "pin_id": pin.id,
                    "slot_index": pin.slot_index,
                    "server_id": pin.server_id,
                    "server_name": server_name,
                    "session_id": sess.id,
                    "pinned_at": pin.pinned_at,
                    "expires_at": pin.expires_at,
                    "monitoring_enabled": bool(pin.monitoring_enabled),
                    "ttl_hours": int(pin.ttl_hours or 8),
                }
            )
        return out

    async def list_active_pins_all(self, *, monitoring_only: bool = False) -> list[dict[str, Any]]:
        now = _utc_naive_now()
        q = (
            select(WatchPin, WatchPinSession)
            .join(WatchPinSession, WatchPin.session_id == WatchPinSession.id)
        )
        rows = (await self.session.execute(q)).all()
        out: list[dict[str, Any]] = []
        for pin, sess in rows:
            if not _pin_still_occupied(pin, now):
                continue
            if monitoring_only and not pin.monitoring_enabled:
                continue
            out.append(
                {
                    "pin_id": pin.id,
                    "admin_id": pin.admin_id,
                    "slot_index": pin.slot_index,
                    "server_id": pin.server_id,
                    "session_id": sess.id,
                }
            )
        return out

    async def get_pin_by_slot(self, admin_id: int, slot_index: int) -> WatchPin | None:
        q = select(WatchPin).where(
            WatchPin.admin_id == admin_id,
            WatchPin.slot_index == slot_index,
        )
        pin = (await self.session.execute(q)).scalar_one_or_none()
        if pin and _pin_still_occupied(pin, _utc_naive_now()):
            return pin
        return None

    async def get_pin_by_server(self, admin_id: int, server_id: int) -> WatchPin | None:
        q = select(WatchPin).where(
            WatchPin.admin_id == admin_id,
            WatchPin.server_id == server_id,
        )
        pin = (await self.session.execute(q)).scalar_one_or_none()
        if pin and _pin_still_occupied(pin, _utc_naive_now()):
            return pin
        return None

    async def create_pin(
        self,
        *,
        admin_id: int,
        server_id: int,
        slot_index: int,
        ttl_hours: int = 8,
    ) -> dict[str, Any]:
        now = _utc_naive_now()
        hours = normalize_watch_ttl_hours(ttl_hours)
        expires = now + timedelta(hours=hours)
        sess = WatchPinSession(
            admin_id=admin_id,
            server_id=server_id,
            slot_index=slot_index,
            started_at=now,
        )
        self.session.add(sess)
        await self.session.flush()
        pin = WatchPin(
            admin_id=admin_id,
            session_id=sess.id,
            slot_index=slot_index,
            server_id=server_id,
            pinned_at=now,
            expires_at=expires,
            monitoring_enabled=True,
            ttl_hours=hours,
        )
        self.session.add(pin)
        await self.session.flush()
        return {
            "pin_id": pin.id,
            "session_id": sess.id,
            "slot_index": slot_index,
            "server_id": server_id,
            "pinned_at": now,
            "expires_at": expires,
            "ttl_hours": hours,
        }

    async def end_session(self, session_id: int, reason: str) -> None:
        q = select(WatchPinSession).where(WatchPinSession.id == session_id)
        sess = (await self.session.execute(q)).scalar_one_or_none()
        if sess and sess.ended_at is None:
            sess.ended_at = _utc_naive_now()
            sess.end_reason = reason

    async def delete_pin(self, pin_id: int) -> None:
        await self.session.execute(delete(WatchPin).where(WatchPin.id == pin_id))

    async def delete_pin_by_slot(self, admin_id: int, slot_index: int) -> WatchPin | None:
        pin = await self.get_pin_by_slot(admin_id, slot_index)
        if not pin:
            return None
        await self.end_session(pin.session_id, "manual")
        await self.delete_pin(pin.id)
        return pin

    async def replace_pin_server(
        self,
        *,
        admin_id: int,
        slot_index: int,
        server_id: int,
        ttl_hours: int = 8,
    ) -> dict[str, Any]:
        existing = await self.get_pin_by_slot(admin_id, slot_index)
        if existing:
            await self.end_session(existing.session_id, "replaced")
            await self.delete_pin(existing.id)
        other = await self.get_pin_by_server(admin_id, server_id)
        if other and other.slot_index != slot_index:
            await self.end_session(other.session_id, "replaced")
            await self.delete_pin(other.id)
        return await self.create_pin(
            admin_id=admin_id,
            server_id=server_id,
            slot_index=slot_index,
            ttl_hours=ttl_hours,
        )

    async def reset_pin_ttl(self, admin_id: int, server_id: int, ttl_hours: int | None = None) -> dict[str, Any]:
        """Re-pin same server: end old session and start fresh monitoring window."""
        pin = await self.get_pin_by_server(admin_id, server_id)
        if not pin:
            raise ValueError("pin not found")
        slot_index = pin.slot_index
        hours = normalize_watch_ttl_hours(ttl_hours if ttl_hours is not None else int(pin.ttl_hours or 8))
        await self.end_session(pin.session_id, "replaced")
        await self.delete_pin(pin.id)
        return await self.create_pin(
            admin_id=admin_id,
            server_id=server_id,
            slot_index=slot_index,
            ttl_hours=hours,
        )

    async def remove_pin_by_slot(self, admin_id: int, slot_index: int) -> WatchPin | None:
        return await self.delete_pin_by_slot(admin_id, slot_index)

    async def count_monitoring_pins_for_server(self, server_id: int) -> int:
        q = select(func.count()).select_from(WatchPin).where(
            WatchPin.server_id == server_id,
            WatchPin.monitoring_enabled.is_(True),
        )
        return int((await self.session.execute(q)).scalar_one() or 0)

    async def _ensure_open_session(self, pin: WatchPin) -> None:
        q = select(WatchPinSession).where(WatchPinSession.id == pin.session_id)
        sess = (await self.session.execute(q)).scalar_one_or_none()
        if sess and sess.ended_at is None:
            return
        now = _utc_naive_now()
        new_sess = WatchPinSession(
            admin_id=pin.admin_id,
            server_id=pin.server_id,
            slot_index=pin.slot_index,
            started_at=now,
        )
        self.session.add(new_sess)
        await self.session.flush()
        pin.session_id = new_sess.id

    async def set_monitoring_enabled(
        self,
        admin_id: int,
        slot_index: int,
        enabled: bool,
        *,
        ttl_hours: int | None = None,
    ) -> WatchPin | None:
        pin = await self.get_pin_by_slot(admin_id, slot_index)
        if not pin:
            return None
        if ttl_hours is not None:
            pin.ttl_hours = normalize_watch_ttl_hours(ttl_hours)
        if enabled:
            await self._ensure_open_session(pin)
            hours = normalize_watch_ttl_hours(int(pin.ttl_hours or 8))
            pin.ttl_hours = hours
            pin.monitoring_enabled = True
            pin.expires_at = _utc_naive_now() + timedelta(hours=hours)
        else:
            pin.monitoring_enabled = False
        await self.session.flush()
        return pin

    async def set_pin_ttl_hours(
        self,
        admin_id: int,
        slot_index: int,
        ttl_hours: int,
    ) -> WatchPin | None:
        pin = await self.get_pin_by_slot(admin_id, slot_index)
        if not pin:
            return None
        hours = normalize_watch_ttl_hours(ttl_hours)
        pin.ttl_hours = hours
        if pin.monitoring_enabled:
            pin.expires_at = _utc_naive_now() + timedelta(hours=hours)
        await self.session.flush()
        return pin

    async def expire_due_pins(self) -> list[int]:
        now = _utc_naive_now()
        q = select(WatchPin).where(
            WatchPin.expires_at <= now,
            WatchPin.monitoring_enabled.is_(True),
        )
        pins = (await self.session.execute(q)).scalars().all()
        server_ids: list[int] = []
        for pin in pins:
            await self.end_session(pin.session_id, "expired")
            pin.monitoring_enabled = False
            server_ids.append(pin.server_id)
        return server_ids

    async def insert_probe_record(self, **fields: Any) -> WatchProbeRecord:
        row = WatchProbeRecord(**fields)
        self.session.add(row)
        await self.session.flush()
        return row

    async def get_latest_probe_for_session(self, session_id: int) -> WatchProbeRecord | None:
        q = (
            select(WatchProbeRecord)
            .where(WatchProbeRecord.session_id == session_id)
            .order_by(WatchProbeRecord.recorded_at.desc())
            .limit(1)
        )
        return (await self.session.execute(q)).scalar_one_or_none()

    async def purge_records_before(self, cutoff: datetime) -> int:
        res = await self.session.execute(
            delete(WatchProbeRecord).where(WatchProbeRecord.recorded_at < cutoff)
        )
        return int(res.rowcount or 0)

    async def first_empty_slot(self, admin_id: int) -> int | None:
        active = await self.list_active_pins_for_admin(admin_id)
        used = {p["slot_index"] for p in active}
        for i in range(4):
            if i not in used:
                return i
        return None

    async def list_probe_records(
        self,
        admin_id: int,
        *,
        server_id: int | None = None,
        session_id: int | None = None,
        probe_status: str | None = None,
        since: datetime | None = None,
        page: int = 1,
        per_page: int = 50,
    ) -> tuple[list[dict[str, Any]], int]:
        filters = [WatchProbeRecord.admin_id == admin_id]
        if server_id is not None:
            filters.append(WatchProbeRecord.server_id == server_id)
        if session_id is not None:
            filters.append(WatchProbeRecord.session_id == session_id)
        if probe_status:
            filters.append(WatchProbeRecord.probe_status == probe_status)
        if since is not None:
            filters.append(WatchProbeRecord.recorded_at >= since)

        count_q = select(func.count()).select_from(WatchProbeRecord).where(*filters)
        total = int((await self.session.execute(count_q)).scalar_one() or 0)

        q = (
            select(WatchProbeRecord, Server.name)
            .join(Server, WatchProbeRecord.server_id == Server.id)
            .where(*filters)
            .order_by(WatchProbeRecord.recorded_at.desc())
            .offset(max(0, page - 1) * per_page)
            .limit(per_page)
        )
        rows = (await self.session.execute(q)).all()
        items: list[dict[str, Any]] = []
        for rec, server_name in rows:
            items.append(
                {
                    "id": rec.id,
                    "recorded_at": rec.recorded_at.isoformat() if rec.recorded_at else None,
                    "session_id": rec.session_id,
                    "server_id": rec.server_id,
                    "server_name": server_name,
                    "source": rec.source,
                    "probe_status": rec.probe_status,
                    "duration_ms": rec.duration_ms,
                    "cpu_pct": rec.cpu_pct,
                    "mem_pct": rec.mem_pct,
                    "disk_pct": rec.disk_pct,
                    "net_up_bps": rec.net_up_bps,
                    "net_down_bps": rec.net_down_bps,
                    "disk_read_bps": rec.disk_read_bps,
                    "disk_write_bps": rec.disk_write_bps,
                    "error": rec.error,
                }
            )
        return items, total

    async def list_metrics_series(
        self,
        admin_id: int,
        *,
        server_ids: list[int],
        since: datetime,
        ok_only: bool = True,
    ) -> list[dict[str, Any]]:
        if not server_ids:
            return []
        filters = [
            WatchProbeRecord.admin_id == admin_id,
            WatchProbeRecord.server_id.in_(server_ids),
            WatchProbeRecord.recorded_at >= since,
        ]
        if ok_only:
            filters.append(WatchProbeRecord.probe_status == "ok")
        q = (
            select(WatchProbeRecord)
            .where(*filters)
            .order_by(WatchProbeRecord.recorded_at.asc())
        )
        rows = (await self.session.execute(q)).scalars().all()
        return [
            {
                "server_id": r.server_id,
                "recorded_at": r.recorded_at.isoformat() if r.recorded_at else None,
                "cpu_pct": r.cpu_pct,
                "mem_pct": r.mem_pct,
                "disk_pct": r.disk_pct,
                "net_up_bps": r.net_up_bps,
                "net_down_bps": r.net_down_bps,
                "disk_read_bps": r.disk_read_bps,
                "disk_write_bps": r.disk_write_bps,
                "probe_status": r.probe_status,
            }
            for r in rows
        ]

    async def list_sessions(self, admin_id: int, *, since: datetime) -> list[dict[str, Any]]:
        q = (
            select(WatchPinSession, Server.name)
            .join(Server, WatchPinSession.server_id == Server.id)
            .where(WatchPinSession.admin_id == admin_id, WatchPinSession.started_at >= since)
            .order_by(WatchPinSession.started_at.desc())
        )
        rows = (await self.session.execute(q)).all()
        out: list[dict[str, Any]] = []
        for sess, server_name in rows:
            ok_q = select(func.count()).select_from(WatchProbeRecord).where(
                WatchProbeRecord.session_id == sess.id,
                WatchProbeRecord.probe_status == "ok",
            )
            fail_q = select(func.count()).select_from(WatchProbeRecord).where(
                WatchProbeRecord.session_id == sess.id,
                WatchProbeRecord.probe_status != "ok",
            )
            ok_count = int((await self.session.execute(ok_q)).scalar_one() or 0)
            fail_count = int((await self.session.execute(fail_q)).scalar_one() or 0)
            out.append(
                {
                    "session_id": sess.id,
                    "server_id": sess.server_id,
                    "server_name": server_name,
                    "slot_index": sess.slot_index,
                    "started_at": sess.started_at.isoformat() if sess.started_at else None,
                    "ended_at": sess.ended_at.isoformat() if sess.ended_at else None,
                    "end_reason": sess.end_reason,
                    "probe_ok_count": ok_count,
                    "probe_fail_count": fail_count,
                }
            )
        return out