import os
import base64
from io import BytesIO
from PIL import Image
import gradio as gr
from openai import OpenAI

def get_client():
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise ValueError("找不到 HF_TOKEN 環境變數，請確認是否已在 Space Settings 中設定 Secrets。")
    return OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=token,
    )

def compress_image_to_base64(image_path, max_size=1024, quality=85):
    """將圖片縮放並壓縮轉為 Base64，避免 413 Payload Too Large 錯誤"""
    img = Image.open(image_path)
    
    # 轉換為 RGB（避免 PNG 的 RGBA 模式在存成 JPEG 時出錯）
    if img.mode in ("RGBA", "P"):
        img = img.convert("RGB")
        
    # 等比例縮放長寬，最長邊不超過 max_size
    img.thumbnail((max_size, max_size))
    
    buffered = BytesIO()
    img.save(buffered, format="JPEG", quality=quality)
    encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{encoded_string}"

def analyze_image(image, prompt):
    if image is None:
        return "請先上傳一張圖片！"
    
    if not prompt or prompt.strip() == "":
        prompt = "Describe this image in one sentence."

    try:
        # 進行壓縮轉檔
        base64_image = compress_image_to_base64(image)
        
        client = get_client()
        completion = client.chat.completions.create(
            model="google/gemma-4-31B-it:novita",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {"url": base64_image}
                        }
                    ]
                }
            ],
            max_tokens=300
        )
        return completion.choices[0].message.content
    except Exception as e:
        return f"發生錯誤：{str(e)}"

demo = gr.Interface(
    fn=analyze_image,
    inputs=[
        gr.Image(type="filepath", label="上傳圖片"),
        gr.Textbox(label="輸入提示詞 (Prompt)", value="Describe this image in one sentence in chinese")
    ],
    outputs=gr.Textbox(label="模型回應"),
    title="Gemma 多模態影像分析助手",
    description="透過 Hugging Face Router 呼叫 gemma-4-31B-it 模型進行影像辨識與分析。"
)

if __name__ == "__main__":
    demo.launch()