import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline
from huggingface_hub import InferenceClient

# --- Image Generation Setup (from first example) ---

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/sdxl-turbo"  # Or your preferred image model

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

image_pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
image_pipe = image_pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024


# --- Image Generation Function (from first example) ---

def infer_image(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    image = image_pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, seed


# --- Chatbot Setup (from second example) ---

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")  # Or your preferred chat model


# --- Chatbot Function (from second example) ---

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": message})

    response = ""

    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = chunk.choices[0].delta.content
        response += token
        yield response

# --- Combined Gradio Interface ---
css = """
#col-container {
    margin: 0 auto;
    max-width: 800px; /* Increased max-width for wider layout */
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# Combined Text-to-Image and Chatbot")

    with gr.Tab("Text-to-Image"):
        with gr.Column(elem_id="col-container"):
            with gr.Row():
                image_prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                image_run_button = gr.Button("Generate Image", scale=0, variant="primary")

            image_result = gr.Image(label="Result", show_label=False)

            with gr.Accordion("Advanced Image Settings", open=False):
                image_negative_prompt = gr.Text(
                    label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt"
                )
                image_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                image_randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Row():
                    image_width = gr.Slider(
                        label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512
                    )
                    image_height = gr.Slider(
                        label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512
                    )
                with gr.Row():
                    image_guidance_scale = gr.Slider(
                        label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=0.0
                    )
                    image_num_inference_steps = gr.Slider(
                        label="Number of inference steps", minimum=1, maximum=50, step=1, value=2
                    )

            image_examples = [
                "A photo of a cat riding a unicorn",
                "A futuristic cityscape at sunset",
                "An abstract painting of emotions",
            ]
            gr.Examples(examples=image_examples, inputs=[image_prompt])

        image_run_button.click(
            infer_image,
            inputs=[
                image_prompt,
                image_negative_prompt,
                image_seed,
                image_randomize_seed,
                image_width,
                image_height,
                image_guidance_scale,
                image_num_inference_steps,
            ],
            outputs=[image_result, image_seed],
        )


    with gr.Tab("Chatbot"):
        chatbot = gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
            ],
        )

if __name__ == "__main__":
    demo.launch()