import gradio as gr
from PIL import Image, ImageOps
import os
import zipfile


def save_image(image, path):
    # 画像の保存 (PNG形式で保存)
    image.save(path, format='PNG')

def process_image(image):
    # 元の画像のサイズ
    original_width, original_height = image.size
    
    # 新しいサイズ
    target_width, target_height = 252, 144
    
    # 分割後の画像を格納するリスト
    outputs = []
    
    # 画像を4つのセクションに分割するための新しいサイズ
    new_width = original_width // 2
    new_height = original_height // 2
    
    # 画像を4つのセクションに分割
    for i in range(2):
        for j in range(2):
            left = new_width * j
            upper = new_height * i
            right = new_width * (j + 1)
            lower = new_height * (i + 1)
            
            # 画像を切り取る
            cropped_image = image.crop((left, upper, right, lower))
            
            # アスペクト比を保ちながら、リサイズ
            cropped_image.thumbnail((target_width, target_height), Image.Resampling.LANCZOS)
            
            # 新しい黒背景を作成
            black_background = Image.new('RGBA', (target_width, target_height), (0, 0, 0, 255))
            
            # 余白の追加を左上、左下、右上、右下に対応
            if i == 0 and j == 0:  # 左上
                left_offset = (target_width - cropped_image.width)  # 右側に余白
                top_offset = (target_height - cropped_image.height)  # 下側に余白
            elif i == 1 and j == 0:  # 左下
                left_offset = (target_width - cropped_image.width)  # 右側に余白
                top_offset = 0  # 上側に余白
            elif i == 0 and j == 1:  # 右上
                left_offset = 0  # 左側に余白
                top_offset = (target_height - cropped_image.height)  # 下側に余白
            elif i == 1 and j == 1:  # 右下
                left_offset = 0  # 左側に余白
                top_offset = 0  # 上側に余白

            # 中心に元画像を貼り付け
            black_background.paste(cropped_image, (left_offset, top_offset))
            
            outputs.append(black_background)
    
    # 出力される4つの画像を返す
    return outputs[0], outputs[1], outputs[2], outputs[3]

def zip_files(zip_files, zip_path):
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file_path in zip_files:
            zipf.write(file_path, arcname=os.path.basename(file_path))

class webui:
    def __init__(self):
        self.demo = gr.Blocks()

    def main(self, image_path):
        image = Image.open(image_path)
        # 拡張子を取り除いたファイル名を取得
        image_name = os.path.splitext(image_path)[0]

        # 画像を処理
        output1, output2, output3, output4 = process_image(image)

        # 保存先のパスを作成
        output1_path = f"{image_name}_1.png"
        output2_path = f"{image_name}_2.png"
        output3_path = f"{image_name}_3.png"
        output4_path = f"{image_name}_4.png"
        
        # 画像を保存（PNG形式で保存）
        save_image(output1, output1_path)
        save_image(output2, output2_path)
        save_image(output3, output3_path)
        save_image(output4, output4_path)

        # 保存したファイルをリストに追加
        zip_files_list = [output1_path, output2_path, output3_path, output4_path]
        zip_path = f"{image_name}.zip"
        zip_files(zip_files_list, zip_path)
        
        # 出力として画像とzipファイルを返す
        outputs = [output1, output2, output3, output4]
        return outputs, zip_path

    def launch(self, share):
        with self.demo:
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type='filepath')
                    submit = gr.Button(value="Start")
                with gr.Row():
                    with gr.Column():
                        with gr.Tab("output"):
                            output_0 = gr.Gallery(format="png")
                        output_file = gr.File()
            submit.click(
                self.main, 
                inputs=[input_image], 
                outputs=[output_0, output_file]
            )

        self.demo.queue()
        self.demo.launch(share=share)

if __name__ == "__main__":
    ui = webui()
    ui.launch(share=True)