main.py

import asyncio
import glob
import logging
import logging.handlers
import os
import re
import yaml
from dataclasses import dataclass
from json import JSONDecodeError

from aiogram import Bot, Dispatcher, F
from aiogram.filters import CommandStart
from aiogram.types import Message
from aiogram.enums import ChatType
from aiogram.client.default import DefaultBotProperties
from aiogram import Router
from dotenv import load_dotenv

# AICODE-NOTE: [CONTEXT] Minimal single-file bot to moderate images by OCR via OpenAI API


logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
def _setup_file_logging() -> None:
    """Configure daily rotating file logging and keep 30 days."""
    log_dir = "logs"
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, "tgimguard.log")
    file_handler = logging.handlers.TimedRotatingFileHandler(
        log_path, when="midnight", backupCount=30, encoding="utf-8"
    )
    file_handler.setLevel(logging.INFO)
    file_formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s() - %(message)s"
    )
    file_handler.setFormatter(file_formatter)
    logging.getLogger().addHandler(file_handler)

_setup_file_logging()

# Load .env early
load_dotenv()

router = Router()


def get_env(name: str) -> str:
    """Get required environment variable or raise."""
    value = os.getenv(name)
    if not value:
        raise RuntimeError(f"Missing required env var: {name}")
    return value


OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_MODEL_ID = os.getenv("OPENAI_MODEL_ID", "gpt-4o-mini")
CHAT_OWNER_ID = os.getenv("CHAT_OWNER_ID")  # optional, but recommended
IMAGE_COMPRESSION_THRESHOLD_KB = int(os.getenv("IMAGE_COMPRESSION_THRESHOLD_KB", "300"))
IMAGE_RESIZE_WIDTH_PX = int(os.getenv("IMAGE_RESIZE_WIDTH_PX", "800"))


@dataclass
class Rule:
    """Moderation rule with patterns and minimum match threshold."""
    name: str
    min_matches: int
    patterns: list[tuple[int, re.Pattern[str]]]  # (pattern_index, compiled_pattern)


def _load_rules() -> list[Rule]:
    """Load moderation rules from rules/*.yaml files.
    
    Each YAML file represents one rule with structure:
    name: rule name
    min_matches: 2
    patterns:
      - type: literal
        value: text
      - type: regex
        value: pattern
    """
    base_dir = os.path.join(os.path.dirname(__file__), "rules")
    if not os.path.exists(base_dir):
        logger.warning("Rules directory does not exist: %s", base_dir)
        return []
    
    rules: list[Rule] = []
    yaml_files = glob.glob(os.path.join(base_dir, "*.yaml"))
    
    for yaml_path in yaml_files:
        try:
            with open(yaml_path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f)
            
            if not data:
                logger.warning("Empty YAML file: %s", yaml_path)
                continue
            
            name = data.get("name", os.path.basename(yaml_path))
            min_matches = data.get("min_matches", 1)
            pattern_defs = data.get("patterns", [])
            
            compiled_patterns: list[tuple[int, re.Pattern[str]]] = []
            for idx, pattern_def in enumerate(pattern_defs):
                pattern_type = pattern_def.get("type", "literal")
                value = pattern_def.get("value", "")
                
                if not value:
                    logger.warning("Empty pattern in %s, skipping", yaml_path)
                    continue
                
                try:
                    if pattern_type == "literal":
                        compiled = re.compile(re.escape(value), re.IGNORECASE)
                    elif pattern_type == "regex":
                        compiled = re.compile(value)
                    else:
                        logger.warning("Unknown pattern type '%s' in %s, skipping", pattern_type, yaml_path)
                        continue
                    
                    compiled_patterns.append((idx, compiled))
                except Exception as e:
                    logger.warning("Failed to compile pattern in %s: %s", yaml_path, e)
            
            if compiled_patterns:
                rule = Rule(name=name, min_matches=min_matches, patterns=compiled_patterns)
                rules.append(rule)
                logger.info("Loaded rule '%s' with %d patterns (min_matches=%d)", 
                           name, len(compiled_patterns), min_matches)
            else:
                logger.warning("No valid patterns in %s", yaml_path)
                
        except Exception as e:
            logger.warning("Failed to load rule from %s: %s", yaml_path, e)
    
    logger.info("Loaded %d rule(s) total", len(rules))
    return rules


RULES = _load_rules()



async def get_chat_owner_id(bot: Bot, chat_id: int) -> int | None:
    """Return creator user id for chat, or None if not found."""
    try:
        admins = await bot.get_chat_administrators(chat_id)
    except Exception as e:
        logger.warning("Failed to fetch chat administrators: %s", e)
        return None
    for admin in admins:
        # Aiogram v3 represents owner as status "creator"
        if getattr(admin, "status", None) == "creator":
            return admin.user.id
    return None


@router.message(CommandStart())
async def start(message: Message) -> None:
    """Reply with brief info."""
    await message.answer("Bot is active. Send an image in a group to be moderated.")


def _detect_mime_type(image_bytes: bytes) -> str:
    """Return best-effort MIME type for given image bytes.

    Checks magic numbers of common formats; falls back to image/jpeg if unknown.
    """
    header = image_bytes[:12]
    if header.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if header[:3] == b"\xff\xd8\xff":
        return "image/jpeg"
    if header.startswith(b"GIF87a") or header.startswith(b"GIF89a"):
        return "image/gif"
    if header.startswith(b"RIFF") and header[8:12] == b"WEBP":
        return "image/webp"
    return "image/jpeg"


def _prepare_image_for_llm(image_bytes: bytes) -> tuple[bytes, str]:
    """Resize to configured width and convert to JPEG if size exceeds threshold.

    Returns (processed_bytes, mime_type).
    """
    threshold_bytes = max(0, IMAGE_COMPRESSION_THRESHOLD_KB) * 1024
    if threshold_bytes and len(image_bytes) > threshold_bytes:
        try:
            import io
            from PIL import Image  # type: ignore

            with Image.open(io.BytesIO(image_bytes)) as img:
                # Convert to RGB for JPEG compatibility
                if img.mode not in ("RGB", "L"):
                    img = img.convert("RGB")

                width, height = img.size
                target_w = int(IMAGE_RESIZE_WIDTH_PX)
                if width > target_w and target_w > 0:
                    new_h = int(height * target_w / max(1, width))
                    img = img.resize((target_w, new_h), Image.LANCZOS)

                out = io.BytesIO()
                img.save(out, format="JPEG")
                processed = out.getvalue()
                logger.info(
                    "Image re-saved: new_size=%d bytes (old_size=%d bytes)",
                    len(processed),
                    len(image_bytes),
                )
                return processed, "image/jpeg"
        except Exception as e:
            logger.warning("Image preprocessing failed, using original: %s", e)

    # Below threshold or processing failed: keep original and detect mime
    return image_bytes, _detect_mime_type(image_bytes)


async def extract_text_via_openai(image_bytes: bytes) -> str:
    """Call OpenAI-compatible API with image to extract visible text via OCR.
    Uses multimodal prompt with input_image content type.
    """
    # Lazy import to keep surface minimal
    from openai import OpenAI

    client = OpenAI(base_url=OPENAI_BASE_URL, api_key=get_env("OPENAI_API_KEY"))

    # Preprocess image per limits and upload as data URL
    # For small images this is fine; for large uploads consider hosting.
    import base64

    processed_bytes, mime_type = _prepare_image_for_llm(image_bytes)
    b64 = base64.b64encode(processed_bytes).decode("ascii")
    data_url = f"data:{mime_type};base64,{b64}"

    system_prompt = (
        "You are an OCR engine. Extract ONLY readable text from the image. "
        "Return plain text without explanations."
    )

    # Retry up to 2 times on JSONDecodeError (total 3 attempts)
    for attempt in range(3):
        try:
            resp = client.chat.completions.create(
                model=OPENAI_MODEL_ID,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "Extract visible text from this image."},
                            {"type": "image_url", "image_url": {"url": data_url}},
                        ],
                    },
                ],
                temperature=0,
            )

            text = resp.choices[0].message.content or ""
            return text.strip()
        except JSONDecodeError:
            if attempt == 2:
                raise
            await asyncio.sleep(1)


def contains_prohibited(text: str) -> tuple[str, list[str]] | None:
    """Check if any rule is triggered. Return (rule_name, matches) or None.
    
    A rule is triggered when the number of unique pattern matches >= min_matches.
    Returns tuple of (rule name, list of matched strings) on trigger, None otherwise.
    """
    if not RULES:
        return None
    
    for rule in RULES:
        matched_pattern_indices = set()
        matches: list[str] = []
        
        for pattern_idx, pattern in rule.patterns:
            match = pattern.search(text)
            if match:
                matched_pattern_indices.add(pattern_idx)
                matches.append(match.group(0))
        
        if len(matched_pattern_indices) >= rule.min_matches:
            logger.info("Rule '%s' triggered: %d/%d patterns matched", 
                       rule.name, len(matched_pattern_indices), len(rule.patterns))
            return (rule.name, matches)
    
    return None


def _format_user_info(user) -> str:
    """Build a single-line string combining available user identifiers."""
    parts: list[str] = []
    username = getattr(user, "username", None)
    if username:
        uname = str(username)
        parts.append(uname if uname.startswith("@") else f"@{uname}")

    first_name = getattr(user, "first_name", None)
    last_name = getattr(user, "last_name", None)
    full_name = " ".join([n for n in [first_name, last_name] if n])
    if full_name:
        parts.append(full_name)

    user_id = getattr(user, "id", None)
    if user_id is not None:
        parts.append(f"id={user_id}")

    return " | ".join(parts) or "unknown user"


def _format_chat_info(chat) -> str:
    """Build a single-line string combining available chat identifiers."""
    parts: list[str] = []
    title = getattr(chat, "title", None)
    if title:
        parts.append(str(title))

    username = getattr(chat, "username", None)
    if username:
        uname = str(username)
        parts.append(uname if uname.startswith("@") else f"@{uname}")

    chat_type = getattr(chat, "type", None)
    if chat_type:
        parts.append(f"type={chat_type}")

    chat_id = getattr(chat, "id", None)
    if chat_id is not None:
        parts.append(f"id={chat_id}")

    return " | ".join(parts) or "unknown chat"


def _chat_tag(chat) -> str:
    """Return a short prefix identifying the chat for log messages."""
    username = getattr(chat, "username", None)
    if username:
        uname = str(username)
        return f"[{uname if uname.startswith('@') else '@' + uname}]"
    chat_id = getattr(chat, "id", None)
    return f"[id={chat_id}]" if chat_id is not None else "[id=?]"


@router.message(F.content_type.in_({"photo", "document"}))
async def on_image(message: Message, bot: Bot) -> None:
    """Handle images sent to group/supergroup chats."""
    if message.chat.type not in {ChatType.GROUP, ChatType.SUPERGROUP}:
        return

    file_id = None
    if message.photo:
        file_id = message.photo[-1].file_id
    elif message.document and (message.document.mime_type or "").startswith("image/"):
        file_id = message.document.file_id

    if not file_id:
        return

    try:
        file = await bot.get_file(file_id)
        file_path = file.file_path
        # Download bytes
        from aiohttp import ClientSession

        tg_file_url = f"https://api.telegram.org/file/bot{bot.token}/{file_path}"
        async with ClientSession() as session:
            async with session.get(tg_file_url) as resp:
                resp.raise_for_status()
                image_bytes = await resp.read()

        text = await extract_text_via_openai(image_bytes)
        logger.info("%s OCR extracted %d chars", _chat_tag(message.chat), len(text))
        preview = " ".join(text.split())[:50]
        logger.info("%s OCR text preview: %s", _chat_tag(message.chat), preview)

        detection = contains_prohibited(text)
        logger.info("%s Image evaluation: %s", _chat_tag(message.chat), "prohibited" if detection else "OK")

        if detection:
            rule_name, matches = detection
            
            # Delete message
            try:
                await bot.delete_message(chat_id=message.chat.id, message_id=message.message_id)
            except Exception as e:
                logger.warning("%s Failed to delete message: %s", _chat_tag(message.chat), e)

            # Ban user
            try:
                await bot.ban_chat_member(chat_id=message.chat.id, user_id=message.from_user.id)
            except Exception as e:
                logger.warning("%s Failed to ban user: %s", _chat_tag(message.chat), e)

            # Notify owner: env first, otherwise autodetect creator
            owner_id = int(CHAT_OWNER_ID) if CHAT_OWNER_ID else await get_chat_owner_id(bot, message.chat.id)
            if owner_id:
                try:
                    user_info = _format_user_info(message.from_user)
                    chat_info = _format_chat_info(message.chat)
                    matches_str = "\n".join(f"  • {m}" for m in matches)
                    await bot.send_message(
                        chat_id=owner_id,
                        text=(
                            f"User {user_info} banned in chat {chat_info}\n"
                            f"Rule triggered: <code>{rule_name}</code>\n"
                            f"Matched strings:\n{matches_str}\n\n"
                            f"Full text from image:\n<blockquote>{text}</blockquote>"
                        ),
                        parse_mode="HTML",
                    )
                except Exception as e:
                    logger.warning("%s Failed to notify owner: %s", _chat_tag(message.chat), e)
    except Exception as e:
        logger.exception("%s Image handling error: %s", _chat_tag(message.chat), e)


async def main() -> None:
    token = get_env("TELEGRAM_BOT_TOKEN")

    dp = Dispatcher()
    dp.include_router(router)

    bot = Bot(token=token, default=DefaultBotProperties(parse_mode=None))

    await dp.start_polling(bot)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except (KeyboardInterrupt, SystemExit):
        pass