scripts/worker.py

# ruff: noqa: E402
import asyncio
import logging
import re
import sys
import time
from pathlib import Path

from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

load_dotenv(Path(__file__).resolve().parents[1] / ".env")
from datetime import datetime

from app.db.models import AIDetection, GroupConfig, MediaAsset, MessageEvent, StatsMessagesDaily, StatsWordcloudTerm
from app.db.session import SessionLocal
from app.services.ai_client import detect_image, detect_text
from app.services.queue import claim_job, complete_job, fail_job, get_pending_topics
from app.services.telegram_moderator import delete_message, punish_user
from app.services.telegram_sender import send_broadcast, send_post, send_scheduled_message

logger = logging.getLogger(__name__)

QUEUE_TOPICS = [
    "ai_text_detect",
    "ai_image_detect",
    "moderation_action",
    "broadcast_send",
    "post_send",
    "schedule_send",
    "schedule_tick",
    "stats_rollup",
]


async def _handle_ai_detection(message_id: str, kind: str) -> None:
    db = SessionLocal()
    try:
        message = db.get(MessageEvent, message_id)
        if not message:
            return
        config = db.query(GroupConfig).filter(GroupConfig.group_id == message.group_id).first()
        ad_cfg = (config.config_json.get("ad_filter") if config else {}) or {}
        if kind == "text" and not ad_cfg.get("text_ai_enabled", True):
            return
        if kind == "image" and not ad_cfg.get("image_ai_enabled", True):
            return
        threshold = float(ad_cfg.get("text_ai_threshold" if kind == "text" else "image_ai_threshold", 0.85))
        score = 0.0
        if kind == "text":
            score = await detect_text(message.text or "")
        else:
            file_id = ""
            if message.media_id:
                asset = db.get(MediaAsset, message.media_id)
                if asset:
                    file_id = asset.telegram_file_id
            score = await detect_image(file_id)
        result = "pass" if score < threshold else "reject"
        db.add(
            AIDetection(
                group_id=message.group_id,
                message_id=message.id,
                type=kind,
                score=score,
                threshold=threshold,
                model_id=None,
                result=result,
            )
        )
        db.commit()
        if result == "reject" and message.chat_id:
            action = ad_cfg.get("action", "delete")
            duration = ad_cfg.get("action_duration")
            if action == "delete":
                await delete_message(
                    chat_id=message.chat_id,
                    message_id=message.telegram_message_id,
                    group_id=str(message.group_id),
                    user_id=message.telegram_user_id,
                    reason="AI广告过滤",
                    source="ai",
                )
            else:
                await punish_user(
                    chat_id=message.chat_id,
                    user_id=message.telegram_user_id,
                    group_id=str(message.group_id),
                    action=action,
                    duration=duration,
                    reason="AI广告过滤",
                    source="ai",
                )
    finally:
        db.close()


def _tokenize(text: str) -> list[str]:
    tokens = re.findall(r"[A-Za-z0-9]+|[\u4e00-\u9fff]+", text)
    return [token for token in tokens if len(token) >= 2]


def _handle_stats_rollup(message_id: str) -> None:
    db = SessionLocal()
    try:
        message = db.get(MessageEvent, message_id)
        if not message:
            return
        day = (message.created_at or datetime.utcnow()).date()
        stats = (
            db.query(StatsMessagesDaily)
            .filter(StatsMessagesDaily.group_id == message.group_id, StatsMessagesDaily.date == day)
            .first()
        )
        if not stats:
            stats = StatsMessagesDaily(group_id=message.group_id, date=day, message_count=0, active_users=0)
            db.add(stats)
        stats.message_count += 1
        active_users = (
            db.query(MessageEvent.telegram_user_id)
            .filter(MessageEvent.group_id == message.group_id)
            .filter(MessageEvent.created_at >= datetime.combine(day, datetime.min.time()))
            .filter(MessageEvent.created_at <= datetime.combine(day, datetime.max.time()))
            .distinct()
            .count()
        )
        stats.active_users = active_users

        if message.text:
            for token in _tokenize(message.text):
                term = (
                    db.query(StatsWordcloudTerm)
                    .filter(
                        StatsWordcloudTerm.group_id == message.group_id,
                        StatsWordcloudTerm.date == day,
                        StatsWordcloudTerm.term == token,
                    )
                    .first()
                )
                if not term:
                    term = StatsWordcloudTerm(group_id=message.group_id, date=day, term=token, count=0)
                    db.add(term)
                term.count += 1
        db.commit()
    finally:
        db.close()


async def _dispatch(topic: str, payload: dict) -> None:
    """Dispatch a single queue message by topic."""
    if topic == "post_send":
        post_id = payload.get("post_id")
        if post_id:
            await send_post(post_id)
    elif topic == "broadcast_send":
        broadcast_id = payload.get("broadcast_id")
        if broadcast_id:
            await send_broadcast(broadcast_id)
    elif topic == "schedule_send":
        schedule_id = payload.get("schedule_id")
        if schedule_id:
            await send_scheduled_message(schedule_id)
    elif topic == "ai_text_detect":
        message_id = payload.get("message_id")
        if message_id:
            await _handle_ai_detection(message_id, "text")
    elif topic == "ai_image_detect":
        message_id = payload.get("message_id")
        if message_id:
            await _handle_ai_detection(message_id, "image")
    elif topic == "stats_rollup":
        message_id = payload.get("message_id")
        if message_id:
            # stats_rollup is sync (DB only), run in thread to avoid blocking
            await asyncio.to_thread(_handle_stats_rollup, message_id)
    else:
        logger.info("[worker] %s %s", topic, payload)


async def worker_loop() -> None:
    logger.info("worker started (PostgreSQL queue)")
    while True:
        try:
            pending_topics = await asyncio.to_thread(get_pending_topics, QUEUE_TOPICS)
            if not pending_topics:
                await asyncio.sleep(0.5)
                continue

            for topic in pending_topics:
                job = await asyncio.to_thread(claim_job, topic)
                if not job:
                    continue
                try:
                    await _dispatch(job.topic, job.payload)
                    await asyncio.to_thread(complete_job, job.id)
                except Exception:
                    logger.exception("worker error processing %s", job.topic)
                    await asyncio.to_thread(fail_job, job.id)
        except Exception:
            logger.exception("worker loop error")
            await asyncio.sleep(1)


def main():
    """Standalone entry point (backward compatible)."""
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    asyncio.run(worker_loop())


if __name__ == "__main__":
    main()