import os
import requests
import json
import sys

# ================= 配置区域 =================
# 你的 Ollama 模型名称
MODEL_NAME = "gemma3:12b"  

# Ollama API 地址
OLLAMA_API_URL = "http://localhost:11434/api/chat"

# 需要处理的文件后缀
EXTENSIONS = ('.wav', '.mp3', '.flac', '.aiff', '.ogg', '.m4a')

# 翻译缓存 (保留在内存中，处理下一个文件夹时如果有重复词，速度会飞快)
translation_cache = {}

# ===========================================

def get_translation_via_requests(filename, model):
    """使用 requests 直接调用 Ollama 接口"""
    if filename in translation_cache:
        return translation_cache[filename]

    system_prompt = (
        "You are a professional translator for Audio Sound Effects (SFX) libraries. "
        "Your task is to translate English filenames into concise Simplified Chinese. "
        "Rules: 1. Only output the Chinese translation. 2. Do not explain. 3. Keep technical terms accurate (e.g., 'Whoosh' -> '嗖嗖声')."
    )
    
    user_prompt = f'Translate this filename: "{filename}"'

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        "stream": False,
        "options": {
            "temperature": 0.1
        }
    }

    try:
        response = requests.post(OLLAMA_API_URL, json=payload, timeout=60)
        response.raise_for_status()
        
        result_json = response.json()
        translated_text = result_json['message']['content'].strip()

        invalid_chars = '<>:"/\|?*\'"“”。'
        for char in invalid_chars:
            translated_text = translated_text.replace(char, '')
        
        if not translated_text:
            return None

        translation_cache[filename] = translated_text
        return translated_text

    except Exception as e:
        # 这里不打印网络报错细节，以免刷屏，只返回None跳过
        return None

def is_already_chinese(text):
    """检查文件名里是否已经包含中文"""
    for char in text:
        if '\u4e00' <= char <= '\u9fff':
            return True
    return False

def process_folder():
    """单个文件夹处理流程"""
    
    # 1. 获取路径
    while True:
        target_dir = input("\n请拖入(或粘贴)你要处理的【文件夹路径】: ").strip()
        target_dir = target_dir.strip('"').strip("'") # 去除引号
        
        if os.path.exists(target_dir):
            break
        else:
            print("❌ 错误：路径不存在，请重新输入。")

    # 2. 选择模式
    print("\n请选择模式:")
    print("1. 仅模拟 (只打印不修改)")
    print("2. 直接执行 (修改文件名)")
    mode_choice = input("请输入数字 (默认1): ").strip()
    
    is_dry_run = True
    if mode_choice == '2':
        is_dry_run = False
        print("\n>>> ⚠️ 警告：即将开始【真实修改】文件名！ <<<")
    else:
        print("\n>>> 🛡️ 模拟模式：不会修改任何文件 <<<")

    # 3. 扫描与处理
    count = 0
    success_count = 0
    print(f"\n正在扫描目录: {target_dir} ...\n")

    for root, dirs, files in os.walk(target_dir):
        for file in files:
            if file.lower().endswith(EXTENSIONS):
                old_name_with_ext = file
                name_part, ext_part = os.path.splitext(old_name_with_ext)

                if is_already_chinese(name_part):
                    continue
                
                count += 1
                # 动态显示进度
                print(f"[{count}] 正在思考: {name_part} ...", end="\r")
                
                cn_name = get_translation_via_requests(name_part, MODEL_NAME)
                
                padding = " " * 30 # 用于覆盖之前的打印内容
                
                if cn_name:
                    # 格式：【中文】英文.wav
                    new_name_with_ext = f"【{cn_name}】{old_name_with_ext}"
                    old_path = os.path.join(root, old_name_with_ext)
                    new_path = os.path.join(root, new_name_with_ext)
                    
                    if is_dry_run:
                        print(f"[模拟] {name_part} -> 【{cn_name}】{padding}")
                    else:
                        try:
                            os.rename(old_path, new_path)
                            print(f"[成功] {new_name_with_ext}{padding}")
                            success_count += 1
                        except Exception as e:
                            print(f"[失败] {e}{padding}")
                else:
                    # 如果翻译失败，不需要刷屏，静默跳过或简单提示
                    pass

    print("\n" + "-"*30)
    print(f"✅ 本次任务完成！扫描: {count} 个，成功重命名: {success_count} 个。")

def main():
    print("=============================================")
    print("   音效文件名智能翻译工具 (Gemma + Ollama)")
    print("=============================================")
    print(f"当前模型: {MODEL_NAME}")

    while True:
        process_folder()
        
        # 循环询问
        print("\n" + "="*30)
        choice = input("🔄 是否继续处理其他目录？(y/n): ").strip().lower()
        
        if choice == 'y':
            print("\n" * 2) # 空两行，视觉上分隔
            continue
        else:
            print("\n👋 再见！")
            break

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n用户强制停止，程序退出。")