import os
import requests
import sys

# ================= 配置区域 =================
MODEL_NAME = "gemma3:12b"  
OLLAMA_API_URL = "http://localhost:11434/api/chat"

# 翻译缓存
translation_cache = {}

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

def get_translation(text, model):
    if text in translation_cache:
        return translation_cache[text]

    system_prompt = (
        "You are a professional translator for Audio/Video asset libraries. "
        "Translate the category/folder name into concise Simplified Chinese. "
        "Rules: 1. Only output the Chinese translation. 2. Do not explain. "
        "3. Keep terms like 'SFX', 'BGM', 'Foley' accurate."
    )
    
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f'Translate folder name: "{text}"'}
        ],
        "stream": False,
        "options": {"temperature": 0.1}
    }

    try:
        response = requests.post(OLLAMA_API_URL, json=payload, timeout=60)
        if response.status_code == 200:
            translated_text = response.json()['message']['content'].strip()
            # 清洗特殊字符
            for char in '<>:"/\|?*\'"“”。':
                translated_text = translated_text.replace(char, '')
            
            if translated_text:
                translation_cache[text] = translated_text
                return translated_text
    except:
        pass
    return None

def is_already_chinese(text):
    """只要包含任意一个中文字符，就认为是中文命名的"""
    for char in text:
        if '\u4e00' <= char <= '\u9fff':
            return True
    return False

def process_folder_renaming():
    while True:
        target_dir = input("\n请拖入(或粘贴)要处理文件夹的【根目录路径】: ").strip().strip('"').strip("'")
        if os.path.exists(target_dir):
            break
        print("❌ 路径不存在，请重试。")

    print("\n请选择模式:")
    print("1. 仅模拟 (推荐先看一眼)")
    print("2. 直接执行")
    is_dry_run = input("请输入数字 (默认1): ").strip() != '2'

    if not is_dry_run:
        print("\n>>> ⚠️ 警告：即将修改文件夹名称！ <<<")
    
    count = 0
    success_count = 0
    skipped_chinese_count = 0
    
    print(f"\n正在扫描并分析目录结构 (倒序模式)...")

    # topdown=False 确保先处理最里面的子文件夹
    for root, dirs, files in os.walk(target_dir, topdown=False):
        for dirname in dirs:
            
            # 【核心保护机制】
            # 如果文件夹名字里已经有中文，直接跳过，并在控制台显示灰色/提示信息
            if is_already_chinese(dirname):
                # 这里为了不刷屏，只在 dry_run 或者特定情况下打印，或者只计数
                # 既然你担心根目录被改，我们可以把跳过的信息打印出来让你放心
                # print(f"[跳过-已含中文] {dirname}") 
                skipped_chinese_count += 1
                continue
            
            count += 1
            print(f"[{count}] 正在思考: {dirname} ...", end="\r")
            
            cn_name = get_translation(dirname, MODEL_NAME)
            padding = " " * 40
            
            if cn_name:
                new_dirname = f"【{cn_name}】{dirname}"
                
                old_path = os.path.join(root, dirname)
                new_path = os.path.join(root, new_dirname)
                
                if is_dry_run:
                    print(f"[模拟] {dirname} -> 【{cn_name}】{dirname}{padding}")
                else:
                    try:
                        os.rename(old_path, new_path)
                        print(f"[成功] {new_dirname}{padding}")
                        success_count += 1
                    except Exception as e:
                        print(f"[失败] {dirname}: {e}{padding}")

    print("\n" + "-"*30)
    print(f"✅ 完成！")
    print(f"   - 翻译并重命名: {success_count} 个文件夹")
    print(f"   - 智能避开已汉化目录: {skipped_chinese_count} 个")

def main():
    print("=== 子目录智能汉化工具 (保护根目录版) ===")
    print(f"当前模型: {MODEL_NAME}")
    
    while True:
        process_folder_renaming()
        if input("\n🔄 是否继续处理其他位置？(y/n): ").lower() != 'y':
            break

if __name__ == "__main__":
    main()