import os
import io
import requests
from uuid import uuid4

import boto3
from botocore.client import Config
from PIL import Image
import gradio as gr


css = """
#col-container {
    margin: 0 auto;
    max-width: 512px;
}
"""

MAX_PIXEL_BUDGET = 1024 * 1024


def upload_to_r2(file_path, object_key, content_type):
    s3 = boto3.client(
        's3',
        endpoint_url=os.getenv('R2_ENDPOINT'),
        aws_access_key_id=os.getenv('R2_ACCESS_KEY_ID'),
        aws_secret_access_key=os.getenv('R2_SECRET_ACCESS_KEY'),
        config=Config(signature_version='s3v4'),
        region_name='auto'
    )

    with open(file_path, 'rb') as f:
        s3.put_object(
            Bucket=os.getenv('R2_BUCKET'),
            Key=object_key,
            Body=f,
            ContentType=content_type
        )

    download_url = s3.generate_presigned_url(
        'get_object',
        Params={
            'Bucket': os.getenv('R2_BUCKET'),
            'Key': object_key
        },
        ExpiresIn=3600  # url expiration time in seconds
    )
    
    return download_url


def process_input(input_image, upscale_factor):
    w, h = input_image.size
    w_original, h_original = w, h
    aspect_ratio = w / h

    was_resized = False

    # compute minimum dimension after upscaling
    min_dimension = min(w, h) * upscale_factor
    
    # if minimum dimension is above 1024, adjust scale factor
    if min_dimension > 1024:
        new_scale = 1024 / min(w, h)
        upscale_factor = min(2, new_scale)  # cap at 2x if needed
        gr.Info(f'Adjusted scale factor to {upscale_factor}x to maintain minimum dimension of 1024 pixels')

    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
        gr.Info(
            f'Requested output image is too large. Resizing input to fit within pixel budget.'
        )
        input_image = input_image.resize(
            (
                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
            )
        )
        was_resized = True

    return input_image, w_original, h_original, was_resized, upscale_factor

def infer(
    input_image,
    upscale_factor,
    image_category,
):
    true_input_image = input_image
    input_image, w_original, h_original, was_resized, adjusted_scale = process_input(
        input_image, upscale_factor
    )

    temp_input_path = 'temp_input.png'
    input_image.save(temp_input_path)
    image_url = upload_to_r2(temp_input_path, str(uuid4()), 'image/png')
    gr.Info('Upscaling image...')
    
    try:
        resp = requests.get(
            os.getenv('ENDPOINT') ,
            headers={
                'Modal-Key': os.getenv('AUTH_KEY'),
                'Modal-Secret': os.getenv('AUTH_SECRET'),
            },
            params={
                'image_url': image_url,
                'image_category': image_category,
                'scale_factor': adjusted_scale,
                'output_format': 'png',
                'upload_to_r2': False
            }
        )

        if resp.status_code != 200:
            raise gr.Error(f'API request failed with status {resp.status_code}: {resp.text}')

        # save the response image
        output_path = 'output.png'
        with open(output_path, 'wb') as f:
            f.write(resp.content)

        output_image = Image.open(output_path)
        
        if was_resized:
            gr.Info(
                f'Resizing output image to targeted {w_original * adjusted_scale}x{h_original * adjusted_scale} size.'
            )
            output_image = output_image.resize((int(w_original * adjusted_scale), int(h_original * adjusted_scale)))

        return output_image

    except Exception as e:
        raise gr.Error(f'Error during upscaling: {str(e)}')
    finally:
        if os.path.exists(temp_input_path):
            os.remove(temp_input_path)
        if os.path.exists(output_path):
            os.remove(output_path)

with gr.Blocks(css=css) as demo:
    gr.Markdown(
        """
    # 🚀 Puncta Lite - AI Image Upscaler
    
    This is a lite version of Puncta's AI image upscaler. For more advanced features and higher quality results, visit [puncta.ai](https://www.puncta.ai/).

    *Note*: This demo is limited to a maximum output resolution of 1024x1024 pixels. For higher resolution upscaling, please visit our full version at puncta.ai.
    """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_im = gr.Image(label='Input Image', type='pil')
        with gr.Column(scale=1):
            output_im = gr.Image(label='Output Image', type='pil')

    with gr.Row():
        with gr.Column(scale=1):
            upscale_factor = gr.Slider(
                label='Upscale Factor',
                minimum=2,
                maximum=4,
                step=1,
                value=2,
            )
        with gr.Column(scale=1):
            image_category = gr.Dropdown(
                label='Image Category',
                choices=['general', 'portrait', 'outdoor', 'digital art'],
                value='general'
            )

    with gr.Row():
        run_button = gr.Button(value='Upscale Image')

    examples = gr.Examples(
        examples=[
            ['examples/dogs.jpg', 2, 'outdoor'],
            ['examples/portrait1.png', 3, 'portrait'],
            ['examples/anime_1.png', 2, 'digital art'],
            ['examples/vintage_family_photo.jpg', 2, 'portrait'],
            ['examples/bus1.png', 3, 'general'],
            ['examples/festival_3.png', 2, 'general'],
            ['examples/general_2.png', 2, 'general'],
            ['examples/anime_2.jpg', 3, 'digital art'],
        ],
        inputs=[
            input_im,
            upscale_factor,
            image_category,
        ],
        fn=infer,
        outputs=output_im,
        cache_examples=True,
    )

    gr.Markdown("**Disclaimer:**")
    gr.Markdown(
        "This demo is for testing purposes only. For commercial use and higher quality results, please visit [puncta.ai](https://www.puncta.ai/)."
    )

    gr.on(
        [run_button.click],
        fn=infer,
        inputs=[
            input_im,
            upscale_factor,
            image_category,
        ],
        outputs=output_im,
        show_api=False,
    )

demo.queue().launch(share=False, show_api=False)