# img_bot.py
import discord, os, io, asyncio, logging, requests, replicate

# ── 필수 환경 변수 ────────────────────────────────────────
TOKEN       = os.getenv("DISCORD_TOKEN")                 # 디스코드 봇 토큰
CHANNEL_ID  = int(os.getenv("DISCORD_CHANNEL_ID"))       # 듣고 응답할 채널
REPL_TOKEN  = (os.getenv("OPENAI_API_KEY") or "").strip()# Replicate 토큰 (같은 변수 사용)
# ─────────────────────────────────────────────────────────

MODEL = (
    "bytedance/sdxl-lightning-4step:"
    "6f7a773af6fc3e8de9d5a3c00be77c17308914bf67772726aff83496ba1e3bbe"
)

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s [%(levelname)s] %(message)s")

intents = discord.Intents.default()
intents.message_content = True

class ImageBot(discord.Client):
    async def on_ready(self):
        logging.info(f"Logged in as {self.user} (id={self.user.id})")

    async def on_message(self, message: discord.Message):
        # 봇 자신·다른 채널 메시지 무시
        if message.author.id == self.user.id or message.channel.id != CHANNEL_ID:
            return

        prompt = message.content.strip()
        if not prompt:
            return

        await message.channel.typing()

        # Replicate 호출(기본값: prompt만 전달)
        def run_replicate():
            return list(replicate.run(MODEL, input={"prompt": prompt}))

        try:
            loop   = asyncio.get_running_loop()
            images = await loop.run_in_executor(None, run_replicate)
        except Exception as e:
            logging.error(f"Replicate error: {e}")
            await message.reply("⚠️ 이미지 생성 실패!")
            return

        files = []
        for idx, item in enumerate(images):
            try:
                data = item.read() if hasattr(item, "read") else requests.get(item).content
                files.append(discord.File(io.BytesIO(data), filename=f"img_{idx}.png"))
            except Exception as e:
                logging.warning(f"IMG {idx} 처리 실패: {e}")

        await message.reply(files=files if files else None,
                            content=None if files else "⚠️ 이미지를 전송할 수 없습니다.")

if __name__ == "__main__":
    replicate.Client(api_token=REPL_TOKEN)  # OPENAI_API_KEY 그대로 사용
    ImageBot(intents=intents).run(TOKEN)