#!/usr/bin/env python3
"""
企业微信会话存档SDK命令行工具
Python重构版本，替代原始的tool_testSdk.cpp
"""

import argparse
import json
import sys
import os
from pathlib import Path
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 添加父目录到路径，以便导入wework_ctypes
sys.path.insert(0, str(Path(__file__).parent))

from wework_ctypes import WeWorkFinanceSDK, WeWorkSDKError

def parse_args():
    """解析命令行参数"""
    # 从环境变量获取默认值
    default_corpid = os.getenv("CORPID", "wwdf65802ca25ec195")
    default_secret = os.getenv("WEWORK_SECRET", "-Ta6WMWxBhfGolWnnlO15nQckj3DRKAowUOdX2fwvzE")
    
    parser = argparse.ArgumentParser(
        description="企业微信会话存档SDK工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 拉取会话存档
  %(prog)s chatmsg 0 100 "" "" 5
  
  # 拉取媒体文件
  %(prog)s mediadata "fileid" "" "" 5 "/tmp/media.file"
  
  # 解密数据
  %(prog)s decryptdata "encrypt_key" "encrypt_chat_msg"
  
  # 使用代理拉取会话存档
  %(prog)s chatmsg 0 100 "http://proxy.example.com:8080" "user:pass" 10
        """
    )
    
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # chatmsg命令
    chatmsg_parser = subparsers.add_parser('chatmsg', help='拉取会话存档')
    chatmsg_parser.add_argument('seq', type=int, help='从指定的seq开始拉取消息，首次使用请使用0')
    chatmsg_parser.add_argument('limit', type=int, help='一次拉取的消息条数，最大值1000条')
    chatmsg_parser.add_argument('proxy', type=str, help='代理链接，如：socks5://10.0.0.1:8081 或 http://10.0.0.1:8081')
    chatmsg_parser.add_argument('passwd', type=str, help='代理账号密码，如：user_name:passwd_123')
    chatmsg_parser.add_argument('timeout', type=int, help='超时时间，单位秒')
    
    # mediadata命令
    mediadata_parser = subparsers.add_parser('mediadata', help='拉取媒体文件')
    mediadata_parser.add_argument('fileid', type=str, help='媒体文件id，从解密后的会话存档中得到')
    mediadata_parser.add_argument('proxy', type=str, help='代理链接')
    mediadata_parser.add_argument('passwd', type=str, help='代理账号密码')
    mediadata_parser.add_argument('timeout', type=int, help='超时时间，单位秒')
    mediadata_parser.add_argument('savefile', type=str, help='媒体文件保存路径')
    
    # decryptdata命令
    decryptdata_parser = subparsers.add_parser('decryptdata', help='解密会话存档内容')
    decryptdata_parser.add_argument('encrypt_key', type=str, help='拉取会话存档返回的encrypt_random_key，使用配置在企业微信管理台的rsa公钥对应的私钥解密后得到encrypt_key')
    decryptdata_parser.add_argument('encrypt_chat_msg', type=str, help='拉取会话存档返回的encrypt_chat_msg')
    
    # 添加全局选项
    parser.add_argument('--corpid', type=str, 
                       default=default_corpid,
                       help=f'企业ID，默认从环境变量CORPID读取 (当前: {default_corpid})')
    parser.add_argument('--secret', type=str,
                       default=default_secret,
                       help=f'聊天内容存档的Secret，默认从环境变量WEWORK_SECRET读取')
    parser.add_argument('--lib-path', type=str,
                       help='libWeWorkFinanceSdk_C.so路径，默认使用lib/sdk/libWeWorkFinanceSdk_C.so')
    
    return parser.parse_args()

def main():
    """主函数"""
    args = parse_args()
    
    if not args.command:
        print("错误: 需要指定命令 (chatmsg, mediadata, decryptdata)")
        print("使用 --help 查看帮助")
        return 1
    
    try:
        # 初始化SDK
        with WeWorkFinanceSDK(args.corpid, args.secret, args.lib_path) as sdk:
            
            if args.command == 'chatmsg':
                # 拉取会话存档
                print(f"拉取会话存档: seq={args.seq}, limit={args.limit}, timeout={args.timeout}s")
                if args.proxy:
                    print(f"使用代理: {args.proxy}")
                
                result = sdk.get_chat_data(
                    seq=args.seq,
                    limit=args.limit,
                    proxy=args.proxy,
                    passwd=args.passwd,
                    timeout=args.timeout
                )
                
                # 尝试解析JSON并美化输出
                try:
                    data = json.loads(result)
                    print("获取聊天数据成功:")
                    print(json.dumps(data, indent=2, ensure_ascii=False))
                    
                    # 如果有聊天数据，显示统计信息
                    if 'chatdata' in data and data['chatdata']:
                        print(f"\n共获取 {len(data['chatdata'])} 条聊天记录")
                        if data['chatdata']:
                            max_seq = max(msg.get('seq', 0) for msg in data['chatdata'])
                            print(f"最大seq值: {max_seq} (下次拉取建议使用此值)")
                    
                except json.JSONDecodeError:
                    print("获取聊天数据成功 (原始数据):")
                    print(result)
                
                return 0
                
            elif args.command == 'mediadata':
                # 拉取媒体文件
                print(f"拉取媒体文件: fileid={args.fileid}, timeout={args.timeout}s")
                print(f"保存到: {args.savefile}")
                if args.proxy:
                    print(f"使用代理: {args.proxy}")
                
                success = sdk.get_media_data(
                    sdk_fileid=args.fileid,
                    save_path=args.savefile,
                    proxy=args.proxy,
                    passwd=args.passwd,
                    timeout=args.timeout
                )
                
                if success:
                    print(f"媒体文件下载完成: {args.savefile}")
                    
                    # 检查文件大小
                    if os.path.exists(args.savefile):
                        size = os.path.getsize(args.savefile)
                        print(f"文件大小: {size} 字节 ({size/1024:.2f} KB)")
                else:
                    print("媒体文件下载失败")
                    return 1
                
                return 0
                
            elif args.command == 'decryptdata':
                # 解密会话存档内容
                print("解密会话存档内容...")
                
                result = sdk.decrypt_data(args.encrypt_key, args.encrypt_chat_msg)
                
                # 尝试解析JSON并美化输出
                try:
                    data = json.loads(result)
                    print("解密成功:")
                    print(json.dumps(data, indent=2, ensure_ascii=False))
                except json.JSONDecodeError:
                    print("解密成功 (原始数据):")
                    print(result)
                
                return 0
                
    except WeWorkSDKError as e:
        print(f"SDK错误: {e}")
        return 1
    except FileNotFoundError as e:
        print(f"文件错误: {e}")
        return 1
    except Exception as e:
        print(f"未知错误: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())