app.py

import os
import discord
import time
from fastapi import FastAPI
import uvicorn
import asyncio
from threading import Thread
from discord import app_commands, ui
# 导入自定义 DNS 库
from dns_resolver import dns_manager
# 导入数据库基础类
from mongo import MongoManager
# 导入学科业务函数和学科列表
from material import (
    CATEGORIES, 
    upload_material_logic, 
    approve_material_logic, 
    delete_material_logic, 
    list_materials_logic
)

# --- 注入 DNS 补丁 ---
# 必须在所有网络连接尝试之前执行
dns_manager.patch_socket()

# 初始化数据库管理器
db_manager = MongoManager(os.getenv("MONGO_URL"), "tonyd")

# 身份组权限列表
mod = [1357078628969746462, 1357079285113819146, 1436537252317626449, 1357079707539214417, 1357079889160835183]

# --- 1. FastAPI 保活 ---
app = FastAPI()
# 使用显式的方法定义，并使用 JSONResponse 确保兼容性
@app.api_route("/", methods=["GET", "HEAD"])
async def status():
    # 使用 JSONResponse 包装，防止 Starlette 误判
    return {"status": "hi"}

def run_flask():
    # Hugging Face 默认使用 7860 端口
    port = int(os.environ.get('PORT', 7860))
    uvicorn.run(app, host='0.0.0.0', port=port)

# --- 2. 辅助函数 ---
def format_pings(role_input: str):
    pass

# --- 3. Modal 表单定义 ---
class AnnounceModal(ui.Modal):
    def __init__(self, channel: discord.TextChannel, ping_roles_str: str = None):
        super().__init__(title="New Announcement")
        self.channel = channel
        self.ping_roles_str = ping_roles_str

    announcement_title = ui.TextInput(label="Title", placeholder="Enter title...", style=discord.TextStyle.short)
    announcement_body = ui.TextInput(label="Body", placeholder="Enter message...", style=discord.TextStyle.long)

    async def on_submit(self, interaction: discord.Interaction):
        content = self.announcement_body.value
        pings = self.ping_roles_str
        if pings: pass
        
        embed = discord.Embed(
            title=self.announcement_title.value,
            description=content,
            color=discord.Color.blue()
        )
        footer = f"\n\n||{pings}||\n\n---\n**Sent by:** ||{interaction.user.mention}||\n**Time:** <t:{int(time.time())}:f>"
        embed.description += footer

        await self.channel.send(embed=embed)
        await interaction.response.send_message(f"✅ Sent to {self.channel.mention}!", ephemeral=False)

class EditAnnounceModal(ui.Modal):
    def __init__(self, message: discord.Message, new_ping_roles_str: str = None):
        super().__init__(title="Edit Announcement")
        self.message = message
        self.new_ping_roles_str = new_ping_roles_str
        
        old_title = message.embeds[0].title if message.embeds else ""
        old_body = message.embeds[0].description.split("\n\n---")[0] if message.embeds else ""

        self.edit_title = ui.TextInput(label="Title", default=old_title, style=discord.TextStyle.short)
        self.edit_body = ui.TextInput(label="Body", default=old_body, style=discord.TextStyle.long)
        
        self.add_item(self.edit_title)
        self.add_item(self.edit_body)

    async def on_submit(self, interaction: discord.Interaction):
        content = self.edit_body.value
        pings = self.new_ping_roles_str
        if pings: pass

        new_embed = discord.Embed(
            title=self.edit_title.value,
            description=content,
            color=discord.Color.blue()
        )
        footer = f"\n\n||{pings}||\n\n---\n**Edited by:** ||{interaction.user.mention}||\n**Time:** <t:{int(time.time())}:f>"
        new_embed.description += footer

        await self.message.edit(embed=new_embed)
        await interaction.response.send_message(f"✅ Updated announcement {self.message.id} !", ephemeral=False)

# --- 4. 机器人主体 ---
class MyBot(discord.Client):
    def __init__(self):
        super().__init__(intents=discord.Intents.default())
        self.tree = app_commands.CommandTree(self)

    async def setup_hook(self):
        await self.tree.sync()
        print(f"✅ Slash commands synced")

client = MyBot()
client.db = db_manager 

# --- 5. 指令定义 ---
@client.tree.command(name="announce", description="Post announcement")
async def announce(interaction: discord.Interaction, channel: discord.TextChannel, role_ids: str = None):
    if not any(role.id in mod for role in interaction.user.roles):
        await interaction.response.send_message("❌ **Access Denied**", ephemeral=True)
        return
    await interaction.response.send_modal(AnnounceModal(channel, role_ids))

@client.tree.command(name="announce_edit", description="Edit announcement")
async def announce_edit(interaction: discord.Interaction, channel: discord.TextChannel, message_id: str, new_role_ids: str = None):
    if not any(role.id in mod for role in interaction.user.roles):
        await interaction.response.send_message("❌ **Access Denied**", ephemeral=True)
        return
    try:
        message = await channel.fetch_message(int(message_id))
        await interaction.response.send_modal(EditAnnounceModal(message, new_role_ids))
    except Exception as e:
        await interaction.response.send_message(f"❌ Error: {e}", ephemeral=False)

@client.tree.command(name="wismer_commands", description="Show all avaliable commands")
async def wismer_commands(interaction: discord.Interaction):
    await interaction.response.send_message(
        '''**Avaliable Commands**:
        `/announce`: Post an announcement.
        `/announce_edit`: Edit an announcement.
        `/mats`: Show all materials.
        `/uploadmat`: Upload a material.
        `/approvemat`: Approve a unapproved material. Only Mod+ can use it.
        `/deletemat`: To delete or reject a material. Only Mod+ can use it.
        `/unappmats`: To show unapprove materials. Only Mod+ can use it.''',
        ephemeral=False
    )

ADMIN_LOG_CHANNEL_ID = 1468391869686878455 

@client.tree.command(name="uploadmat", description="Upload a study link")
@app_commands.choices(category=[app_commands.Choice(name=c, value=c) for c in CATEGORIES])
async def uploadmat(interaction: discord.Interaction, category: str, title: str, url: str):
    await interaction.response.defer(ephemeral=True)
    if not url.startswith(("http://", "https://")):
        return await interaction.followup.send("❌ Invalid URL!", ephemeral=True)
    try:
        new_id = await upload_material_logic(client.db, category, title, url, interaction.user.id)
        await interaction.followup.send(f"✅ Submitted! ID: `#{new_id}`.", ephemeral=True)
        log_channel = client.get_channel(ADMIN_LOG_CHANNEL_ID)
        if log_channel:
            embed = discord.Embed(title="🔔 New Material Pending", color=discord.Color.orange())
            embed.add_field(name="Subject", value=category, inline=True)
            embed.add_field(name="Index ID", value=f"`#{new_id}`", inline=True)
            await log_channel.send(embed=embed)
    except Exception as e:
        await interaction.followup.send(f"❌ Error: {e}", ephemeral=True)

@client.tree.command(name="approvemat", description="[Mod Only] Approve material")
@app_commands.choices(category=[app_commands.Choice(name=c, value=c) for c in CATEGORIES])
async def approvemat(interaction: discord.Interaction, category: str, index_id: str):
    await interaction.response.defer(ephemeral=False)
    if not any(role.id in mod for role in interaction.user.roles):
        return await interaction.followup.send("❌ Access Denied", ephemeral=True)
    await approve_material_logic(client.db, category, index_id)
    await interaction.followup.send(f"✅ Approved **{category}** `#{index_id}`.")

@client.tree.command(name="deletemat", description="[Mod Only] Delete material")
@app_commands.choices(category=[app_commands.Choice(name=c, value=c) for c in CATEGORIES])
async def deletemat(interaction: discord.Interaction, category: str, index_id: str):
    await interaction.response.defer(ephemeral=False)
    if not any(role.id in mod for role in interaction.user.roles):
        return await interaction.followup.send("❌ Access Denied", ephemeral=True)
    await delete_material_logic(client.db, category, index_id)
    await interaction.followup.send(f"🗑️ Deleted ID `#{index_id}`.")

@client.tree.command(name="mats", description="View approved materials")
@app_commands.choices(category=[app_commands.Choice(name=c, value=c) for c in CATEGORIES])
async def mats(interaction: discord.Interaction, category: str):
    await interaction.response.defer(ephemeral=False)
    materials = await list_materials_logic(client.db, category, status="approved")
    if not materials:
        return await interaction.followup.send(f"📭 No materials in **{category}**.")
    list_text = "\n".join([f"`#{m['_id']}` [{m['title']}]({m['url']})" for m in materials])
    embed = discord.Embed(title=f"📚 {category}", description=list_text, color=discord.Color.blue())
    await interaction.followup.send(embed=embed)

@client.tree.command(name="unappmats", description="[Mod Only] View pending")
@app_commands.choices(category=[app_commands.Choice(name=c, value=c) for c in CATEGORIES])
async def unappmats(interaction: discord.Interaction, category: str):
    await interaction.response.defer(ephemeral=False)
    if not any(role.id in mod for role in interaction.user.roles):
        return await interaction.followup.send("❌ Access Denied", ephemeral=True)
    materials = await list_materials_logic(client.db, category, status="pending")
    if not materials:
        return await interaction.followup.send(f"✅ No pending items.")
    list_text = "\n".join([f"`#{m['_id']}` **{m['title']}**\n🔗 {m['url']}" for m in materials])
    embed = discord.Embed(title=f"⏳ {category} Pending", description=list_text, color=discord.Color.orange())
    await interaction.followup.send(embed=embed)

# --- 6. 启动逻辑 ---
async def start_bot():
    token = os.getenv('DISCORD_TOKEN')
    while True:
        try:
            print("正在尝试连接 Discord...")
            await client.start(token)
        except Exception as e:
            print(f"❌ 连接失败: {e}，15秒后重试...")
            # 报错后尝试关闭客户端，清理 aiohttp session 缓存
            try:
                await client.close()
            except:
                pass
            await asyncio.sleep(15)

if __name__ == "__main__":
    # 启动 Web 保活
    t = Thread(target=run_flask)
    t.daemon = True
    t.start()
    
    # 启动机器人异步循环
    try:
        asyncio.run(start_bot())
    except KeyboardInterrupt:
        pass