import gradio as gr
import openai
import os
import requests
import base64
from PIL import Image
import io

openai_key = os.getenv('OpenAIKey')
getimg_key = os.getenv('GetImgAIKey')
gpt_model = "gpt-4o-mini"
img_endpoint = "stable-diffusion-xl"
img_model = "juggernaut-xl-v10"
step_count = 14
img_height = 512
img_width = 512

# Helper functions
def get_openai_response(api_key, model, prompt):
    client = openai.OpenAI(api_key=api_key)
    response = client.ChatCompletion.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message['content']

def generate_response(prompt):
    response = get_openai_response(api_key=openai_key, model=gpt_model, prompt=prompt)
    return response

def generate_image(prompt):
    url = f"https://api.getimg.ai/v1/{img_endpoint}/text-to-image"
    headers = {
        "authorization": f"Bearer {getimg_key}",
        "accept": "application/json",
        "content-type": "application/json"
    }
    payload = {
        "prompt": prompt,
        "model": img_model,
        "width": img_width,
        "height": img_height,
        "steps": step_count,
        "output_format": "jpeg",
        "response_format": "b64"
    }

    response = requests.post(url, headers=headers, json=payload)
    if response.status_code == 200:
        # Get the base64 encoded image from the response
        image_base64 = response.json().get('image')
        if image_base64:
            # Decode the base64 image and convert it to an image format Gradio can use
            image_data = base64.b64decode(image_base64)
            image = Image.open(io.BytesIO(image_data))
            return image
        else:
            return "Error: Image generation failed, no image data returned."
    else:
        return f"Error: {response.status_code} - {response.text}"

# Gradio Blocks
with gr.Blocks(css="styles.css", theme='TMBoeren/ThomasMore') as demo:
    gr.Markdown("<h1 class='gr-title'>Get ChatGPT response</h1>")
    gr.Markdown("<span class='gr-description'>Enter your question to get responses from ChatGPT</span>")
    
    input_box = gr.Textbox(label="Your Question", lines=2, elem_id="input_textbox")
    submit_btn = gr.Button(value="Submit")
    output_box = gr.Textbox(label="Response", lines=5, elem_id="output_textbox")
    
    submit_btn.click(generate_response, inputs=[input_box], outputs=[output_box])

    gr.Markdown("<h1 class='gr-title'>Generate image</h1>")
    gr.Markdown("<span class='gr-description'>Enter your prompt to generate an image</span>")
    
    image_prompt_textbox = gr.Textbox(label="Your prompt", lines=2, elem_id="image_prompt_textbox")
    image_submit_btn = gr.Button(value="Submit")
    image_output = gr.Image(label="Generated Image", elem_id="image_output", width=img_width, height=img_height)

    image_submit_btn.click(generate_image, inputs=[image_prompt_textbox], outputs=[image_output])

demo.launch(show_api=False, allowed_paths=["fonts"])