import discord
import logging
import os
from huggingface_hub import InferenceClient
import asyncio
import requests
from io import BytesIO

# 로깅 설정
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

# 추론 API 클라이언트 설정
hf_client = InferenceClient("Kwai-Kolors/Kolors", token=os.getenv("HF_TOKEN"))

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

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

    async def on_message(self, message):
        if message.author == self.user:
            return
        if not self.is_message_in_specific_channel(message):
            return
        if self.is_processing:
            return
        self.is_processing = True
        try:
            response = await generate_image(message)
            await message.channel.send(file=response)
        finally:
            self.is_processing = False

    def is_message_in_specific_channel(self, message):
        # 메시지가 지정된 채널이거나, 해당 채널의 쓰레드인 경우 True 반환
        return message.channel.id == SPECIFIC_CHANNEL_ID or (
            isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
        )

async def generate_image(message):
    user_input = message.content
    user_mention = message.author.mention
    system_message = f"{user_mention}, 텍스트를 기반으로 이미지를 생성합니다."

    logging.debug(f'User input: {user_input}')

    loop = asyncio.get_event_loop()
    response = await loop.run_in_executor(None, lambda: hf_client.text_to_image(
        prompt=user_input, num_return_sequences=1))

    image_url = response[0]['generated_image']
    logging.debug(f'Generated image URL: {image_url}')

    image_response = requests.get(image_url)
    image = BytesIO(image_response.content)
    image_file = discord.File(image, filename="generated_image.png")

    return image_file

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