import os
import base64
import numpy as np
from PIL import Image, ImageChops, ImageDraw

import io
import requests
import replicate
import gradio as gr

from dotenv import load_dotenv, find_dotenv

# Locate the .env file
dotenv_path = find_dotenv()
load_dotenv(dotenv_path)
REPLICATE_API_TOKEN = os.getenv('REPLICATE_API_TOKEN')

def sdxl_api(weighted_prompt, weight_amt, prompt, starter_img, prompt_strength):
    input = {
        "seed": 1235,
        "image": starter_img,
        "prompt": "high quality render of (" + weighted_prompt + ")" + str(weight_amt) + " " +  prompt + ", minimalist and simple on a white background",
        "negative_prompt": "worst quality, low quality, illustration, 2d, painting, cartoons, sketch, logo, buttons, markings, text, wires, complex, screws, nails, construction, partial, multiple, pattern, words",
        "prompt_weighting": True,
        #"width": 768, 
        #"height": 768,
        "apply_watermark" : False,
        "scheduler":"K_EULER_ANCESTRAL",
        "prompt_strength":prompt_strength

    }

    output = replicate.run(
        "batouresearch/sdxl-weighting-prompts:66175a2993706e1721076d5c7f92f0c81ec6d065ec20717527f05dd8528a1fc7",
        input=input
    )

    response = requests.get(output[0])
    return Image.open(io.BytesIO(response.content))

def img2imgstarter(prompt):
    input = {
        "prompt": "high quality 3D render of a simple " + prompt + ", front-view, minimalist and simple mockup on a white background",
        "output_format": "jpg",
        "output_quality": 75,
        "steps": 14
    }
    
    try:
        output = replicate.run(
            "stability-ai/stable-diffusion-3",
            input=input
        )
    except Exception as e:
        raise gr.Error(f"Error: {e}")

    try:
        image_url = output[0]
        response = requests.get(image_url)
        img1 = Image.open(io.BytesIO(response.content))
        
        # Save the starter image to a bytes buffer
        buffered = io.BytesIO()
        img1.save(buffered, format="JPEG")
        
        # Encode the starter image to base64
        return "data:image/jpeg;base64," + base64.b64encode(buffered.getvalue()).decode('utf-8')

    except Exception as e:
        raise gr.Error(f"Image download failed: {e}")

def main(text1, text2, prompt, dropdown_value, image_input):

    if dropdown_value=="prompt+image":
        starter_image_pil = Image.fromarray(image_input.astype('uint8'))

        # Resize the starter image if either dimension is larger than 768 pixels
        if starter_image_pil.size[0] > 768 or starter_image_pil.size[1] > 768:
            # Calculate the new size while maintaining the aspect ratio
            if starter_image_pil.size[0] > starter_image_pil.size[1]:
                # Width is larger than height
                new_width = 768
                new_height = int((768 / starter_image_pil.size[0]) * starter_image_pil.size[1])
            else:
                # Height is larger than width
                new_height = 768
                new_width = int((768 / starter_image_pil.size[1]) * starter_image_pil.size[0])
            
            # Resize the image
            starter_image_pil = starter_image_pil.resize((new_width, new_height), Image.LANCZOS)
        # Save the starter image to a bytes buffer
        buffered = io.BytesIO()
        starter_image_pil.save(buffered, format="JPEG")
        
        # Encode the starter image to base64
        starter_img = "data:image/jpeg;base64," + base64.b64encode(buffered.getvalue()).decode('utf-8')
        prompt_strength=.82
    else:
        starter_img = img2imgstarter(prompt)
        prompt_strength = .95

    outputs = [text1]
    output_amts = [2.25, 1.8, 1.5]
    for amt in output_amts:
        outputs.append(sdxl_api(text1, amt, prompt, starter_img, prompt_strength))

    outputs.append(sdxl_api("", 1, prompt, starter_img, .5))

    output_amts.reverse()
    for amt in output_amts:
        outputs.append(sdxl_api(text2, amt, prompt, starter_img, prompt_strength))

    outputs.append(text2)
    return outputs


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            text1 = gr.Textbox(value="organic-shaped", label="Word 1 (Left)")
            text2 = gr.Textbox(value="geometric-shaped", label="Word 2 (Right)")
        with gr.Column(scale=1):
            dropdown = gr.Dropdown(choices=["prompt", "prompt+image"], value="prompt", label="Input Type")
            prompt = gr.Textbox(label="Prompt")
            image_input = gr.Image(label="Image Input", visible=False)
            
    submit_btn = gr.Button("Submit")
    
    with gr.Row():
        output1 = gr.Textbox(label="Word 1")
        output2 = gr.Image(label="Output 1")
        output3 = gr.Image(label="Output 2")
    with gr.Row():
        output4 = gr.Image(label="Output 3")
        output5 = gr.Image(label="Output 4")
        output6 = gr.Image(label="Output 5")
    with gr.Row():
        output7 = gr.Image(label="Output 6")
        output8 = gr.Image(label="Output 7")
        output9 = gr.Textbox(label="Word 2")

    submit_btn.click(main, inputs=[text1, text2, prompt, dropdown, image_input], outputs=[output1, output2, output3, output4, output5, output6, output7, output8, output9])
    
    def update_visibility(selected):
        return gr.update(visible=(selected in ["prompt", "prompt+image"])), gr.update(visible=(selected in ["prompt+image"]))
    
    dropdown.change(fn=update_visibility, inputs=dropdown, outputs=[prompt, image_input])

demo.launch(share=False)