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

def image_to_base64(img):
    buffered = BytesIO()
    # Downscale for mobile brochure feel (max width 600px)
    img.thumbnail((600, 600))
    img.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def generate_brochure(image_files, info_file):
    if not image_files or not info_file:
        return None
    
    # 1. Read and split the info text
    with open(info_file.name, 'r') as f:
        text_content = f.read().split('\n\n') # Assumes double new-line separates info
    
    md_content = "# Project Brochure\n\n"
    
    # 2. Logic to group 3 images and 3 text blocks
    # We loop through images in steps of 3
    for i in range(0, len(image_files), 3):
        md_content += '<div style="max-width: 400px; margin: auto; border: 1px solid #ccc; padding: 10px;">\n\n'
        md_content += f"## Page { (i//3) + 1 }\n\n"
        
        # Get the next 3 images and next 3 text blocks
        current_images = image_files[i:i+3]
        current_texts = text_content[i:i+3]
        
        for img_path, txt in zip(current_images, current_texts):
            img = Image.open(img_path)
            b64 = image_to_base64(img)
            
            # Add to Markdown with HTML for mobile-style sizing
            md_content += f'<img src="data:image/png;base64,{b64}" style="width:100%; border-radius:8px;">\n\n'
            md_content += f"{txt}\n\n---\n\n"
        
        md_content += '</div>\n<div style="page-break-after: always;"></div>\n\n'

    # 3. Save the result
    output_path = "brochure.md"
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(md_content)
        
    return output_path

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 📱 Mobile Brochure Generator")
    gr.Markdown("Upload images and an `info.txt` file to generate a self-contained Markdown brochure.")
    
    with gr.Row():
        with gr.Column():
            img_input = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
            txt_input = gr.File(label="Upload info.txt", file_types=[".txt"])
            gen_btn = gr.Button("Generate Brochure", variant="primary")
            
        with gr.Column():
            out_file = gr.File(label="Download .md Brochure")

    gen_btn.click(
        fn=generate_brochure,
        inputs=[img_input, txt_input],
        outputs=out_file
    )

if __name__ == "__main__":
    # Note: Using the new Gradio 6.0+ launch style
    demo.launch(theme=gr.themes.Soft())