# img_bot.py
# “이미지 첨부 → 실패 시 URL 출력” 완결판

import os, io, re, random, asyncio, logging, subprocess, base64
from urllib.parse import urljoin, quote_plus

import discord
import requests
import replicate
from transformers import pipeline as transformers_pipeline
from gradio_client import Client

try:
    from PIL import Image                # WEBP → PNG 변환
    PIL_OK = True
except Exception:
    PIL_OK = False                       # Pillow 미설치 시 변환 생략

# ────────────────── 환경 변수 ──────────────────
TOKEN       = os.getenv("DISCORD_TOKEN")
CHANNEL_ID  = int(os.getenv("DISCORD_CHANNEL_ID"))
REPL_TOKEN  = (os.getenv("OPENAI_API_KEY") or "").strip()
HF_TOKEN    = (os.getenv("HF_TOKEN") or "").strip()

if not TOKEN or not CHANNEL_ID:
    raise RuntimeError("DISCORD_TOKEN 과 DISCORD_CHANNEL_ID 환경 변수를 모두 지정하세요.")
if not REPL_TOKEN:
    raise RuntimeError("OPENAI_API_KEY 에 Replicate Personal Access Token 값을 넣어주세요.")

os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN          # (Replicate 미사용)

# ────────────────── Gradio 서버 ─────────────────
GRADIO_URL = "http://211.233.58.201:7971"
GRADIO_API = "/process_and_save_image"

# ────────────────── 번역 파이프라인 ─────────────
translator = transformers_pipeline(
    "translation",
    model="Helsinki-NLP/opus-mt-ko-en",
    device=-1,
    **({"token": HF_TOKEN} if HF_TOKEN else {})
)

async def ko2en_async(text: str) -> str:
    """한글 포함 시 영어 번역."""
    if not re.search(r"[가-힣]", text):
        return text
    loop = asyncio.get_running_loop()
    try:
        return await loop.run_in_executor(
            None,
            lambda: translator(text, max_length=256, num_beams=1)[0]["translation_text"].strip()
        )
    except Exception as e:
        logging.warning(f"번역 실패, 원문 사용: {e}")
        return text

# ────────────────── 로깅 ────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler()]
)

# ────────────────── Discord 클라이언트 ──────────
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})")
        try:
            subprocess.Popen(["python", "web.py"])
            logging.info("web.py server has been started.")
        except Exception as e:
            logging.warning(f"web.py 실행 실패: {e}")

    async def on_message(self, message: discord.Message):
        if message.author.id == self.user.id or message.channel.id != CHANNEL_ID:
            return

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

        prompt_en = await ko2en_async(prompt_raw)
        await message.channel.typing()

        # ───────────── Gradio 호출 ──────────────
        def generate_image():
            client = Client(GRADIO_URL)
            return client.predict(
                height=768,
                width=768,
                steps=30,
                scales=3.5,
                prompt=prompt_en,
                seed=random.randint(0, 2**32 - 1),
                api_name=GRADIO_API
            )

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

        # 리스트 결과 → 첫 요소
        if isinstance(result, list):
            result = result[0]

        # ───────────── 이미지 데이터 확보 ─────────────
        data = None
        remote_path = None               # 실패 시 사용자에게 보여줄 URL
        candidate_urls: list[str] = []

        def add_remote(p: str):
            url = urljoin(GRADIO_URL, f"/gradio_api/file={quote_plus(p)}")
            candidate_urls.append(url)
            return url

        # dict 결과
        if isinstance(result, dict):
            u = result.get("url")
            if u:
                if u.startswith("/"):
                    u = urljoin(GRADIO_URL, u)
                candidate_urls.append(u)
                remote_path = u
            p = result.get("path")
            if p:
                remote_path = add_remote(p)
                if os.path.isfile(p):
                    try:
                        with open(p, "rb") as f:
                            data = f.read()
                    except Exception:
                        pass

        # str 결과
        elif isinstance(result, str):
            s = result.strip()
            if s.startswith("data:"):
                try:
                    data = base64.b64decode(re.sub(r"^data:image/[^;]+;base64,", "", s))
                except Exception as e:
                    logging.warning(f"base64 디코딩 실패: {e}")
            elif s.startswith("http"):
                candidate_urls.append(s)
                remote_path = s
            else:
                remote_path = add_remote(s)
                if s.startswith("/"):
                    candidate_urls.append(urljoin(GRADIO_URL, s))
                if os.path.isfile(s):
                    try:
                        with open(s, "rb") as f:
                            data = f.read()
                    except Exception:
                        pass

        # URL 다운로드 (헤더 무시, 200 OK 만 확인)
        for u in candidate_urls:
            if data:
                break
            try:
                r = requests.get(u, timeout=10)
                if r.ok and r.content:
                    data = r.content
                    break
            except Exception as e:
                logging.warning(f"URL 다운로드 실패({u}): {e}")

        # ───────────── Discord 전송 ─────────────
        if data:
            if PIL_OK:  # WEBP → PNG 변환
                try:
                    buf = io.BytesIO()
                    Image.open(io.BytesIO(data)).convert("RGB").save(buf, format="PNG")
                    buf.seek(0)
                    await message.reply(files=[discord.File(buf, filename="image.png")])
                    return
                except Exception as e:
                    logging.warning(f"PNG 변환 실패: {e}")
            # 변환 실패 또는 Pillow 없음 → 원본 전송
            await message.reply(files=[discord.File(io.BytesIO(data), filename="image.webp")])
            return

        # ───────────── 실패 시 URL 출력 ─────────────
        if remote_path:
            await message.reply(content=remote_path)
        else:
            await message.reply("⚠️ 이미지를 전송할 수 없습니다.")

# ────────────────── 실행 ────────────────────────
if __name__ == "__main__":
    replicate.Client(api_token=REPL_TOKEN)   # 구조 유지
    ImageBot(intents=intents).run(TOKEN)