import discord
import logging
import os
from gradio_client import Client

# 로깅 설정
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])

# 인텐트 설정
intents = discord.Intents.default()
intents.message_content = True  # 메시지 내용 수신 인텐트 활성화
intents.messages = True
intents.guilds = True
intents.guild_messages = True

# 특정 채널 ID 설정
SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))

class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')

    async def on_message(self, message):
        if message.author == self.user or message.channel.id != SPECIFIC_CHANNEL_ID:
            return
        if message.content.startswith("!generate"):  # 프롬프트 시작 명령어
            prompt = message.content[len("!generate "):]  # 명령어 제거 후 프롬프트 추출
            image_path = await self.generate_image(prompt)
            if image_path:
                await message.channel.send(file=discord.File(image_path))
            else:
                await message.channel.send("이미지를 생성하지 못했습니다.")

    async def generate_image(self, prompt):
        client = Client("ginipick/komodel")
        result = client.predict(
            prompt=prompt,
            negative="low quality, low quality, (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)",
            width=1024,
            height=1024,
            scale=7,
            steps=50,
            api_name="/generate_image"
        )
        if 'image_path' in result:
            return result['image_path']  # API가 반환한 이미지 경로
        return None

if __name__ == "__main__":
    discord_token = os.getenv('DISCORD_TOKEN')
    client = MyClient(intents=intents)
    client.run(discord_token)