import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline
import spaces

# 기본 설정
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# 모델 로드 (FLUX.1-schnell)
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=dtype
).to(device)

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

@spaces.GPU()
def generate_image(prompt, seed, randomize_seed, width, height, steps, guidance_scale):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device).manual_seed(seed)
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=steps,
        generator=generator,
        guidance_scale=guidance_scale
    ).images[0]
    return image, seed

# CSS 스타일 (좌측 사이드바 및 전체 레이아웃 참고)
css = """
body {
    background: linear-gradient(135deg, #667eea, #764ba2);
    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
    color: #333;
    margin: 0;
    padding: 0;
}
.gradio-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 30px 40px;
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
    margin: 40px auto;
    max-width: 1200px;
}
.gradio-container h1 {
    color: #333;
    text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.2);
}
.sidebar {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
}
button, .btn {
    background: linear-gradient(90deg, #ff8a00, #e52e71);
    border: none;
    color: #fff;
    padding: 12px 24px;
    text-transform: uppercase;
    font-weight: bold;
    letter-spacing: 1px;
    border-radius: 5px;
    cursor: pointer;
    transition: transform 0.2s ease-in-out;
}
button:hover, .btn:hover {
    transform: scale(1.05);
}
"""

# 예제 프롬프트 (각 탭별로 다양한 예시)
example_prompts = {
    "Flowchart": [
        "A hand-drawn style flowchart depicting a software release pipeline with clear nodes for development, testing, deployment, and maintenance.",
        "An illustrated business process flowchart for a customer service workflow, with decision points and clear labels.",
        "A creative flowchart showing the steps of a marketing campaign from ideation to execution."
    ],
    "Infographic": [
        "A visually appealing infographic displaying global sales data with vibrant colors, icons, and modern design elements.",
        "An infographic illustrating startup growth metrics with graphs, charts, and minimalist design.",
        "A data-driven infographic showcasing key performance indicators for a corporate strategy, with clear sections and illustrations."
    ],
    "Mockup": [
        "A sketch-style UX mockup for a mobile banking app login flow, featuring clean lines and minimalist design.",
        "A wireframe mockup for an e-commerce website homepage, with user-friendly navigation and modern layout.",
        "A prototype mockup for a productivity dashboard with a focus on intuitive user interface elements."
    ],
    "Diagram": [
        "An educational diagram of a supply chain process, with clear labels and vibrant, friendly illustrations.",
        "A business diagram showing the flow of information between departments in an organization, with modern icons and a clean layout.",
        "A conceptual diagram of a data pipeline, illustrating each step with simple, engaging visuals."
    ],
    "Design": [
        "A sleek industrial design concept for a modern office chair, featuring ergonomic curves and minimalist aesthetics.",
        "A futuristic design for a high-tech smart conference room, blending modern materials with interactive displays.",
        "A creative product design for a smart coffee machine with a metallic finish and touch interface."
    ]
}

# 간단한 예제 핸들러 (예제 버튼 클릭 시 프롬프트 업데이트)
def set_prompt(example_text):
    return example_text

with gr.Blocks(css=css, title="비즈니스 에이전트 이미지 생성기") as demo:
    gr.Markdown(
        """
        <div style="text-align:center;">
            <h1>비즈니스 에이전트 이미지 생성기</h1>
            <p>비즈니스에 필요한 다양한 요소를 탭별로 구분하여 프로세스별 의사결정 및 디자인을 생성합니다.</p>
            <p><strong>Gini's AI Spaces:</strong> Flowchart, Infographic, Mockup, Diagram, Design</p>
        </div>
        """
    )
    
    with gr.Row():
        # 메인 영역: 탭별로 분리된 생성 인터페이스
        with gr.Column(scale=8):
            with gr.Tabs():
                ## Flowchart 탭
                with gr.Tab("Flowchart"):
                    flow_prompt = gr.Textbox(label="Flowchart Prompt", placeholder="Enter a flowchart description...", lines=5)
                    flow_generate = gr.Button("Generate Flowchart")
                    flow_image = gr.Image(label="Generated Flowchart")
                    with gr.Accordion("Example Prompts", open=False):
                        for ex in example_prompts["Flowchart"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=flow_prompt)
                
                ## Infographic 탭
                with gr.Tab("Infographic"):
                    info_prompt = gr.Textbox(label="Infographic Prompt", placeholder="Enter an infographic description...", lines=5)
                    info_generate = gr.Button("Generate Infographic")
                    info_image = gr.Image(label="Generated Infographic")
                    with gr.Accordion("Example Prompts", open=False):
                        for ex in example_prompts["Infographic"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=info_prompt)
                
                ## Mockup 탭
                with gr.Tab("Mockup"):
                    mock_prompt = gr.Textbox(label="Mockup Prompt", placeholder="Enter a mockup description...", lines=5)
                    mock_generate = gr.Button("Generate Mockup")
                    mock_image = gr.Image(label="Generated Mockup")
                    with gr.Accordion("Example Prompts", open=False):
                        for ex in example_prompts["Mockup"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=mock_prompt)
                
                ## Diagram 탭
                with gr.Tab("Diagram"):
                    diag_prompt = gr.Textbox(label="Diagram Prompt", placeholder="Enter a diagram description...", lines=5)
                    diag_generate = gr.Button("Generate Diagram")
                    diag_image = gr.Image(label="Generated Diagram")
                    with gr.Accordion("Example Prompts", open=False):
                        for ex in example_prompts["Diagram"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=diag_prompt)
                
                ## Design 탭
                with gr.Tab("Design"):
                    des_prompt = gr.Textbox(label="Design Prompt", placeholder="Enter a design concept...", lines=5)
                    des_generate = gr.Button("Generate Design")
                    des_image = gr.Image(label="Generated Design")
                    with gr.Accordion("Example Prompts", open=False):
                        for ex in example_prompts["Design"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=des_prompt)
        
        # 좌측 사이드바: 공통 생성 파라미터
        with gr.Sidebar(label="Parameters", open=True):
            gr.Markdown("### Generation Parameters")
            seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            width_slider = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            height_slider = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            steps_slider = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=20)
            guidance_slider = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=15.0, step=0.5, value=7.5)
    
    # 각 탭별 생성 버튼에 이벤트 연결
    flow_generate.click(
        fn=generate_image,
        inputs=[flow_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[flow_image, seed_slider]
    )
    
    info_generate.click(
        fn=generate_image,
        inputs=[info_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[info_image, seed_slider]
    )
    
    mock_generate.click(
        fn=generate_image,
        inputs=[mock_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[mock_image, seed_slider]
    )
    
    diag_generate.click(
        fn=generate_image,
        inputs=[diag_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[diag_image, seed_slider]
    )
    
    des_generate.click(
        fn=generate_image,
        inputs=[des_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[des_image, seed_slider]
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )