import os, spaces
import gradio as gr
from huggingface_hub import HfApi, login
from datetime import datetime

@spaces.GPU
def upload_file_to_hf(file_path):
    """
    将ZIP文件上传到TalkTalkAI/RVC Hugging Face模型仓库。
    
    参数:
        file_path (str): 要上传的ZIP文件路径
    
    返回:
        str: 成功或错误消息
    """
    try:
        # 从环境变量获取API令牌
        api_token = os.environ.get("hf_token")
        if not api_token:
            return "错误: 未设置hf_token环境变量"
        
        # 设置仓库ID
        repo_id = "TalkTalkAI/RVC"
        
        # Log in to Hugging Face
        login(token=api_token, add_to_git_credential=False)
        
        # Initialize the Hugging Face API
        api = HfApi()
        
        # 获取文件名并确保是ZIP文件
        filename = os.path.basename(file_path)
        if not filename.lower().endswith('.zip'):
            return "错误: 仅支持ZIP格式文件"
        
        # 在文件名中添加时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        base_name = os.path.splitext(filename)[0]
        path_in_repo = f"{base_name}_{timestamp}.zip"
        
        # 创建提交消息
        commit_message = f"上传模型 {path_in_repo}"
        
        # Upload the file
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=path_in_repo,
            repo_id=repo_id,
            commit_message=commit_message
        )
        
        return f"https://huggingface.co/TalkTalkAI/RVC/resolve/main/{path_in_repo}"
    
    except Exception as e:
        return f"上传文件时出错: {str(e)}"

# Create the Gradio interface
with gr.Blocks(title="RVC模型上传") as demo:
    gr.Markdown("# RVC模型上传工具")
    gr.Markdown("""
    将ZIP文件上传到书梦 AI 歌手模型仓库。
    
    上传文件:
        从本地上传您的 AI 歌手 ZIP 文件，其中包含 .pth 和 .index 模型文件
    
    返回结果:
        您的 AI 歌手 ZIP 文件下载链接，可直接在[书梦](https://www.doingdream.com/)使用
    """)
    
    with gr.Group():
        file = gr.File(label="选择要上传的模型ZIP文件", file_types=[".zip"])
        
    upload_button = gr.Button("上传模型", variant="primary")
    result = gr.Textbox(label="模型下载链接", interactive=False, lines=4)
    
    # 定义上传函数（适用于Gradio）
    def gradio_upload(file):
        if not file:
            return "错误: 未选择任何文件"
        
        if not file.name.lower().endswith('.zip'):
            return "错误: 仅支持ZIP格式文件"
        
        result = upload_file_to_hf(file.name)
        
        # 检查是否上传成功并包含URL
        if "成功上传" in result:
            # 使结果更易于复制粘贴
            gr.Info("上传成功！模型链接已生成")
        
        return result
    
    # Connect the button to the function
    upload_button.click(
        fn=gradio_upload,
        inputs=[file],
        outputs=result
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()