import os
import glob
import subprocess
import pandas as pd
from datasets import Dataset, Features, Image as HFImage, Value
from huggingface_hub import HfApi, DatasetCard
import tempfile
from PIL import Image

# ==========================================
# 配置区
# ==========================================
REPO_ID = "waveshaker/dataset_collections" 
REPO_TYPE = "dataset"
HF_TOKEN = ""  # 请填入你的 Token
# ==========================================

def clean_txt_newlines(current_dir):
    """1. 递归清理所有 txt 文件中的换行符"""
    print("\n" + "="*50)
    print("🧹 [任务 1/5] 正在清理所有 TXT 文件中的换行符...")
    print("="*50)
    
    txt_files = glob.glob(os.path.join(current_dir, "**", "*.txt"), recursive=True)
    if not txt_files:
        print("  ✅ 没有找到需要处理的 TXT 文件。")
        return

    cleaned_count = 0
    for txt_path in txt_files:
        # 跳过根目录下的 txt（可能是说明文档）
        if os.path.dirname(txt_path) == current_dir:
            continue
            
        try:
            with open(txt_path, 'r', encoding='utf-8') as f:
                content = ''.join(line.strip() for line in f)
            
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(content)
            cleaned_count += 1
        except Exception as e:
            print(f"  ⚠️ 处理 {os.path.basename(txt_path)} 失败: {e}")
            
    print(f"  🎉 TXT 清理完成！共处理了 {cleaned_count} 个文件。")

def remove_mask_metadata(current_dir):
    """2. 递归剥离 masklabel.png 的元数据，保护隐私"""
    print("\n" + "="*50)
    print("🛡️ [任务 2/5] 正在剥离 Mask 蒙版图 (masklabel.png) 的元数据...")
    print("="*50)
    
    png_files = glob.glob(os.path.join(current_dir, "**", "*.png"), recursive=True)
    mask_files = [f for f in png_files if "masklabel" in os.path.basename(f).lower()]
    
    if not mask_files:
        print("  ✅ 没有找到 masklabel.png 文件，跳过此步骤。")
        return

    success_count = 0
    for file_path in mask_files:
        try:
            with Image.open(file_path) as img:
                # 读取像素数据，彻底断开原元数据
                data = list(img.getdata())
                clean_img = Image.new(img.mode, img.size)
                clean_img.putdata(data)
                # 覆盖保存
                clean_img.save(file_path, "PNG")
            success_count += 1
            print(f"  🔒 已清理元数据: {os.path.basename(file_path)}")
        except Exception as e:
            print(f"  ⚠️ 处理失败 [{os.path.basename(file_path)}]: {e}")

    print(f"  🎉 元数据清理完成！共处理了 {success_count} 张 Mask 图。")

def convert_png_to_jpg(current_dir):
    """3. 批量转换常规 png 为 jpg，并删除原 png（绝对跳过 masklabel）"""
    print("\n" + "="*50)
    print("🎬 [任务 3/5] 正在批量转换常规 PNG 为 JPG...")
    print("="*50)
    
    png_files = glob.glob(os.path.join(current_dir, "**", "*.png"), recursive=True)
    if not png_files:
        print("  ✅ 没有找到需要转换的 PNG 文件。")
        return

    converted_count = 0
    for png_path in png_files:
        # 🚨 核心过滤：跳过 masklabel 文件，保证蒙版依然是 png
        if "masklabel" in png_path.lower():
            continue
            
        jpg_path = os.path.splitext(png_path)[0] + ".jpg"
        
        try:
            subprocess.run(
                ["ffmpeg", "-y", "-i", png_path, "-q:v", "2", jpg_path],
                stdout=subprocess.DEVNULL, 
                stderr=subprocess.DEVNULL,
                check=True
            )
            os.remove(png_path)
            converted_count += 1
            print(f"  🔄 转换: {os.path.basename(png_path)} -> {os.path.basename(jpg_path)}")
        except subprocess.CalledProcessError:
            print(f"  ⚠️ 转换失败: {png_path}")
        except FileNotFoundError:
            print("  ❌ 错误: 找不到 ffmpeg！请确保已安装并配置了环境变量。")
            return

    print(f"  🎉 转换完成！共转换 {converted_count} 张图片。")

def generate_parquet_and_update_yaml(api, current_dir, region_prefix):
    """4. 本地打包生成 parquet，拉取云端 YAML，追加并重排"""
    print("\n" + "="*50)
    print("📦 [任务 4/5] 正在本地打包 Parquet 并更新 YAML 配置...")
    print("="*50)
    
    local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")
    
    print(f"  ☁️ 正在读取云端 {REPO_ID} 的现有配置...")
    try:
        remote_readme_path = api.hf_hub_download(repo_id=REPO_ID, filename="README.md", repo_type=REPO_TYPE)
        with open(remote_readme_path, "r", encoding="utf-8") as f:
            content = f.read()
        card = DatasetCard(content)
        print("  ✅ 成功读取云端 README。")
    except Exception as e:
        print(f"  ⚠️ 找不到云端 README，将新建: {e}")
        card = DatasetCard("")

    yaml_data = card.data if card.data else {}
    if "configs" not in yaml_data: yaml_data["configs"] = []
    
    yaml_data["dataset_info"] = {
        "features": [
            {"name": "image", "dtype": "image"},
            {"name": "file_name", "dtype": "string"},
            {"name": "text", "dtype": "string"}
        ]
    }

    existing_config_names = [c.get("config_name") for c in yaml_data["configs"]]
    new_configs_added = False
    valid_folders = []

    for sub_folder_name in os.listdir(current_dir):
        sub_folder_path = os.path.join(current_dir, sub_folder_name)
        
        if os.path.isdir(sub_folder_path) and not sub_folder_name.startswith('.'):
            jpg_files = glob.glob(os.path.join(sub_folder_path, "*.jpg")) + glob.glob(os.path.join(sub_folder_path, "*.jpeg"))
            if not jpg_files:
                continue
                
            print(f"  📦 正在打包: {sub_folder_name}")
            data_rows = []
            
            for jpg_path in jpg_files:
                file_name = os.path.basename(jpg_path)
                base_name = os.path.splitext(file_name)[0]
                txt_path = os.path.join(sub_folder_path, f"{base_name}.txt")
                
                text_content = ""
                if os.path.exists(txt_path):
                    with open(txt_path, "r", encoding="utf-8") as f:
                        text_content = f.read().strip()

                try:
                    with open(jpg_path, "rb") as f:
                        image_bytes = {"bytes": f.read(), "path": None}
                except Exception as e:
                    print(f"    ⚠️ 跳过损坏图片 {file_name}: {e}")
                    continue

                data_rows.append({
                    "image": image_bytes,
                    "file_name": file_name,
                    "text": text_content
                })

            if data_rows:
                df = pd.DataFrame(data_rows)
                features = Features({
                    "image": HFImage(),  # 👈 这里加上了 HF
                    "file_name": Value("string"),
                    "text": Value("string")
                })
                ds = Dataset.from_pandas(df, features=features)
                ds.to_parquet(os.path.join(sub_folder_path, "train.parquet"))
                valid_folders.append(sub_folder_name)
                print(f"    ✅ 生成 train.parquet ({len(data_rows)} 条数据)")

                config_name = f"{local_folder_name}_{sub_folder_name}".lower()
                path_str = f"{region_prefix}/{local_folder_name}/{sub_folder_name}/train.parquet"
                
                if config_name not in existing_config_names:
                    yaml_data["configs"].append({
                        "config_name": config_name,
                        "data_files": [{"split": "train", "path": path_str}]
                    })
                    new_configs_added = True
                    print(f"    ➕ 已添加配置: {config_name}")

    def get_priority(config):
        data_files = config.get("data_files", [])
        path = data_files[0].get("path", "") if data_files else ""
        # 直接使用原始大小写，确保与云端路径完全一致
        priority_dict = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}
        path_prefix = path.split('/')[0] if path else ""
        return priority_dict.get(path_prefix, 99)

    yaml_data["configs"] = sorted(yaml_data["configs"], key=lambda c: (get_priority(c), c.get("config_name", "")))
    print("  🔄 已根据 [CN > KR > JP > Fictional > WE] 的顺序重排 YAML。")

    if new_configs_added:
        card.data = yaml_data
        with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', suffix='.md', delete=False) as temp_readme:
            temp_readme.write(str(card))
            return temp_readme.name, valid_folders
    return None, valid_folders

def upload_files_and_readme(api, current_dir, region_prefix, temp_readme_path, valid_folders):
    """5. 上传文件及 README"""
    print("\n" + "="*50)
    print("🚀 [任务 5/5] 正在上传文件到 Hugging Face...")
    print("="*50)
    
    local_folder_name = os.path.basename(current_dir).replace(" ", "_").replace("-", "_")
    
    for sub_folder_name in valid_folders:
        sub_folder_path = os.path.join(current_dir, sub_folder_name)
        path_in_repo = f"{region_prefix}/{local_folder_name}/{sub_folder_name}"
        
        print(f"  📦 开始上传: {sub_folder_name} -> {path_in_repo}")
        try:
            api.upload_folder(
                folder_path=sub_folder_path,
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                path_in_repo=path_in_repo,
                commit_message=f"Update files for {path_in_repo}",
                # 上传过滤：保留 masklabel.png，jpg，txt，parquet
                allow_patterns=["*.jpg", "*.jpeg", "*.txt", "*.png", "*.parquet"]
            )
            print(f"  🎉 {sub_folder_name} 上传完成！")
        except Exception as e:
            print(f"  ❌ {sub_folder_name} 上传失败: {e}")

    if temp_readme_path:
        print("\n  📝 正在推送最新的 README.md...")
        try:
            api.upload_file(
                path_or_fileobj=temp_readme_path,
                path_in_repo="README.md",
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                commit_message="Auto-update README YAML configs"
            )
            print("  🎉 README.md 已成功更新到云端！")
        except Exception as e:
            print(f"  ❌ README.md 上传失败: {e}")
        finally:
            os.remove(temp_readme_path)

def main():
    api = HfApi(token=HF_TOKEN)
    current_dir = os.getcwd() 
    
    print("\n" + "="*50)
    print("🌟 欢迎使用 5合1 数据集自动化流水线 🌟")
    print("="*50)
    
    print("请选择该数据集的地区/分类前缀：")
    print("1. CN | 2. KR | 3. JP | 4. Fictional | 5. WE | 6. 自定义")
    choice = input("\n👉 请输入数字或直接输入前缀 (默认回车为 KR): ").strip().upper()
    if not choice: choice = "2"
    mapping = {"1": "CN", "2": "KR", "3": "JP", "4": "Fictional", "5": "WE"}
    region_prefix = mapping.get(choice, choice).replace(" ", "_").replace("-", "_")
    print(f"\n✅ 已选择前缀: {region_prefix}")

    # 任务 1：清理 txt 换行符
    clean_txt_newlines(current_dir)

    # 任务 2：清理 Mask 元数据
    remove_mask_metadata(current_dir)

    # 任务 3：转换 PNG -> JPG
    convert_png_to_jpg(current_dir)

    # 任务 4：打包 Parquet 并更新 YAML
    temp_readme_path, valid_folders = generate_parquet_and_update_yaml(api, current_dir, region_prefix)

    # 任务 5：上传
    upload_files_and_readme(api, current_dir, region_prefix, temp_readme_path, valid_folders)

    print("\n" + "="*50)
    print("🎉 5合1 全流程处理完毕！")
    print("="*50)

if __name__ == "__main__":
    main()