import os
import time
import logging
import gradio as gr
import random
import requests
import json

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

NUM_IMAGES = 1
REFRESH_INTERVAL_S = 1.0
MIN_IMAGE_SIZE = 768
MAX_IMAGE_SIZE = 1024
MAX_SEED = 100

API_KEY = os.getenv('DREAMLOOK_API_KEY')
API_ENDPOINT = os.getenv('DREAMLOOK_API_ENDPOINT', default='https://api.dreamlook.ai')
DREAMLOOK_API_ENDPOINT_IMAGE_CREATE_PATH = os.getenv('DREAMLOOK_API_ENDPOINT_IMAGE_CREATE_PATH', 'image_gen')
DREAMLOOK_API_MODEL_CHECKPOINT = os.getenv('DREAMLOOK_API_MODEL_CHECKPOINT')

API_CREATE_ENDPOINT = f"{API_ENDPOINT}/{DREAMLOOK_API_ENDPOINT_IMAGE_CREATE_PATH}"
API_STATUS_ENDPOINT = f"{API_ENDPOINT}/jobs/image-gen/"

examples = [
    "A screen-printed poster from 1975 with bold, expressive contours features a pensive figure in red and green, set against a dark, moody background. \"Kunstmuseum Bern Sammlung im Obersteg\"",
    "Colorful offset-printed poster from 1990 featuring animated, cartoonish characters with exaggerated expressions. Background features a gradient of blue and light orange hues. Bold, handwritten-style text. \"NOTRE ÉGLISE\" \"TOUTE UNE VIE À SOUTENIR\" \"ÉGLISE NATIONALE PROTESTANTE DE GENÈVE\"",
    "Offset printed poster in black and white from 1983 featuring a photo of a joyful couple embracing in front of a train, with bold text above and below the image. \"J'ai un cours du soir en ville.\" \"Et chaque heure un train.\" \"Vos CFF\"",
    "Lithographic poster from 1911, showcasing three well-dressed men in trench coats and top hats with rich, muted colors and stylized outlining, emphasizing modern men's fashion. \"PKZ\" \"Moderne Herrenbekleidung\" \"Burger & Zimmermann\" \"Luzern - Reußbrücke - Kramgasse 2\"",
]

NEGATIVE_PROMPT = "lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, text, signature, watermark, username, blurry"


def infer(prompt, negative_prompt, seed, randomize_seed, guidance_scale):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    headers = {
        "content-type": "application/json",
        "authorization": f"Bearer {API_KEY}"
    }

    payload = {
        # "base_model": "stable-diffusion-xl-v1-0"
        "checkpoint_id": DREAMLOOK_API_MODEL_CHECKPOINT,
        "prompt": prompt[:500],
        "negative_prompt": NEGATIVE_PROMPT,
        "num_samples": NUM_IMAGES,
        "width": 704,
        "height": 1024,
        "num_inference_steps": 30,
        "enable_hrf": False,
        "scheduler_type": "dpmpp-2m-karras",
        "seed": seed,
        "model_type": "sdxl-v1",
        "guidance_scale": guidance_scale,
    }
    logging.debug(f"Payload: {json.dumps(payload, indent=2)}")

    response = requests.post(API_CREATE_ENDPOINT, headers=headers, data=json.dumps(payload))
    logging.debug(f"API response status code: {response.status_code}")
    logging.debug(f"API response: {response.text}")

    if response.status_code == 200:
        job_id = response.json()['job_id']
        logging.info(f"Job submitted successfully. Job ID: {job_id}")

        image_url = None

        while True:
            status_response = requests.get(f"{API_STATUS_ENDPOINT}{job_id}", headers=headers)
            status_data = status_response.json()
            state = status_data.get('state')
            yield state, image_url
            logging.info(f"Job state: {state}")

            if state == 'success':
                image_url = status_data['image_gen_results'][0]['generated_images'][0]['url']
                logging.info(f"Image generation succeeded. Image URL: {image_url}")
                yield state, image_url
                return
            elif state == 'failed':
                failure_reason = status_data.get('failure_reason', 'Unknown')
                logging.error(f"Job failed with reason: {failure_reason}")
                return f"Error: Job failed with reason: {failure_reason}"

            time.sleep(REFRESH_INTERVAL_S)
    else:
        return f"Error: {response.status_code}, {response.text}"


with gr.Blocks() as demo:
    gr.Markdown("# Perspectives on AI")
    gr.Markdown(
        "In collaboration with the Alliance Graphique Internationale [\"Perspectives on AI\"](https://agi-open.com) exhibition, [dreamlook.ai](https://dreamlook.ai) has created an image generation model trained from thousands of posters from the [Basel Poster Collection](https://www.recherche-plakatsammlungbasel.ch/objects) to enable its exploration through prompts.")
    # gr.Markdown("Write a prompt then click on \"Generate image\"")
    gr.Markdown("[Read more...](https://dreamlook.ai/perspectives-on-ai)")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                info="Describe the image you want to generate, or click on a sample prompt below.",
                placeholder="Enter your prompt",
                max_lines=10,

            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Enter a negative prompt (optional)",
                visible=False
            )
            seed = gr.Slider(
                label="Seed",
                info="The random number used to create your image. Change it to explore variants.",
                minimum=0,
                maximum=MAX_SEED,
                value=42,
                visible=True
            )
            randomize_seed = gr.Checkbox(
                label="Randomize Seed",
                value=False,
                visible=False
            )
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                info="Defined how much your prompt is taken into account when generating the image.",
                minimum=5.0,
                maximum=15.0,
                value=10.0
            )
            # num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=10, maximum=30, value=20)
            gr.Examples(
                examples=examples,
                inputs=[prompt],
                label="Sample prompts",
            )
            run_button = gr.Button(
                "Generate image",
            )

        with gr.Column():
            status_box = gr.Textbox(label="Status", value="", interactive=False)
            result = gr.Image(label="Result")

    run_button.click(
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed, guidance_scale],
        outputs=[status_box, result]
    )

demo.launch()