punishment_manager.py

from datetime import datetime, timedelta
from typing import Optional, Dict
import discord

class PunishmentManager:
    def __init__(self, db):
        self.db = db
        
    async def get_punishment(self, user_id: int, guild_id: int) -> Optional[Dict]:
        """Retorna informações de punição de um usuário"""
        return await self.db.fetch_one(
            "SELECT * FROM punishments WHERE user_id = ? AND guild_id = ?",
            (user_id, guild_id)
        )
        
    async def add_warning(self, user_id: int, guild_id: int, warned_by: int, 
                         max_warnings: int = 3, ban_duration_hours: int = 24) -> tuple[int, bool]:
        """Adiciona um aviso ao usuário e retorna (total_avisos, foi_banido)"""
        punishment = await self.get_punishment(user_id, guild_id)
        
        if not punishment:
            await self.db.execute(
                """INSERT INTO punishments (guild_id, user_id, warnings, banned_by)
                VALUES (?, ?, 1, ?)""",
                (guild_id, user_id, warned_by)
            )
            return 1, False
            
        new_warnings = punishment['warnings'] + 1
        
        if new_warnings >= max_warnings:
            ban_until = datetime.utcnow() + timedelta(hours=ban_duration_hours)
            await self.db.execute(
                """UPDATE punishments 
                SET warnings = ?, banned_until = ?, reason = ?, banned_by = ?
                WHERE user_id = ? AND guild_id = ?""",
                (new_warnings, ban_until, "Excesso de avisos", warned_by, user_id, guild_id)
            )
            return new_warnings, True
        else:
            await self.db.execute(
                """UPDATE punishments SET warnings = ?, banned_by = ?
                WHERE user_id = ? AND guild_id = ?""",
                (new_warnings, warned_by, user_id, guild_id)
            )
            return new_warnings, False
            
    async def ban_user(self, user_id: int, guild_id: int, duration_hours: int, 
                      reason: str, banned_by: int):
        """Bane um usuário temporariamente"""
        ban_until = datetime.utcnow() + timedelta(hours=duration_hours)
        
        await self.db.execute(
            """INSERT INTO punishments 
            (guild_id, user_id, banned_until, reason, banned_by)
            VALUES (?, ?, ?, ?, ?)
            ON CONFLICT(user_id, guild_id) DO UPDATE SET
            banned_until = ?, reason = ?, banned_by = ?""",
            (guild_id, user_id, ban_until, reason, banned_by, 
             ban_until, reason, banned_by)
        )
        
    async def unban_user(self, user_id: int, guild_id: int):
        """Remove o banimento de um usuário"""
        await self.db.execute(
            """UPDATE punishments 
            SET banned_until = NULL, reason = NULL
            WHERE user_id = ? AND guild_id = ?""",
            (user_id, guild_id)
        )
        
    async def is_banned(self, user_id: int, guild_id: int) -> tuple[bool, Optional[str]]:
        """Verifica se um usuário está banido e retorna (está_banido, motivo)"""
        punishment = await self.get_punishment(user_id, guild_id)
        
        if not punishment or not punishment['banned_until']:
            return False, None
            
        ban_until = datetime.fromisoformat(punishment['banned_until'])
        
        if datetime.utcnow() < ban_until:
            return True, punishment['reason']
        else:
            await self.unban_user(user_id, guild_id)
            return False, None
            
    async def clear_warnings(self, user_id: int, guild_id: int):
        """Limpa os avisos de um usuário"""
        await self.db.execute(
            """UPDATE punishments SET warnings = 0
            WHERE user_id = ? AND guild_id = ?""",
            (user_id, guild_id)
        )
        
    def format_punishment_embed(self, user: discord.Member, punishment: Dict, 
                               color: int) -> discord.Embed:
        """Formata um embed com informações de punição"""
        embed = discord.Embed(
            title=f"⚠️ Punições de {user.display_name}",
            color=color,
            timestamp=datetime.utcnow()
        )
        
        embed.set_thumbnail(url=user.display_avatar.url)
        
        embed.add_field(
            name="Avisos",
            value=f"`{punishment['warnings']}`",
            inline=True
        )
        
        if punishment['banned_until']:
            ban_until = datetime.fromisoformat(punishment['banned_until'])
            embed.add_field(
                name="Banido até",
                value=f"<t:{int(ban_until.timestamp())}:F>",
                inline=True
            )
            
            if punishment['reason']:
                embed.add_field(
                    name="Motivo",
                    value=punishment['reason'],
                    inline=False
                )
        else:
            embed.add_field(
                name="Status",
                value="✅ Sem banimento ativo",
                inline=True
            )
            
        return embed