"""
Base64 转图片工具
将 base64 字符串转换为图片文件
"""
import sys
import os
import base64

def main():
    if len(sys.argv) < 3:
        print("错误: 需要提供 base64 字符串和输出路径", file=sys.stderr)
        print("用法: python base64-to-img.py <base64字符串或文件路径> <输出图片路径>", file=sys.stderr)
        sys.exit(1)
    
    base64_input = sys.argv[1]
    output_path = sys.argv[2]
    
    try:
        # 判断输入是文件路径还是 base64 字符串
        base64_data = None
        
        if os.path.isfile(base64_input):
            # 如果是文件路径，读取文件内容
            with open(base64_input, 'r', encoding='utf-8') as f:
                base64_data = f.read().strip()
        else:
            # 否则当作 base64 字符串处理
            base64_data = base64_input.strip()
        
        # 清理 Base64 数据：移除可能存在的 data URL 前缀
        if ',' in base64_data:
            # 如果包含逗号，说明可能有 data:image/...;base64, 前缀
            base64_index = base64_data.index(',')
            base64_data = base64_data[base64_index + 1:]
        
        base64_data = base64_data.strip()
        
        # 解码 base64
        try:
            image_data = base64.b64decode(base64_data)
        except Exception as e:
            print(f"错误: Base64 解码失败 - {e}", file=sys.stderr)
            sys.exit(1)
        
        # 创建输出目录（如果不存在）
        output_dir = os.path.dirname(output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)
        
        # 写入图片文件
        with open(output_path, 'wb') as f:
            f.write(image_data)
        
        print(f"成功: Base64 已转换为图片: {output_path}")
        
    except Exception as error:
        print(f"错误: {error}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == '__main__':
    main()

