import os
import gradio as gr
import requests
import base64
from PIL import Image
from io import BytesIO
import tempfile
import zipfile

API_KEY = os.getenv("SEGMIND_API_KEY")
API_URL = "https://api.segmind.com/v1/nomos-upscaler"

def image_file_to_base64(image_path):
    with open(image_path, 'rb') as f:
        image_data = f.read()
    return base64.b64encode(image_data).decode('utf-8')

def image_url_to_base64(image_url):
    response = requests.get(image_url)
    image_data = response.content
    return base64.b64encode(image_data).decode('utf-8')

def upscale_single_image(base64_image):
    data = {
        "image": base64_image,
        "image_format": "png",
        "image_quality": 95,
        "base64": False
    }
    headers = {'x-api-key': API_KEY}
    response = requests.post(API_URL, json=data, headers=headers)
    if response.status_code == 200:
        return Image.open(BytesIO(response.content))
    else:
        raise Exception(f"API Error: {response.status_code} - {response.text}")

def upscale_images(files, urls, state):
    results = []
    filepaths = []
    temp_dir = tempfile.mkdtemp()
    # Process uploaded files
    if files:
        for file_path in files:
            try:
                base64_img = image_file_to_base64(file_path)
                upscaled_img = upscale_single_image(base64_img)
                out_path = os.path.join(temp_dir, f"upscaled_{os.path.basename(file_path)}.png")
                upscaled_img.save(out_path, format="PNG")
                results.append(upscaled_img)
                filepaths.append(out_path)
            except Exception as e:
                results.append(f"Error: {str(e)}")
    # Process URLs
    if urls:
        url_list = [u.strip() for u in urls.splitlines() if u.strip()]
        for idx, url in enumerate(url_list):
            try:
                base64_img = image_url_to_base64(url)
                upscaled_img = upscale_single_image(base64_img)
                out_path = os.path.join(temp_dir, f"upscaled_url_{idx+1}.png")
                upscaled_img.save(out_path, format="PNG")
                results.append(upscaled_img)
                filepaths.append(out_path)
            except Exception as e:
                results.append(f"Error: {str(e)}")
    # Update state with new filepaths
    state = filepaths
    return results, state

def zip_upscaled_images(state):
    if not state:
        return None
    temp_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    with zipfile.ZipFile(temp_zip.name, 'w') as zipf:
        for fpath in state:
            zipf.write(fpath, arcname=os.path.basename(fpath))
    return temp_zip.name

with gr.Blocks() as demo:
    gr.Markdown("# Bulk Image Upscaler\nUpload multiple images or paste image URLs (one per line). Download all upscaled images as a zip.")
    state = gr.State([])
    with gr.Row():
        file_input = gr.File(label="Upload Images", file_count="multiple", type="filepath")
        url_input = gr.Textbox(label="Image URLs (one per line)", lines=5)
    output_gallery = gr.Gallery(label="Upscaled Images", show_label=True)
    run_btn = gr.Button("Upscale Images")
    download_zip_btn = gr.Button("Download All as ZIP")
    zip_file_output = gr.File(label="Download ZIP")
    
    run_btn.click(
        upscale_images,
        inputs=[file_input, url_input, state],
        outputs=[output_gallery, state]
    )
    download_zip_btn.click(
        zip_upscaled_images,
        inputs=[state],
        outputs=zip_file_output
    )

if __name__ == "__main__":
    # Make sure to set the environment variable SEGMIND_API_KEY before running
    # For example, in Linux/Mac: export SEGMIND_API_KEY=your_api_key
    demo.launch()