app/database.py

"""Database layer — Postgres via asyncpg with a tiny chat-settings cache."""
from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Optional

import asyncpg

from app.config import DEFAULTS

LOGGER = logging.getLogger(__name__)

SCHEMA_PATH = Path(__file__).parent.parent / "scripts" / "init_db.sql"


@dataclass
class ChatSettings:
    chat_id: int
    enabled: bool = DEFAULTS["enabled"]
    check_admins: bool = DEFAULTS["check_admins"]
    threshold: int = DEFAULTS["threshold"]

    filter_photos: bool = DEFAULTS["filter_photos"]
    filter_stickers: bool = DEFAULTS["filter_stickers"]
    filter_videos: bool = DEFAULTS["filter_videos"]
    filter_gifs: bool = DEFAULTS["filter_gifs"]
    filter_video_notes: bool = DEFAULTS["filter_video_notes"]
    filter_web_preview: bool = DEFAULTS["filter_web_preview"]

    block_porn: bool = DEFAULTS["block_porn"]
    block_hentai: bool = DEFAULTS["block_hentai"]
    block_sexy: bool = DEFAULTS["block_sexy"]

    antispam_enabled: bool = DEFAULTS["antispam_enabled"]
    warn_user: bool = DEFAULTS["warn_user"]
    max_violations: int = DEFAULTS["max_violations"]
    violation_window_minutes: int = DEFAULTS["violation_window_minutes"]
    ban_duration_minutes: int = DEFAULTS["ban_duration_minutes"]
    warn_delete_seconds: int = DEFAULTS["warn_delete_seconds"]

    delete_message_template: str = DEFAULTS["delete_message_template"]

    def blocked_categories(self) -> list[str]:
        out: list[str] = []
        if self.block_porn:
            out.append("porn")
        if self.block_hentai:
            out.append("hentai")
        if self.block_sexy:
            out.append("sexy")
        return out


@dataclass
class _CacheEntry:
    settings: ChatSettings
    fetched_at: float = field(default_factory=time.monotonic)


class Database:
    """Thin asyncpg wrapper. Pool is lazy."""

    CACHE_TTL_SECONDS = 60.0

    def __init__(self, dsn: str):
        self._dsn = dsn
        self._pool: Optional[asyncpg.Pool] = None
        self._cache: Dict[int, _CacheEntry] = {}

    async def connect(self) -> None:
        if self._pool is not None:
            return
        # Conservative pool — Render free has tight memory
        self._pool = await asyncpg.create_pool(
            dsn=self._dsn,
            min_size=1,
            max_size=4,
            command_timeout=10.0,
        )
        await self._init_schema()
        LOGGER.info("Database pool ready")

    async def close(self) -> None:
        if self._pool is not None:
            await self._pool.close()
            self._pool = None

    async def _init_schema(self) -> None:
        sql = SCHEMA_PATH.read_text(encoding="utf-8")
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            await conn.execute(sql)

    # ------------------------------------------------------------------
    # Chat lifecycle
    # ------------------------------------------------------------------
    async def upsert_chat(
        self,
        chat_id: int,
        title: Optional[str],
        added_by: Optional[int],
        member_count: Optional[int] = None,
        initial_status: str = "pending",
    ) -> str:
        """Insert chat or update title. Returns the resulting status string.

        Behaviour:
        - First insert uses ``initial_status`` (so the caller can pass 'approved'
          when the bot is added by the owner themselves).
        - On conflict (re-add) we keep prior status if it was 'approved' or
          'rejected' (owner already decided), otherwise reset to 'pending'.
        """
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                INSERT INTO chats (chat_id, title, added_by, status, status_at, member_count)
                VALUES ($1, $2, $3, $4, NOW(), $5)
                ON CONFLICT (chat_id) DO UPDATE
                    SET title = EXCLUDED.title,
                        member_count = COALESCE(EXCLUDED.member_count, chats.member_count),
                        updated_at = NOW(),
                        status = CASE
                            WHEN chats.status IN ('approved','rejected') THEN chats.status
                            ELSE 'pending'
                        END,
                        status_at = CASE
                            WHEN chats.status IN ('approved','rejected') THEN chats.status_at
                            ELSE NOW()
                        END
                RETURNING status
                """,
                chat_id,
                title,
                added_by,
                initial_status,
                member_count,
            )
            await conn.execute(
                """
                INSERT INTO chat_settings (chat_id) VALUES ($1)
                ON CONFLICT (chat_id) DO NOTHING
                """,
                chat_id,
            )
        self._cache.pop(chat_id, None)
        return row["status"] if row else initial_status

    async def set_chat_status(
        self,
        chat_id: int,
        status: str,
        decided_by: Optional[int] = None,
    ) -> None:
        assert self._pool is not None
        if status not in ("pending", "approved", "rejected", "removed"):
            raise ValueError(f"Invalid chat status: {status}")
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                UPDATE chats
                   SET status = $1, status_at = NOW(), status_by = $2, updated_at = NOW()
                 WHERE chat_id = $3
                """,
                status, decided_by, chat_id,
            )
        self._cache.pop(chat_id, None)

    async def get_chat(self, chat_id: int) -> Optional[Dict[str, Any]]:
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT chat_id, title, added_by, status, status_at, status_by, member_count "
                "FROM chats WHERE chat_id = $1",
                chat_id,
            )
        return dict(row) if row else None

    async def is_chat_approved(self, chat_id: int) -> bool:
        chat = await self.get_chat(chat_id)
        return bool(chat and chat["status"] == "approved")

    async def list_chats_for_admin(self, admin_user_id: int) -> list[Dict[str, Any]]:
        """Best-effort list of chats that the bot knows about, used by /api/chats.

        We can't query Telegram membership cheaply for every chat, so we return
        chats where:
          - status is approved (so the user can pick from active ones), AND
          - the user is the original `added_by` OR the global owner.
        Membership re-verification happens per-chat on the dedicated endpoints.
        """
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            rows = await conn.fetch(
                """
                SELECT chat_id, title, status, member_count, added_at, status_at
                FROM chats
                WHERE status = 'approved'
                ORDER BY status_at DESC
                """,
            )
        return [dict(r) for r in rows]

    async def mark_chat_removed(self, chat_id: int) -> None:
        await self.set_chat_status(chat_id, "removed")

    async def delete_chat(self, chat_id: int) -> None:
        """Hard delete; kept for tests and admin tooling."""
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            await conn.execute("DELETE FROM chats WHERE chat_id = $1", chat_id)
        self._cache.pop(chat_id, None)

    # ------------------------------------------------------------------
    # Settings
    # ------------------------------------------------------------------
    async def get_settings(self, chat_id: int) -> ChatSettings:
        cached = self._cache.get(chat_id)
        if cached and (time.monotonic() - cached.fetched_at) < self.CACHE_TTL_SECONDS:
            return cached.settings

        assert self._pool is not None
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT * FROM chat_settings WHERE chat_id = $1",
                chat_id,
            )

        if row is None:
            settings = ChatSettings(chat_id=chat_id)
        else:
            settings = ChatSettings(
                chat_id=chat_id,
                enabled=row["enabled"],
                check_admins=row["check_admins"],
                threshold=row["threshold"],
                filter_photos=row["filter_photos"],
                filter_stickers=row["filter_stickers"],
                filter_videos=row["filter_videos"],
                filter_gifs=row["filter_gifs"],
                filter_video_notes=row["filter_video_notes"],
                filter_web_preview=row["filter_web_preview"],
                block_porn=row["block_porn"],
                block_hentai=row["block_hentai"],
                block_sexy=row["block_sexy"],
                antispam_enabled=row["antispam_enabled"],
                warn_user=row["warn_user"],
                max_violations=row["max_violations"],
                violation_window_minutes=row["violation_window_minutes"],
                ban_duration_minutes=row["ban_duration_minutes"],
                warn_delete_seconds=row["warn_delete_seconds"],
                delete_message_template=row["delete_message_template"],
            )

        self._cache[chat_id] = _CacheEntry(settings=settings)
        return settings

    async def update_setting(self, chat_id: int, key: str, value: Any) -> None:
        # Whitelist of mutable columns to avoid SQL injection via dynamic key
        allowed = {
            "enabled", "check_admins", "threshold",
            "filter_photos", "filter_stickers", "filter_videos",
            "filter_gifs", "filter_video_notes", "filter_web_preview",
            "block_porn", "block_hentai", "block_sexy",
            "antispam_enabled", "warn_user", "max_violations",
            "violation_window_minutes", "ban_duration_minutes",
            "warn_delete_seconds", "delete_message_template",
        }
        if key not in allowed:
            raise ValueError(f"Cannot update unknown setting: {key}")

        assert self._pool is not None
        sql = (
            f"UPDATE chat_settings SET {key} = $1, updated_at = NOW() "
            f"WHERE chat_id = $2"
        )
        async with self._pool.acquire() as conn:
            await conn.execute(sql, value, chat_id)
        self._cache.pop(chat_id, None)

    # ------------------------------------------------------------------
    # Violations
    # ------------------------------------------------------------------
    async def record_violation(
        self,
        chat_id: int,
        user_id: int,
        category: Optional[str],
        score: Optional[float],
        media_kind: str,
    ) -> int:
        """Inserts a violation, returns count of recent violations within window."""
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO violations (chat_id, user_id, category, score, media_kind)
                VALUES ($1, $2, $3, $4, $5)
                """,
                chat_id, user_id, category, score, media_kind,
            )
            settings = await self.get_settings(chat_id)
            window_minutes = settings.violation_window_minutes
            count = await conn.fetchval(
                """
                SELECT COUNT(*) FROM violations
                WHERE chat_id = $1 AND user_id = $2
                  AND occurred_at > NOW() - ($3 || ' minutes')::INTERVAL
                """,
                chat_id, user_id, str(window_minutes),
            )
        return int(count)

    # ------------------------------------------------------------------
    # Events (timeline backing the mini-app stats + log screens)
    # ------------------------------------------------------------------
    async def record_event(
        self,
        chat_id: int,
        kind: str,
        user_id: Optional[int] = None,
        user_name: Optional[str] = None,
        category: Optional[str] = None,
        score: Optional[float] = None,
        media_kind: Optional[str] = None,
        meta: Optional[Dict[str, Any]] = None,
    ) -> None:
        import json
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO events
                    (chat_id, kind, user_id, user_name, category, score, media_kind, meta)
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb)
                """,
                chat_id, kind, user_id, user_name, category, score, media_kind,
                json.dumps(meta) if meta else None,
            )

    async def stats_for_chat(self, chat_id: int) -> Dict[str, Any]:
        """Aggregate stats for a single chat in one DB round-trip.

        Returns the shape consumed by /api/stats: d24, week, week_delta_pct,
        total, by_category, daily (last 7d), top_violators (last 30d top 5).
        """
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            row = await conn.fetchrow(
                """
                WITH del AS (
                    SELECT * FROM events WHERE chat_id = $1 AND kind = 'delete'
                )
                SELECT
                    (SELECT COUNT(*) FROM del WHERE occurred_at > NOW() - INTERVAL '24 hours') AS d24,
                    (SELECT COUNT(*) FROM del WHERE occurred_at > NOW() - INTERVAL '7 days')   AS week,
                    (SELECT COUNT(*) FROM del WHERE occurred_at > NOW() - INTERVAL '14 days'
                                              AND occurred_at <= NOW() - INTERVAL '7 days')   AS prev_week,
                    (SELECT COUNT(*) FROM del)                                                  AS total
                """,
                chat_id,
            )
            cats = await conn.fetch(
                """
                SELECT COALESCE(category,'unknown') AS key, COUNT(*) AS count
                FROM events
                WHERE chat_id = $1 AND kind = 'delete'
                  AND occurred_at > NOW() - INTERVAL '7 days'
                GROUP BY 1
                ORDER BY count DESC
                """,
                chat_id,
            )
            daily = await conn.fetch(
                """
                SELECT to_char(date_trunc('day', occurred_at), 'YYYY-MM-DD') AS day,
                       COUNT(*) AS count
                FROM events
                WHERE chat_id = $1 AND kind = 'delete'
                  AND occurred_at > NOW() - INTERVAL '7 days'
                GROUP BY 1
                ORDER BY 1
                """,
                chat_id,
            )
            top = await conn.fetch(
                """
                SELECT user_id,
                       COALESCE(MAX(user_name), 'unknown') AS name,
                       COUNT(*) AS count
                FROM events
                WHERE chat_id = $1 AND kind = 'delete' AND user_id IS NOT NULL
                  AND occurred_at > NOW() - INTERVAL '30 days'
                GROUP BY user_id
                ORDER BY count DESC
                LIMIT 5
                """,
                chat_id,
            )

        d24 = int(row["d24"] or 0)
        week = int(row["week"] or 0)
        prev = int(row["prev_week"] or 0)
        total = int(row["total"] or 0)
        if prev > 0:
            delta = round((week - prev) * 100.0 / prev)
        elif week > 0:
            delta = 100
        else:
            delta = 0
        return {
            "d24": d24,
            "week": week,
            "week_delta_pct": delta,
            "total": total,
            "by_category": [{"key": r["key"], "count": int(r["count"])} for r in cats],
            "daily": [{"day": r["day"], "count": int(r["count"])} for r in daily],
            "top_violators": [
                {"user_id": int(r["user_id"]), "name": r["name"], "count": int(r["count"])}
                for r in top
            ],
        }

    async def events_for_chat(
        self,
        chat_id: int,
        kinds: Optional[list[str]] = None,
        limit: int = 100,
    ) -> list[Dict[str, Any]]:
        assert self._pool is not None
        if kinds:
            sql = (
                "SELECT id, occurred_at, kind, user_id, user_name, category, score, media_kind, meta "
                "FROM events WHERE chat_id = $1 AND kind = ANY($2::text[]) "
                "ORDER BY occurred_at DESC LIMIT $3"
            )
            params: tuple = (chat_id, kinds, limit)
        else:
            sql = (
                "SELECT id, occurred_at, kind, user_id, user_name, category, score, media_kind, meta "
                "FROM events WHERE chat_id = $1 ORDER BY occurred_at DESC LIMIT $2"
            )
            params = (chat_id, limit)
        async with self._pool.acquire() as conn:
            rows = await conn.fetch(sql, *params)
        out = []
        for r in rows:
            d = dict(r)
            if d.get("occurred_at") is not None:
                d["occurred_at"] = d["occurred_at"].isoformat()
            out.append(d)
        return out

    # ------------------------------------------------------------------
    # User language (RU/AZ/EN)
    # ------------------------------------------------------------------
    async def get_user_lang(self, user_id: int) -> str:
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            row = await conn.fetchval(
                "SELECT lang FROM users WHERE user_id = $1", user_id,
            )
        return row or "en"

    async def set_user_lang(self, user_id: int, lang: str) -> None:
        if lang not in ("ru", "az", "en"):
            raise ValueError(f"Unsupported language: {lang}")
        assert self._pool is not None
        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO users (user_id, lang) VALUES ($1, $2)
                ON CONFLICT (user_id) DO UPDATE
                    SET lang = EXCLUDED.lang, updated_at = NOW()
                """,
                user_id, lang,
            )