image_generator.py

import os
import sys
import shutil
import uuid
import random
import logging
from PIL import Image, ImageDraw, ImageFont

# 設定日誌記錄，便於開發與除錯
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# ==============================================================================
# 改為從環境變數讀取多組金鑰，以支援自動備援與驗證機制
# ==============================================================================
GEMINI_API_KEYS = [
    os.environ.get("GOOGLE_API_KEY_1"),
    os.environ.get("GOOGLE_API_KEY_2"),
    os.environ.get("GOOGLE_API_KEY")
]
GEMINI_API_KEYS = list(dict.fromkeys(filter(None, GEMINI_API_KEYS)))

HF_TOKENS = [
    os.environ.get("HF_TOKEN_1"),
    os.environ.get("HF_TOKEN_2"),
    os.environ.get("HF_TOKEN")
]
HF_TOKENS = list(dict.fromkeys(filter(None, HF_TOKENS)))

_working_gemini_key = None
_working_hf_token = None

def check_gemini_key(api_key):
    if not api_key:
        return False
    try:
        import google.generativeai as genai
        genai.configure(api_key=api_key)
        try:
            for _ in genai.list_models():
                break
            return True
        except Exception:
            # 備份快速檢測機制
            model = genai.GenerativeModel('gemini-2.5-flash')
            model.generate_content("ping", generation_config={"max_output_tokens": 1})
            return True
    except Exception as e:
        logger.warning(f"Gemini API key 驗證失敗 (...{api_key[-4:] if len(api_key) > 4 else ''}): {e}")
        return False

def check_hf_token(token):
    if not token:
        return False
    try:
        import requests
        headers = {"Authorization": f"Bearer {token}"}
        r = requests.get("https://huggingface.co/api/whoami-v2", headers=headers, timeout=5)
        return r.status_code == 200
    except Exception as e:
        logger.warning(f"HF Token 驗證失敗 (...{token[-4:] if len(token) > 4 else ''}): {e}")
        return False

def get_working_gemini_key():
    global _working_gemini_key
    if _working_gemini_key and check_gemini_key(_working_gemini_key):
        return _working_gemini_key
    
    for key in GEMINI_API_KEYS:
        if check_gemini_key(key):
            _working_gemini_key = key
            logger.info(f"成功啟用並快取 Gemini 金鑰: ...{key[-4:] if len(key) > 4 else ''}")
            return key
            
    fallback = GEMINI_API_KEYS[0] if GEMINI_API_KEYS else None
    if fallback:
        logger.warning("所有 Gemini 金鑰皆驗證失敗，回退至第一個金鑰嘗試")
    return fallback

def get_working_hf_token():
    global _working_hf_token
    if _working_hf_token and check_hf_token(_working_hf_token):
        return _working_hf_token
    
    for token in HF_TOKENS:
        if check_hf_token(token):
            _working_hf_token = token
            logger.info(f"成功啟用並快取 HF Token: ...{token[-4:] if len(token) > 4 else ''}")
            return token
            
    fallback = HF_TOKENS[0] if HF_TOKENS else None
    if fallback:
        logger.warning("所有 HF Tokens 皆驗證失敗，回退至第一個 Token 嘗試")
    return fallback

# 固定的生圖寬度與高度 (1024x1024)
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024

# 預設的俗俗早安圖問候語，當 Gemini API 金鑰未提供時作為 Fallback 備用
PRESET_GREETINGS = [
    "早安！今天也要假裝系統正常！",
    "早安！人生卡住時，請先重開機！",
    "早安！努力不一定成功，但會很累！",
    "早安！腦袋尚未開機，請多包涵！",
    "早安！今天的你，比昨天更像待機中！",
    "早安！不要放棄，反正也沒人接住你！",
    "早安！上課像登入失敗，但人要有光！",
    "早安！作業不會寫，太陽還是會升起！",
    "早安！群組安靜不是冷漠，是大家都沒電！",
    "早安！願你今天的 Wi-Fi 比人生穩定！"
]

# 頂部大字 Banner 的預設選用標題
PRESET_HEADERS = [
    "早安",
    "平安喜樂",
    "吉祥如意",
    "福氣滿滿",
    "吉祥安康",
    "順心如意",
    "今日開機",
    "福氣載入",
    "平安上線",
    "人生重整",
    "群組復活",
    "努力省電",
    "腦袋離線",
    "系統正常",
    "快樂重啟"
]

# ==============================================================================
# 怪圖風格素材池：讓每次早安圖更有「好怪，再看一眼」的變化
# ==============================================================================

WEIRD_STYLE_PRESETS = [
    "a low-budget 1980s Taiwanese TV shopping commercial set with fake luxury props and overly dramatic lighting",
    "a surreal plastic toy diorama where tiny objects become huge and huge objects become tiny",
    "a chaotic internet meme collage with awkward cheerful energy and impossible scale",
    "a retro Windows 95 clip-art poster mixed with Taiwanese neon sign aesthetics",
    "an overdecorated karaoke music video background with fog machine, rainbow spotlights, and fake glamour",
    "a fake inspirational corporate poster gone wrong, treating a meaningless action like a world-saving mission",
    "a children's craft project made by an overconfident robot, with stickers, glitter, paper cutouts, and messy charm",
    "a surreal supermarket advertisement where the subject is displayed like a sacred limited-edition product",
    "a bizarre morning news broadcast set reporting a completely unimportant event as breaking news",
    "a dreamlike school bulletin board poster with too many decorations and zero design logic"
]

WEIRD_VISUAL_GAGS = [
    "the subject is posing heroically while doing something completely useless",
    "the subject is being celebrated by tiny objects like a world-saving hero",
    "the subject is surrounded by floating breakfast items as if they are planets",
    "the subject is receiving an official award for a meaningless achievement",
    "the subject is dramatically illuminated while failing at a very simple task",
    "the subject is sitting on a plastic throne made of random everyday objects",
    "the subject appears in a breaking-news scene about a tiny inconvenience",
    "the subject is leading a ceremonial parade for no clear reason",
    "the subject is treated like a legendary mascot of emotional damage",
    "the subject looks extremely confident despite clearly misunderstanding the situation"
]

WEIRD_COMPOSITIONS = [
    "fisheye lens perspective, slightly too close to the subject, funny and uncomfortable but not scary",
    "symmetrical poster composition with the subject placed too seriously in the center",
    "busy collage composition with many strange but harmless background details",
    "dramatic low-angle hero shot for a completely ridiculous subject",
    "fake advertisement layout with exaggerated product-poster energy",
    "crowded scene with tiny background characters reacting in confusion",
    "overly majestic composition with sparkles, fog, and unnecessary visual importance",
    "awkward family-photo composition, but everyone is an object or mascot"
]

ELDER_CARD_ELEMENTS = [
    "dramatic golden sunrise",
    "glowing clouds",
    "oversized lotus flowers",
    "sparkles",
    "rainbow gradients",
    "shiny flowers",
    "neon blessing-card borders",
    "floating hearts",
    "cheap clip-art stars",
    "butterflies",
    "glitter effects",
    "overly saturated colorful background"
]

COLOR_MOODS = [
    "extremely oversaturated rainbow colors",
    "neon pink and gold morning glow",
    "cheap fluorescent green and purple color palette",
    "glossy candy-like colors",
    "overexposed golden sunlight with chaotic colorful shadows",
    "unnecessarily cheerful pastel colors with dramatic contrast"
]

def build_weird_recipe():
    elder_elements = random.sample(ELDER_CARD_ELEMENTS, k=random.randint(3, 5))
    return {
        "style": random.choice(WEIRD_STYLE_PRESETS),
        "gag": random.choice(WEIRD_VISUAL_GAGS),
        "composition": random.choice(WEIRD_COMPOSITIONS),
        "elder_elements": ", ".join(elder_elements),
        "color_mood": random.choice(COLOR_MOODS)
    }

# 經典長輩圖必備霓虹彩虹漸層字體配色系列
RAINBOW_COLORS = [
    "#FF0033",  # 霓虹紅
    "#FF6600",  # 霓虹橘
    "#FFFF00",  # 鮮艷黃
    "#33FF33",  # 螢光綠
    "#00FFFF",  # 霓虹藍/青色
    "#CC33FF",  # 亮紫色
    "#FF3399"   # 桃紅色
]

def translate_to_english(text):
    try:
        import requests
        url = f"https://translate.googleapis.com/translate_a/single?client=gtx&sl=auto&tl=en&dt=t&q={requests.utils.quote(text)}"
        r = requests.get(url, timeout=5)
        if r.status_code == 200:
            translated = r.json()[0][0][0]
            logger.info(f"自動翻譯成功: '{text}' -> '{translated}'")
            return translated
    except Exception as e:
        logger.error(f"翻譯 API 呼叫失敗: {e}")
    return text

def get_chinese_font(font_size, font_name="kaiu"):
    windir = os.environ.get('WINDIR', 'C:\\Windows')
    paths = []
    
    paths.extend([
        "kaiu.ttf",                  
        "static/kaiu.ttf",           
        os.path.join(windir, 'Fonts', 'kaiu.ttf'),    
        os.path.join(windir, 'Fonts', 'msjhbd.ttc'),  
        os.path.join(windir, 'Fonts', 'msjh.ttc'),    
        os.path.join(windir, 'Fonts', 'simsun.ttc'),  
        "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc",
        "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
        "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
        "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
        "/usr/share/fonts/truetype/droid/DroidSansFallback.ttf"
    ])
    
    for path in paths:
        if os.path.exists(path):
            try:
                if path.endswith('.ttc'):
                    logger.info(f"成功載入 TTC 字型: {path}")
                    return ImageFont.truetype(path, font_size, index=0)
                else:
                    logger.info(f"成功載入 TTF 字型: {path}")
                    return ImageFont.truetype(path, font_size)
            except Exception as e:
                logger.warning(f"載入字型 {path} 失敗: {e}")
                continue
                
    logger.warning("未偵測到任何系統字型，將使用 Pillow 預設字型。")
    try:
        return ImageFont.load_default()
    except Exception:
        return None

def wrap_text(text, font, max_width, draw):
    lines = []
    current_line = ""
    for char in text:
        test_line = current_line + char
        w = draw.textlength(test_line, font=font)
        if w <= max_width:
            current_line = test_line
        else:
            if current_line:
                lines.append(current_line)
            current_line = char
    if current_line:
        lines.append(current_line)
    return lines

def generate_text_with_gemini(api_key, query):
    if not api_key:
        logger.info("未偵測到 Gemini API Key。將使用內建隨機賀詞。")
        return random.choice(PRESET_GREETINGS)
        
    try:
        import google.generativeai as genai
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.5-flash')
        
        prompt = (
             f"請針對「{query}」寫一句繁體中文早安圖文字。\n"
             f"風格要求：台灣長輩圖、荒謬、38、假勵志、欠吐槽、適合群組聊天。\n"
             f"語氣要像一個很有自信但其實很荒謬的早安機器人。\n"
             f"字數限制：12 到 22 字。\n"
             f"內容方向：可以嘲諷抽象人生、早八、上班、作業、拖延、群組沉默、機器人自己。\n"
             f"安全限制：不可攻擊特定真人、族群、性別、外貌、疾病、政治立場或宗教。\n"
             f"不要髒話，不要人身攻擊，不要太黑暗。\n"
             f"要讓群組成員看了想回：『這什麼鬼啦』。\n"
             f"範例：\n"
             f"- 早安，今天也要假裝系統正常\n"
             f"- 早安，努力不一定成功但很耗電\n"
             f"- 早安，人生卡住請重新整理\n"
             f"- 早安，腦袋離線也是一種生活\n"
             f"- 早安，群組沒人講話我先亂講\n"
             f"請只回傳一句中文早安文字，不要解釋，不要加引號。"
        )
        response = model.generate_content(prompt)
        text = response.text.strip().replace('"', '').replace('「', '').replace('」', '')
        logger.info(f"Gemini 生成的長輩賀詞: {text}")
        return text
    except Exception as e:
        logger.error(f"Gemini 賀詞生成錯誤: {e}")
        return random.choice(PRESET_GREETINGS)

def expand_prompt_with_gemini(api_key, query):
    query_en = translate_to_english(query)
    weird_recipe = build_weird_recipe()
    
    default_expanded = (
        f"A funny, absurd, and strangely memorable Taiwanese good morning meme image featuring {query_en} as the main subject. "
        f"The image should create a 'this is so weird that I need to look again' feeling, while remaining harmless and funny. "
        f"Weird visual style: {weird_recipe['style']}. "
        f"Visual gag: {weird_recipe['gag']}. "
        f"Composition: {weird_recipe['composition']}. "
        f"Taiwanese elder-card elements to include: {weird_recipe['elder_elements']}. "
        f"Color mood: {weird_recipe['color_mood']}. "
        f"The subject must be clearly visible and related to {query_en}, but the scene should feel surreal, ridiculous, overly confident, and awkwardly cheerful. "
        f"Do not always use the same lotus-and-sunrise layout. Make the scene visually different each time. "
        f"No readable text inside the image, because Traditional Chinese text will be added later by the program. "
        f"No real public figures, no political symbols, no hate symbols, no violence, no sexual content, no horror, no offensive stereotypes. "
        f"High resolution, colorful, detailed, humorous, absurd, meme-like, strange but not scary."
    )
    
    if not api_key:
        logger.info("未偵測到 Gemini API Key。使用預設怪圖封裝 Prompt。")
        return default_expanded
        
    try:
        import google.generativeai as genai
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.5-flash')
        
        prompt = (
            f"Please translate the Chinese phrase '{query}' into English, "
            f"then expand it into a detailed image generation prompt for a funny, absurd Taiwanese good morning meme image.\n\n"

            f"Goal:\n"
            f"Make the image feel strange, ridiculous, and visually memorable, like something so weird that group chat members would look twice and want to roast it.\n\n"

            f"Random weirdness recipe for this image:\n"
            f"- Main subject: {query_en}\n"
            f"- Weird visual style: {weird_recipe['style']}\n"
            f"- Visual gag: {weird_recipe['gag']}\n"
            f"- Composition: {weird_recipe['composition']}\n"
            f"- Elder-card elements: {weird_recipe['elder_elements']}\n"
            f"- Color mood: {weird_recipe['color_mood']}\n\n"

            f"Requirements:\n"
            f"1. The main subject must clearly be related to: '{query_en}'.\n"
            f"2. The image must look different from a normal good morning card. Avoid repeating the same lotus-sunrise-flower layout every time.\n"
            f"3. Combine Taiwanese elder-style good morning card aesthetics with absurd internet meme humor.\n"
            f"4. Make the scene strange but harmless: funny, awkwardly cheerful, overly serious, and visually ridiculous.\n"
            f"5. Include one clear visual gag where the subject looks extremely confident while doing something meaningless, useless, or misunderstood.\n"
            f"6. Keep it colorful, detailed, oversaturated, flamboyant, and easy for group chat members to roast.\n"
            f"7. Do not include real public figures, political symbols, hate symbols, violence, sexual content, horror, or offensive stereotypes.\n"
            f"8. Do not generate readable text inside the image, because Traditional Chinese text will be added later by the program.\n\n"

            f"Only return the final English image prompt. Do not include explanations, labels, or quotation marks."
        )

        response = model.generate_content(prompt)
        expanded = response.text.strip().replace('"', '').replace('「', '').replace('」', '')
        logger.info(f"Gemini 擴寫的怪圖英文生圖提示詞: {expanded}")
        return expanded

    except Exception as e:
        logger.error(f"Gemini 提示詞擴寫錯誤: {e}")
        return default_expanded        


def draw_rainbow_text(draw, text, x_center, y_start, font, stroke_width=6, stroke_fill="#000000"):
    if not font:
        return y_start
        
    total_w = 0
    char_widths = []
    for char in text:
        w = draw.textlength(char, font=font)
        char_widths.append(w)
        total_w += w
        
    start_x = x_center - (total_w / 2)
    current_x = start_x
    
    for idx, char in enumerate(text):
        color = RAINBOW_COLORS[idx % len(RAINBOW_COLORS)]
        
        draw.text(
            (current_x, y_start),
            char,
            font=font,
            fill=color,
            stroke_width=stroke_width,
            stroke_fill=stroke_fill
        )
        current_x += char_widths[idx]
        
    bbox = font.getbbox("早")
    char_h = bbox[3] - bbox[1] if bbox else 50
    return y_start + char_h + 10

def overlay_text_on_image(image_path, greeting_text, output_path):
    try:
        img = Image.open(image_path).convert("RGB")
        if img.size != (DEFAULT_WIDTH, DEFAULT_HEIGHT):
            img = img.resize((DEFAULT_WIDTH, DEFAULT_HEIGHT), Image.Resampling.LANCZOS)
            
        draw = ImageDraw.Draw(img)
        
        header_font = get_chinese_font(120, "kaiu")
        body_font = get_chinese_font(52, "kaiu")
        footer_font = get_chinese_font(32, "kaiu")
        
        header_text = f"🌸 {random.choice(PRESET_HEADERS)} 🌸"
        header_y = 60
        if header_font:
            draw.text(
                (DEFAULT_WIDTH / 2, header_y),
                header_text,
                font=header_font,
                fill="#FFFF00",
                anchor="ma",
                stroke_width=8,
                stroke_fill="#000000"
            )
            
        body_y_start = 740
        if body_font:
            lines = wrap_text(greeting_text, body_font, 920, draw)
            current_y = body_y_start
            for line in lines:
                current_y = draw_rainbow_text(
                    draw,
                    line,
                    x_center=DEFAULT_WIDTH / 2,
                    y_start=current_y,
                    font=body_font,
                    stroke_width=6,
                    stroke_fill="#000000"
                )
                
        footer_text = random.choice([
            "認同請分享，不認同也沒差！",
            "看到請回，群組需要復活！",
            "今日能量不足，請互相充電！",
            "早安成功，腦袋另行通知！",
            "分享此圖，福氣自動重開機！",
            "本圖無用，但非常努力！"
        ])
        footer_y = 930
        if footer_font:
            draw.text(
                (DEFAULT_WIDTH / 2, footer_y),
                footer_text,
                font=footer_font,
                fill="#33FF33",
                anchor="ma",
                stroke_width=5,
                stroke_fill="#000000"
            )
            
        img.save(output_path, "JPEG", quality=90)
        logger.info(f"早安圖加工疊加完成，存檔至: {output_path}")
        return True
    except Exception as e:
        logger.error(f"Pillow 文字疊加失敗: {e}")
        try:
            shutil.copy(image_path, output_path)
            return True
        except Exception:
            return False

def generate_good_morning_image(user_input, api_key=None, custom_greeting=None):
    resolved_api_key = api_key or get_working_gemini_key()

    query = user_input.strip()
    if query.startswith("@create image"):
        query = query.replace("@create image", "", 1).strip()
    if not query:
        query = "蓮花"
        
    logger.info(f"開始處理早安圖主題任務: {query}")
    
    if custom_greeting:
        greeting_text = custom_greeting
    else:
        greeting_text = generate_text_with_gemini(resolved_api_key, query)
        
    expanded_prompt = expand_prompt_with_gemini(resolved_api_key, query)
    
    logger.info("呼叫 Gradio Space 端點 (stabilityai/stable-diffusion-3.5-large)...")
    try:
        from gradio_client import Client
        token = get_working_hf_token()
        client = Client("stabilityai/stable-diffusion-3.5-large", token=token)
        
        result = client.predict(
            prompt=expanded_prompt,
            negative_prompt="blurry, low quality, bad anatomy, deformed, ugly, disfigured, poor lighting, out of frame, cropped",
            seed=0,
            randomize_seed=True,
            width=DEFAULT_WIDTH,
            height=DEFAULT_HEIGHT,
            guidance_scale=4.5,
            num_inference_steps=40,
            api_name="/infer",
        )
        
        if result and isinstance(result, tuple) and len(result) > 0:
            temp_image_path = result[0]
        elif isinstance(result, str):
            temp_image_path = result
        else:
            raise ValueError(f"無法識別的 Gradio API 返回格式: {result}")
            
        logger.info(f"SD3.5 基礎底圖繪製成功，暫存於: {temp_image_path}")
        
        data_dir = os.path.abspath("data")
        os.makedirs(data_dir, exist_ok=True)
        
        image_id = str(uuid.uuid4())
        final_image_name = f"{image_id}.jpg"
        final_image_path = os.path.join(data_dir, final_image_name)
        
        success = overlay_text_on_image(temp_image_path, greeting_text, final_image_path)
        if not success:
            raise RuntimeError("Pillow 後製疊加文字發生崩潰失敗")
            
        return {
            "success": True,
            "id": image_id,
            "query": query,
            "greeting_text": greeting_text,
            "expanded_prompt": expanded_prompt,
            "image_path": final_image_path,
            "filename": final_image_name
        }
        
    except Exception as e:
        logger.error(f"整條生圖流水線執行失敗: {e}")
        return {
            "success": False,
            "error": str(e)
        }