import gradio as gr
from PIL import Image, ImageEnhance, ImageFilter
import io
import zipfile

def batch_upscale(images, scale, enhance):
    if not images:
        return None, None, "画像をアップロードしてください"
    
    processed_images = []
    results_info = []
    
    for i, img in enumerate(images):
        # 高品質アップスケール
        width, height = img.size
        new_width = int(width * scale)
        new_height = int(height * scale)
        
        upscaled = img.resize((new_width, new_height), Image.LANCZOS)
        
        # 追加補正
        if enhance:
            upscaled = upscaled.filter(ImageFilter.UnsharpMask(radius=2, percent=150, threshold=3))
            enhancer = ImageEnhance.Contrast(upscaled)
            upscaled = enhancer.enhance(1.15)
            enhancer = ImageEnhance.Sharpness(upscaled)
            upscaled = enhancer.enhance(1.35)
        
        processed_images.append(upscaled)
        
        # サイズ情報
        buf = io.BytesIO()
        upscaled.save(buf, format="PNG", quality=95)
        size_kb = len(buf.getvalue()) / 1024
        results_info.append(f"画像{i+1}: {size_kb:.1f} KB")
    
    # ZIPファイル作成
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zip_file:
        for i, img in enumerate(processed_images):
            img_byte = io.BytesIO()
            img.save(img_byte, format="PNG", quality=95)
            zip_file.writestr(f"upscaled_{i+1}.png", img_byte.getvalue())
    
    return processed_images, zip_buffer.getvalue(), "\n".join(results_info)

with gr.Blocks(title="Kazの高画質化アプリ") as demo:
    gr.Markdown("# 🖼️ KazのAI画像高画質化アプリ\n**複数画像一括処理版**")
    
    with gr.Row():
        with gr.Column():
            input_gallery = gr.Gallery(
                label="📸 複数の画像をドラッグ&ドロップ（一括処理）",
                height=500,
                object_fit="contain"
            )
            scale = gr.Radio([2, 4, 8], value=4, label="🔍 拡大倍率（8xは時間がかかります）")
            enhance = gr.Checkbox(label="シャープ・コントラスト強化", value=True)
            btn = gr.Button("🚀 全画像を高画質化", variant="primary", size="large")
        
        with gr.Column():
            output_gallery = gr.Gallery(label="✨ 処理結果", height=500)
            output_zip = gr.File(label="📦 全画像をZIPでダウンロード")
    
    status = gr.Textbox(label="処理結果", lines=5)
    
    btn.click(
        fn=batch_upscale,
        inputs=[input_gallery, scale, enhance],
        outputs=[output_gallery, output_zip, status]
    )

    gr.Markdown("### Tips\n・複数画像を同時にアップロードできます\n・8x + 強化ONでかなり高画質になります\n・処理後、ZIPをダウンロードしてください")

demo.launch()