import discord
import logging
import os
import uuid
import torch
import subprocess
from huggingface_hub import snapshot_download
from diffusers import StableDiffusion3Pipeline, StableDiffusion3Img2ImgPipeline
from transformers import pipeline

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

# Hugging Face 모델 다운로드
huggingface_token = os.getenv("HF_TOKEN")
model_path = snapshot_download(
    repo_id="stabilityai/stable-diffusion-3-medium",
    revision="refs/pr/26",
    repo_type="model",
    ignore_patterns=[".md", "..gitattributes"],
    local_dir="stable-diffusion-3-medium",
    token=huggingface_token,
)

# 모델 로드 함수
def load_pipeline(pipeline_type):
    logging.debug(f'Loading pipeline: {pipeline_type}')
    if pipeline_type == "text2img":
        return StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_fast=True)
    elif pipeline_type == "img2img":
        return StableDiffusion3Img2ImgPipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_fast=True)

# 디바이스 설정
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 번역 파이프라인 설정
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# 디스코드 봇 클래스
class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False
        self.text2img_pipeline = load_pipeline("text2img").to(device)
        self.text2img_pipeline.enable_attention_slicing()  # 메모리 최적화

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        subprocess.Popen(["python", "web.py"])
        logging.info("web.py 서버가 시작되었습니다.")

    async def on_message(self, message):
        if message.author == self.user:
            return
        if message.content.startswith('!image '):
            self.is_processing = True
            try:
                prompt = message.content[len('!image '):]
                prompt_en = translate_prompt(prompt)
                logging.debug(f'Translated prompt: {prompt_en}')
                image_path = await self.generate_image(prompt_en)
                await message.channel.send(file=discord.File(image_path, 'generated_image.png'))
            finally:
                self.is_processing = False

    async def generate_image(self, prompt):
        generator = torch.Generator(device=device).manual_seed(torch.seed())
        images = self.text2img_pipeline(prompt, num_inference_steps=50, generator=generator)["images"]
        image_path = f'/tmp/{uuid.uuid4()}.png'
        images[0].save(image_path)
        return image_path

# 프롬프트 번역 함수
def translate_prompt(prompt):
    logging.debug(f'Translating prompt: {prompt}')
    translation = translator(prompt, max_length=512)
    translated_text = translation[0]['translation_text']
    logging.debug(f'Translated text: {translated_text}')
    return translated_text

# 디스코드 토큰 및 봇 실행
if __name__ == "__main__":
    discord_token = os.getenv('DISCORD_TOKEN')
    discord_client = MyClient(intents=intents)
    discord_client.run(discord_token)