import numpy as np
import streamlit as st
import requests
from PIL import Image
from io import BytesIO
from streamlit_drawable_canvas import st_canvas

# API information (your actual token)
url = "https://api.magicstudio.com/magiceraser/erase"
access_token = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhZGF0YSI6bnVsbCwiZ3JhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'  # Your token

# Helper function for downloading the processed image
def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
    pil_format = "JPEG" if fmt == "jpg" else "PNG"
    file_format = "jpg" if fmt == "jpg" else "png"
    mime = "image/jpeg" if fmt == "jpg" else "image/png"

    buf = BytesIO()
    pil_image.save(buf, format=pil_format)

    return st.download_button(
        label=label,
        data=buf.getvalue(),
        file_name=f'{filename}.{file_format}',
        mime=mime,
    )

# Set the title of the app
st.title("AI Photo Object Removal")

# Upload the image
uploaded_file = st.file_uploader("Upload an image to remove objects from", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    # Load the uploaded image
    img_input = Image.open(uploaded_file).convert("RGBA")

    # Resize image if too large
    max_size = 2000
    img_width, img_height = img_input.size
    if img_width > max_size or img_height > max_size:
        if img_width > img_height:
            new_width = max_size
            new_height = int((max_size / img_width) * img_height)
        else:
            new_height = max_size
            new_width = int((max_size / img_height) * img_width)
        img_input = img_input.resize((new_width, new_height))

    # Display the image and let the user draw the mask
    stroke_width = st.slider("Brush size", 1, 100, 50)
    st.write("**Draw over the parts of the image you want to remove.**")

    # Create a drawing canvas for the mask
    canvas_result = st_canvas(
        stroke_color="rgba(255, 0, 255, 1)",  # Pink color for the brush
        stroke_width=stroke_width,
        background_image=img_input,
        update_streamlit=True,
        height=img_input.height,
        width=img_input.width,
        drawing_mode="freedraw",
        key="canvas",
    )

    if canvas_result.image_data is not None:
        # Convert the drawing into a mask
        mask_image = Image.fromarray((canvas_result.image_data[:, :, 3] > 0).astype(int) * 255)
        mask_image = mask_image.resize(img_input.size).convert("L")  # Convert to grayscale

        # Display the mask for the user to review
        st.write("**Generated Mask:**")
        st.image(mask_image)

        if st.button('Submit'):
            with st.spinner("AI is processing..."):
                try:
                    # Convert both image and mask to binary format for API upload
                    img_buffer = BytesIO()
                    mask_buffer = BytesIO()

                    img_input.save(img_buffer, format='PNG')
                    mask_image.save(mask_buffer, format='PNG')

                    # Prepare the files for the API
                    files = [
                        ('image_file', ('image.png', img_buffer.getvalue(), 'image/png')),
                        ('mask_file', ('mask.png', mask_buffer.getvalue(), 'image/png')),
                    ]

                    headers = {
                        'accessToken': access_token
                    }

                    # Send POST request to MagicEraser API
                    response = requests.post(url, headers=headers, files=files)

                    if response.status_code == 200:
                        # Process and display the result
                        result_image = Image.open(BytesIO(response.content)).convert("RGB")
                        st.write("AI has finished the job!")
                        st.image(result_image)

                        # Download button for the output image
                        uploaded_name = uploaded_file.name.split('.')[0]
                        image_download_button(
                            pil_image=result_image,
                            filename=f"{uploaded_name}_output",
                            fmt="jpg",
                            label="Download Image"
                        )
                    else:
                        st.error(f"Error: {response.status_code} - {response.text}")
                except Exception as e:
                    st.error(f"Error processing the image: {e}")