import gradio as gr
import requests
import base64
from PIL import Image
import io
import json
import os
import numpy
#from dotenv import load_dotenv

#dotenv_path = '/.env'
#load_dotenv(dotenv_path)

# Get the API key from the environment variable
API_KEY = os.getenv('BASETEN_API_KEY')
API_url = os.getenv('BASETEN_COMFY_MODEL_URL')

def call_api_and_generate_image(negative_prompt, positive_prompt,seed,source_image_base64):
    if not API_KEY:
        return None, "Error: API key not found in environment variables. Please set BASETEN_API_KEY."

    # Convert the source image to base64
    #buffered = io.BytesIO()
    #source_image.save(buffered, format="PNG")
    #img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')

    resp = requests.post(
        API_url,
        headers={"Authorization": f"Api-Key {API_KEY}"},
        json={
            'workflow_values': {
                'negative_prompt': negative_prompt,
                'positive_prompt': positive_prompt,
                'seed': seed,
                'source_image_fs': source_image_base64
            }
        }
    )
    
    result = resp.json()
    print(result)
    base64_data = result['result'][0]['data']
    image_data = base64.b64decode(base64_data)
    image = Image.open(io.BytesIO(image_data))
    return image#, json.dumps(result, indent=2)

def generate_image(positive_prompt, negative_prompt,source_image):
    seed = numpy.random.randint(0, 2**32 - 1)
    source_image_base64 = None
    if source_image is not None:
        with open(source_image, "rb") as image_file:
            source_image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
    try:
        return call_api_and_generate_image(negative_prompt, positive_prompt,seed,source_image_base64)
    except Exception as e:
        return None#, f"Error: {str(e)}"

def clear_fields():
    return "", "", None  # Clear prompt and outputs

with gr.Blocks(theme='freddyaboulton/test-blue') as demo:
    gr.Markdown("<center><h2>Arjun's Face Swapper</h2></center>")
    gr.Markdown("Hi there! Specify a man/woman in any setting below and upload your image to see yourself in that setting!")
    prompt = gr.Textbox(label='Positve Prompt', lines=2, max_lines=5, placeholder = 'Describe your image here, for e.g., superman flying in the sky')
    neg_prompt = gr.Textbox(label='Negative Prompt', lines=2, max_lines=10, value="Low quality, pixelated, poor anatomy, bad fingers")
    source_image = gr.Image(label='Source Image for Face Swap', type="filepath")
    with gr.Group():
        with gr.Row():
                submit_btn = gr.Button(value="Submit", elem_id="generate_button", variant="primary", size="sm")
                clear_btn = gr.ClearButton(value="Clear Question and AI Response", elem_id="clear_button", variant="secondary", size="sm")
        gr.Markdown("<center><h3>AI Response</h3></center>")
        image_output_box = gr.Image(type="pil", label="Final Generated Image")
        #json_output = gr.Textbox(label="API Response JSON", lines=10)
    submit_btn.click(fn=generate_image, inputs = [prompt,neg_prompt,source_image], outputs=[image_output_box])
    clear_btn.click(fn=clear_fields,outputs=[prompt,neg_prompt,image_output_box])

demo.launch()