import os
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from transformers import pipeline

# 1) 이미지 캡셔닝 파이프라인
captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")

# 2) 밈 + 자막 합성 함수
def make_meme_with_subtitle(image, top_text, bottom_text):
    if image is None:
        raise ValueError("이미지가 업로드되지 않았습니다.")
    
    # 이미지 변환
    img = image.convert("RGB")
    w, h = img.size

    # ① 자동 캡션 생성
    try:
        out = captioner(img, max_length=40, num_beams=5)
        raw = out[0].get("generated_text") or out[0].get("text") or ""
        subtitle = raw.strip().capitalize() if raw else "No caption found"
    except Exception as e:
        subtitle = "No caption (error)"

    # ② 자막용 검은 띠 생성
    bar_h = int(h * 0.15)
    bar = Image.new("RGB", (w, bar_h), color="black")
    combined_h = h + bar_h
    combined = Image.new("RGB", (w, combined_h))
    combined.paste(img, (0, 0))
    combined.paste(bar, (0, h))

    draw = ImageDraw.Draw(combined)

    # ③ 자막 폰트 설정
    try:
        font_path = "arialbd.ttf"  # 시스템에 해당 폰트가 없을 경우 오류
        base_font_size = int(bar_h * 0.5)
        font = ImageFont.truetype(font_path, size=base_font_size)
    except:
        font = ImageFont.load_default()

    # ④ 자막 그리기 (하단 중앙)
    tw, th = draw.textsize(subtitle, font=font)
    tx = (w - tw) // 2
    ty = h + (bar_h - th) // 2
    draw.text((tx, ty), subtitle, font=font, fill="white")

    # ⑤ 밈 텍스트 폰트 설정
    meme_font_size = int(h * 0.07)
    try:
        meme_font = ImageFont.truetype(font_path, size=meme_font_size)
    except:
        meme_font = ImageFont.load_default()

    # 상단 텍스트
    if top_text:
        text = top_text.upper()
        tw, th = draw.textsize(text, font=meme_font)
        draw.text(((w - tw) // 2, 10), text, font=meme_font, fill="white", stroke_width=2, stroke_fill="black")

    # 하단 텍스트 (이미지 내부)
    if bottom_text:
        text = bottom_text.upper()
        tw, th = draw.textsize(text, font=meme_font)
        draw.text(((w - tw) // 2, h - th - 10), text, font=meme_font, fill="white", stroke_width=2, stroke_fill="black")

    return combined

# 3) Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📸 밈 생성 + 자동 영어 자막 합성")
    gr.Markdown("이미지를 업로드하고, 상단·하단 밈 텍스트를 입력하면 자동 생성된 영어 자막과 함께 밈 스타일 이미지를 반환합니다.")

    img_in = gr.Image(type="pil", label="Upload Image")  # filepath → pil
    top_txt = gr.Textbox(label="Top Text (optional)", placeholder="e.g. WHEN YOU REALIZE...", lines=1)
    bottom_txt = gr.Textbox(label="Bottom Text (optional)", placeholder="e.g. ...IT'S MONDAY AGAIN", lines=1)
    btn = gr.Button("Generate Meme")
    out_img = gr.Image(label="Meme with Subtitle")
    
    btn.click(fn=make_meme_with_subtitle, inputs=[img_in, top_txt, bottom_txt], outputs=out_img)

# 4) 앱 실행
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))