import gradio as gr
import os
from sora_gen import SoraImageGenerator
import tempfile
from dotenv import load_dotenv

# 加载.env文件（如果存在）
load_dotenv()

# 从环境变量读取认证信息
AUTH_TOKEN = os.environ.get("SORA_AUTH_TOKEN", "")  # 认证Token
USERNAME = os.environ.get("SORA_USERNAME", "sora_user")  # 用户名

# 打印环境变量状态
if AUTH_TOKEN:
    print("已从环境变量读取认证Token")
else:
    print("警告: 环境变量中未设置SORA_AUTH_TOKEN，将使用界面输入")
    
print(f"使用用户名: {USERNAME}")

def generate_images(prompt, num_images, width, height, token=None):
    try:
        # 优先使用界面输入的token，如果为空则使用环境变量
        token_to_use = token if token else AUTH_TOKEN
        if not token_to_use:
            return ["错误: 未提供认证Token"]
            
        # 使用token和用户名实例化生成器
        generator = SoraImageGenerator(
            proxy_host=None, 
            proxy_port=None, 
            auth_token=token_to_use,
            username=USERNAME  # 传入环境变量中的用户名
        )
        
        # 生成图像
        result = generator.generate_image(prompt, int(num_images), int(width), int(height))
        
        if isinstance(result, list):
            return result
        else:
            return [result]  # 错误信息作为列表返回
    except Exception as e:
        return [f"错误: {str(e)}"]

def upload_and_remix(prompt, num_images, image, token=None):
    try:
        # 优先使用界面输入的token，如果为空则使用环境变量
        token_to_use = token if token else AUTH_TOKEN
        if not token_to_use:
            return ["错误: 未提供认证Token"]
            
        # 在临时目录保存上传的图像
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
            temp_path = tmp_file.name
            image.save(temp_path)
        
        # 实例化生成器，传入用户名
        generator = SoraImageGenerator(
            proxy_host=None, 
            proxy_port=None, 
            auth_token=token_to_use,
            username=USERNAME  # 传入环境变量中的用户名
        )
        
        # 上传图像
        upload_result = generator.upload_image(temp_path)
        
        # 清理临时文件
        try:
            os.remove(temp_path)
        except:
            pass
            
        if not isinstance(upload_result, dict) or 'id' not in upload_result:
            return [f"上传失败: {upload_result}"]
        
        media_id = upload_result['id']
        
        # 执行重混
        result = generator.generate_image_remix(prompt, media_id, int(num_images))
        
        if isinstance(result, list):
            return result
        else:
            return [result]
    except Exception as e:
        return [f"错误: {str(e)}"]

# 创建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("# Sora图像生成工具")
    
    # 显示环境变量配置状态
    env_status = []
    if AUTH_TOKEN:
        env_status.append("✅ Token已从环境变量读取")
    else:
        env_status.append("❌ 未设置Token环境变量 (需手动输入)")
        
    env_status.append(f"👤 用户名: {USERNAME}")
    
    gr.Markdown("### 环境配置:\n" + "\n".join(env_status))
    
    with gr.Tab("文本生成图像"):
        with gr.Row():
            with gr.Column():
                txt_prompt = gr.Textbox(label="提示词", placeholder="一只戴着墨镜的可爱柯基犬在沙滩上晒太阳")
                txt_num = gr.Slider(minimum=1, maximum=4, step=1, value=1, label="图像数量")
                txt_width = gr.Slider(minimum=256, maximum=1024, step=8, value=720, label="宽度")
                txt_height = gr.Slider(minimum=256, maximum=1024, step=8, value=480, label="高度")
                # 仅当环境变量未设置时显示Token输入框
                txt_token = gr.Textbox(
                    label="认证Token (Bearer xxx...)" if not AUTH_TOKEN else "认证Token (已从环境变量读取)",
                    placeholder="Bearer eyJhbGciOiJSUzI1NiI..." if not AUTH_TOKEN else "使用环境变量中的Token",
                    type="password",
                    visible=not bool(AUTH_TOKEN)  # 如果环境变量有值则隐藏
                )
                txt_generate = gr.Button("生成图像")
            
            with gr.Column():
                txt_output = gr.Gallery(label="生成结果", columns=2)
        
        txt_generate.click(generate_images, inputs=[txt_prompt, txt_num, txt_width, txt_height, txt_token], outputs=txt_output)
    
    with gr.Tab("图像重混"):
        with gr.Row():
            with gr.Column():
                remix_prompt = gr.Textbox(label="提示词", placeholder="把它变成水彩画风格")
                remix_num = gr.Slider(minimum=1, maximum=4, step=1, value=1, label="图像数量")
                remix_image = gr.Image(type="pil", label="上传原始图像")
                # 仅当环境变量未设置时显示Token输入框
                remix_token = gr.Textbox(
                    label="认证Token (Bearer xxx...)" if not AUTH_TOKEN else "认证Token (已从环境变量读取)",
                    placeholder="Bearer eyJhbGciOiJSUzI1NiI..." if not AUTH_TOKEN else "使用环境变量中的Token",
                    type="password",
                    visible=not bool(AUTH_TOKEN)  # 如果环境变量有值则隐藏
                )
                remix_generate = gr.Button("重混图像")
            
            with gr.Column():
                remix_output = gr.Gallery(label="重混结果", columns=2)
        
        remix_generate.click(upload_and_remix, inputs=[remix_prompt, remix_num, remix_image, remix_token], outputs=remix_output)

# Hugging Face Spaces启动方式
if __name__ == "__main__":
    demo.launch()