import gradio as gr
import requests
import torch
from diffusers import StableDiffusionPipeline

# --- Load Diffusers model ---
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")  # use GPU if available

# --- Functions ---
def greet(name):
    return f"Hello {name}, welcome to AeroVortex 🚀"

def ask_openassistant(prompt):
    url = "https://api.open-assistant.io"  # replace with actual endpoint
    try:
        response = requests.post(url, json={"prompt": prompt})
        return response.json().get("reply", "No response")
    except Exception as e:
        return f"Error contacting OpenAssistant: {e}"

def upload_file(file):
    return f"File {file.name} uploaded successfully to AeroVortex storage."

def generate_image(prompt):
    try:
        image = pipe(prompt).images[0]
        return image
    except Exception as e:
        return f"Error generating image: {e}"

def draft_document(text):
    return f"Drafted document with content: {text[:50]}..."

# --- Gradio Dashboard ---
with gr.Blocks() as demo:
    gr.Markdown("# 🌐 AeroVortex Dashboard")

    with gr.Tab("Greeting"):
        name = gr.Textbox(label="Enter your name")
        output = gr.Textbox(label="Greeting")
        name.submit(greet, inputs=name, outputs=output)

    with gr.Tab("Ask OpenAssistant"):
        prompt = gr.Textbox(label="Enter your prompt")
        reply = gr.Textbox(label="OpenAssistant reply")
        prompt.submit(ask_openassistant, inputs=prompt, outputs=reply)

    with gr.Tab("File Upload"):
        file = gr.File(label="Upload a file")
        file_output = gr.Textbox(label="Upload status")
        file.upload(upload_file, inputs=file, outputs=file_output)

    with gr.Tab("Image Generation"):
        img_prompt = gr.Textbox(label="Enter image prompt")
        img_output = gr.Image(label="Generated image")
        img_prompt.submit(generate_image, inputs=img_prompt, outputs=img_output)

    with gr.Tab("Document Drafting"):
        doc_text = gr.Textbox(label="Enter draft text", lines=5)
        doc_output = gr.Textbox(label="Draft status")
        doc_text.submit(draft_document, inputs=doc_text, outputs=doc_output)

demo.launch()