import discord
import logging
import os
import asyncio
import subprocess
from huggingface_hub import InferenceClient
from PIL import Image
import io

# 로깅 설정
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(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_ready(self):
        logging.info(f'{self.user} has logged in!')
        subprocess.Popen(["python", "web.py"])  # Start the web.py server as a separate process
        logging.info("Web.py server has been started.")

    async def on_message(self, message):
        if message.author == self.user or not message.content:
            return
        if message.channel.id != SPECIFIC_CHANNEL_ID and not isinstance(message.channel, discord.Thread):
            return
        if self.is_processing:
            return
        self.is_processing = True
        try:
            image_path = await generate_image(message.content)
            if image_path:
                await send_image(message.channel, image_path)
            else:
                await message.channel.send("Failed to generate the image.")
        finally:
            self.is_processing = False

async def generate_image(prompt):
    try:
        # 모델 호출
        response = hf_client(inputs={"inputs": prompt}, options={"wait_for_model": True})  # options는 필요에 따라 조정 가능
        image_data = response.get('images', [None])[0]  # 응답에서 이미지 데이터 추출
        if image_data:
            image_bytes = io.BytesIO(base64.b64decode(image_data))
            image = Image.open(image_bytes)
            image_path = "output.png"
            image.save(image_path)
            return image_path
        else:
            raise ValueError("No image data received.")
    except Exception as e:
        logging.error(f"Failed to generate image: {str(e)}")
        return None



async def send_image(channel, image_path):
    file = discord.File(image_path)
    await channel.send(file=file)

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