import requests
import base64
from PIL import Image
from io import BytesIO
import gradio as gr
from gradio_imageslider import ImageSlider
import tempfile

# API token and URL for background removal
OCTOAI_TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6IjNkMjMzOTQ5In0.eyJzdWIiOiJiNWQ4NzYzYi0yZjIwLTRiOTctYjI3ZS1jNDdjNjhlNGZhNWYiLCJ0eXBlIjoidXNlckFjY2Vzc1Rva2VuIiwidGVuYW50SWQiOiIxZGRhOWQ2MC01NjQ3LTRmYjQtYWMyZi1iYzM3ODQ0YzU4MjEiLCJ1c2VySWQiOiJmNjc1ZDRhOC1mZjAwLTRkODctYjY4ZC02ODI1YzFiOGE0MmQiLCJhcHBsaWNhdGlvbklkIjoiYTkyNmZlYmQtMjFlYS00ODdiLTg1ZjUtMzQ5NDA5N2VjODMzIiwicm9sZXMiOlsiRkVUQ0gtUk9MRVMtQlktQVBJIl0sInBlcm1pc3Npb25zIjpbIkZFVENILVBFUk1JU1NJT05TLUJZLUFQSSJdLCJhdWQiOiIzZDIzMzk0OS1hMmZiLTRhYjAtYjdlYy00NmY2MjU1YzUxMGUiLCJpc3MiOiJodHRwczovL2lkZW50aXR5Lm9jdG8uYWkiLCJpYXQiOjE3MjY3MTg2MDR9.hLlTJDu5O8BnZnIaj4CZGCwLPh1GqDEsOS2U5bTY6bx3ZpZKvYLQ6m5DVslC-gdCEof243jJV9trkoRKWlO4qXPRnvlZC6gW664tSpSOOjEJ_0rHmZuytXoexWGYZq5RCOyMS2NkZelYYCKT8AXy59YImx3toNNw2aw6ZUWTSX2pmQXXvYnBlRcYESoDF_UK6249F7by77gfYnKw4lDkYurINdX6fMRi0Dd4q_v7X7urB9G1pIZs-Cw5igfCtFRu1xgKkGPHfF_9HSXfDBKlr0s-BXpVR2gnFwshTJzq4lkbbv9kDVN7mb_hSYb6hZbiyx-L7aFRHVHlvZKim9D1cw"
url = "https://image.octoai.run/background-removal"
headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {OCTOAI_TOKEN}"
}

# Function to send request to the API for background removal
from PIL import ImageColor

def hex_to_rgba(hex_color):
    """
    Convert a hex color string (e.g., "#ffcccb") to an RGBA tuple.
    """
    return ImageColor.getcolor(hex_color, "RGBA")

# Function to send request to the API for background removal
def remove_bg(image, background_choice, bg_color, bg_image):
    # Convert PIL image to base64
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

    payload = {
        "driver": "rembg",
        "model": "isnet-general-use",
        "init_image": img_str,
    }

    response = requests.post(url, json=payload, headers=headers)

    if response.status_code == 200:
        result = response.json()
        image_b64 = result.get("image_b64", None)
        
        if image_b64:
            # Convert base64 string back to an image
            image_data = base64.b64decode(image_b64)
            removed_bg_image = Image.open(BytesIO(image_data)).convert("RGBA")

            # Apply background based on user's choice
            if background_choice == "Solid Color":
                # Convert the hex color to an RGBA tuple
                rgba_color = hex_to_rgba(bg_color)
                background = Image.new("RGBA", removed_bg_image.size, rgba_color)
            elif background_choice == "Image" and bg_image is not None:
                background = bg_image.resize(removed_bg_image.size).convert("RGBA")
            else:
                background = Image.new("RGBA", removed_bg_image.size, (255, 255, 255, 255))  # default to white

            # Composite the foreground image onto the chosen background
            final_image = Image.alpha_composite(background, removed_bg_image)

            # Save the image to a temporary file and return it
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            final_image.save(temp_file.name, format="PNG")
            
            return image, final_image, temp_file.name  # Return original, processed image, and file path
    else:
        print(f"Error: {response.status_code}, {response.text}")
        return None, None, None

# Gradio UI with image input, background options, and slider for before/after comparison
with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 480px !important}") as demo:
    with gr.Row():
        input_image = gr.Image(type="pil", label="Input Image")  # Input image
        background_choice = gr.Radio(choices=["Solid Color", "Image"], label="Background Option")  # Choice of solid color or image

    with gr.Row():
        bg_color = gr.ColorPicker(label="Pick a Background Color", value="#ffcccb")  # Color picker for solid background
        bg_image = gr.Image(type="pil", label="Upload Background Image")  # Upload background image

    image_slider = ImageSlider(label="Before/After")  # Before/After slider
    download_button = gr.Button("Download Processed Image")
    download_file = gr.File(label="Download")

    # Function to process the image when the button is clicked
    def process(image, background_choice, bg_color, bg_image):
        original, final_image, file_path = remove_bg(image, background_choice, bg_color, bg_image)
        return [original, final_image], file_path

    # Trigger the process function on button click
    download_button.click(process, inputs=[input_image, background_choice, bg_color, bg_image], outputs=[image_slider, download_file])

    # Example for testing
    gr.Examples(
        examples=[
            ["example_image.jpg", "Solid Color", "#ffcccb", None],  # Example with solid pink background
            ["example_image2.jpg", "Image", "#ffffff", "background_image.jpg"]  # Example with an uploaded background image
        ],
        inputs=[input_image, background_choice, bg_color, bg_image],
        outputs=[image_slider, download_file],
        fn=process,
        cache_examples=True
    )

# Launch the Gradio app
demo.launch()