from __future__ import annotations
from os import pipe

import gradio as gr
import numpy as np
from model import Model
from diffusers import StableDiffusionPipeline
import gradio as gr
import torch
import requests 
import json 

TITLE = '# Fictionista 🧚'
DESCRIPTION = '''#
This is an unofficial demo for [https://github.com/NVlabs/stylegan3](https://github.com/NVlabs/stylegan3).
'''


model = Model()
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32) #Change to torch.float16 for faster inference and using cpu
#pipe = pipe.to("cuda")  # Remove or comment out this line if using CPU

sg = None

with gr.Blocks(css='style.css') as image_gen_block:
    gr.Markdown(TITLE)
    gr.Markdown(DESCRIPTION)

    with gr.Tabs():
        with gr.TabItem('Character'):
            with gr.Row():
                with gr.Column():
                    model_name = gr.Dropdown(list(model.MODEL_NAME_DICT.keys()),
                                             value='FemaleHero-256-T',
                                             label='Model')
                    seed = gr.Slider(0,
                                     np.iinfo(np.uint32).max,
                                     step=1,
                                     value=0,
                                     label='Seed')
                    psi = gr.Slider(0,
                                    2,
                                    step=0.05,
                                    value=0.7,
                                    label='Truncation psi')
                    tx = gr.Slider(-1,
                                   1,
                                   step=0.05,
                                   value=0,
                                   label='Translate X')
                    ty = gr.Slider(-1,
                                   1,
                                   step=0.05,
                                   value=0,
                                   label='Translate Y')
                    angle = gr.Slider(-180,
                                      180,
                                      step=5,
                                      value=0,
                                      label='Angle')
                    run_button = gr.Button('Run')
                with gr.Column():
                    result = gr.Image(label='Result', elem_id='Character')

                        # City generation tab
        with gr.TabItem('City'):
            with gr.Row():
                generate_city_button = gr.Button("Generate City")
            with gr.Row():
                city_output = gr.Image(label="Generated City", elem_id="city_output")

        with gr.TabItem('Story'):
            with gr.Row():
                api_key_input_sg = gr.Textbox(label="OpenAI API Key")
                prompt_input = gr.Textbox(label='Prompt')
            generate_button = gr.Button('Generate Story')

            with gr.Row():
                story_output = gr.Textbox(label='Generated Story',
                                          placeholder='Click "Generate Story" to see the story',
                                          readonly=True)



    def generate_story(prompt, API_KEY, model="gpt-4", temperature=1, max_tokens=None):
        API_ENDPOINT = "https://api.openai.com/v1/chat/completions"

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}",
        }
        prompt = [{"role": "user", "content":prompt}]
        data = {
            "model": model,
            "messages": prompt,
            "temperature": temperature,
        }

        if max_tokens is not None:
            data["max_tokens"] = max_tokens

        response = requests.post(API_ENDPOINT, headers=headers, data=json.dumps(data))

        if response.status_code == 200:
            return response.json()["choices"][0]["message"]["content"]
        else:
            raise Exception(f"Error {response.status_code}: {response.text}")
        

    def generate_city():
        city_prompt = f"render 2048K ultra realistic render high definition a very technologically advanced city, cyberpunk, hyper realistic, detailed, ray tracing, POV-Ray, metallic, Sci-Fi"
        city_image = pipe(city_prompt).images[0]
        return city_image

    generate_city_button.click(fn=generate_city, inputs=None, outputs=city_output)
    model_name.change(fn=model.set_model, inputs=model_name, outputs=None)
    run_button.click(fn=model.generate_image,
                    inputs=[
                        model_name,
                        seed,
                        psi,
                        tx,
                        ty,
                        angle,
                    ],
                    outputs=result)
    generate_button.click(fn=generate_story, inputs=[prompt_input, api_key_input_sg], outputs=story_output)




image_gen_block.queue().launch(show_api=False)