import gradio as gr
import requests
from PIL import Image
import io
import os
import random

# 使用环境变量获取 Token
HF_TOKEN = os.environ.get('HF_TOKEN')
API_URL = "https://api-inference.huggingface.co/models/rupeshs/LCM-runwayml-stable-diffusion-v1-5"
headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}

def generate_image(prompt, style):
    """专注于图像生成，简化界面"""
    
    # 构建风格化提示词
    style_prompts = {
        "梵高风格": f"{prompt}, 梵高风格, 油画笔触, 浓烈色彩",
        "水墨风格": f"{prompt}, 中国传统水墨风格, 黑白灰调, 留白意境", 
        "数字艺术": f"{prompt}, 数字艺术, 未来科技感, 赛博朋克",
        "写实风格": f"{prompt}, 摄影风格, 真实感, 高细节"
    }
    
    full_prompt = style_prompts.get(style, prompt)
    
    # 生成图像
    image = None
    status = "等待生成..."
    
    if HF_TOKEN:
        try:
            response = requests.post(API_URL, headers=headers, json={"inputs": full_prompt})
            if response.status_code == 200:
                image = Image.open(io.BytesIO(response.content))
                status = "✅ 生成成功！"
            else:
                status = f"❌ 生成失败: {response.status_code}"
        except Exception as e:
            status = f"❌ 错误: {str(e)}"
    else:
        status = "❌ 未配置API Token"
    
    return image, status, full_prompt

# 创建简洁的Gradio界面
with gr.Blocks(theme=gr.themes.Soft(), title="AI艺术生成演示") as demo:
    gr.Markdown("""
    # 🎨 AI艺术生成演示
    ## 技术之善 - AI项目制作与伦理探索
    
    **重点展示**: Stable Diffusion 图像生成能力
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="创作描述",
                placeholder="描述你想要生成的画面...",
                value="星空下的城市夜景",
                lines=2
            )
            
            style_radio = gr.Radio(
                choices=["梵高风格", "水墨风格", "数字艺术", "写实风格"],
                label="艺术风格",
                value="梵高风格"
            )
            
            generate_btn = gr.Button("生成图像", variant="primary", size="lg")
        
        with gr.Column():
            image_output = gr.Image(
                label="生成的艺术作品",
                height=400,
                show_download_button=True
            )
            
            status_output = gr.Textbox(
                label="生成状态",
                interactive=False
            )
            
            prompt_used = gr.Textbox(
                label="实际使用的提示词",
                interactive=False
            )
    
    # 示例
    gr.Markdown("### 💡 示例提示词")
    examples = gr.Examples(
        examples=[
            ["星空下的城市夜景", "梵高风格"],
            ["山水间的古代亭台", "水墨风格"],
            ["未来科技的都市景观", "数字艺术"],
            ["森林中的小鹿", "写实风格"]
        ],
        inputs=[prompt_input, style_radio]
    )
    
    # 页脚
    if HF_TOKEN:
        gr.Markdown("---\n**系统状态**: 🟢 已配置API - 可真实生成图像")
    else:
        gr.Markdown("---\n**系统状态**: 🔴 演示模式 - 请在设置中配置HF_TOKEN环境变量")
    
    # 事件处理
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, style_radio],
        outputs=[image_output, status_output, prompt_used]
    )

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