#!/usr/bin/env python3
"""
非计费模块代码混淆脚本

使用PyArmor对非计费相关的代码进行混淆，保护核心业务逻辑
保留计费模块的可读性，便于客户集成和维护
"""
import os
import shutil
import subprocess
from pathlib import Path
from typing import List, Set

# 计费相关的核心文件（不混淆）
BILLING_FILES = {
    # Models
    "backend/app/models/billing.py",
    "backend/app/models/consumption.py",
    "backend/app/models/invoice.py",
    "backend/app/models/invoice_info.py",
    "backend/app/models/invoice_info_history.py",
    "backend/app/models/user.py",
    
    # Services
    "backend/app/services/billing_service.py",
    "backend/app/services/balance_service.py",
    "backend/app/services/recharge_service.py",
    "backend/app/services/billing_calculator.py",
    "backend/app/services/recharge_handler.py",
    
    # Routers
    "backend/app/routers/billing_router.py",
    "backend/app/routers/invoice_router.py",
    
    # Schemas
    "backend/app/schemas/billing_schema.py",
    "backend/app/schemas/invoice_schema.py",
    
    # Dependencies
    "backend/app/dependencies/balance_check.py",
    
    # Core
    "backend/app/core/config.py",
    "backend/app/core/async_database.py",
    "backend/app/core/async_logger.py",
    "backend/app/core/redis.py",
    "backend/app/database.py",
    
    # Middleware (部分需要保留)
    "backend/app/middleware/auth_log_middleware.py",
    "backend/app/middleware/error_handler.py",
}

# 需要混淆的非计费模块目录
NON_BILLING_MODULES = [
    "backend/app/models/conversation.py",
    "backend/app/models/ai_picture.py",
    "backend/app/models/ai_video.py",
    "backend/app/models/audio.py",
    "backend/app/models/image_translation.py",
    "backend/app/models/translation.py",
    "backend/app/models/zhiwen.py",
    "backend/app/models/research.py",
    "backend/app/models/ocr.py",
    "backend/app/models/photo_answer.py",
    "backend/app/models/tingwu.py",
    "backend/app/models/toolbox.py",
    "backend/app/models/search.py",
    "backend/app/models/api_call_log.py",
    "backend/app/models/platform_api_key.py",
    "backend/app/models/parsed_pricing.py",
    "backend/app/models/model.py",
    "backend/app/models/config.py",
    "backend/app/models/log.py",
    "backend/app/models/admin.py",
    
    # Services
    "backend/app/services/llm_service.py",
    "backend/app/services/image_service.py",
    "backend/app/services/video_service.py",
    "backend/app/services/audio_billing.py",
    "backend/app/services/image_billing.py",
    "backend/app/services/video_billing.py",
    "backend/app/services/translation_billing.py",
    "backend/app/services/image_translation_billing.py",
    "backend/app/services/ocr_billing_calculator.py",
    "backend/app/services/translation_service.py",
    "backend/app/services/image_translation_service.py",
    "backend/app/services/multimodal_translation_service.py",
    "backend/app/services/zhiwen_service.py",
    "backend/app/services/zhiwen_client.py",
    "backend/app/services/research_service.py",
    "backend/app/services/qwen_deep_research_service.py",
    "backend/app/services/ocr_service.py",
    "backend/app/services/qwen_ocr.py",
    "backend/app/services/photo_answer_service.py",
    "backend/app/services/tingwu_service.py",
    "backend/app/services/tingwu_client.py",
    "backend/app/services/tingwu_realtime_client.py",
    "backend/app/services/tingwu_params_adapter.py",
    "backend/app/services/asr_service.py",
    "backend/app/services/tts_service.py",
    "backend/app/services/voice_clone_service.py",
    "backend/app/services/system_voice_service.py",
    "backend/app/services/conversation_service.py",
    "backend/app/services/message_service.py",
    "backend/app/services/model_service.py",
    "backend/app/services/local_model_service.py",
    "backend/app/services/openai_compat_service.py",
    "backend/app/services/dashscope_client.py",
    "backend/app/services/anytrans_client.py",
    "backend/app/services/edututor_client.py",
    "backend/app/services/admin_auth_service.py",
    "backend/app/services/admin_user_service.py",
    "backend/app/services/admin_model_service.py",
    "backend/app/services/admin_stats_service.py",
    "backend/app/services/admin_order_service.py",
    "backend/app/services/admin_bill_service.py",
    "backend/app/services/admin_consumption_service.py",
    "backend/app/services/admin_invoice_service.py",
    "backend/app/services/auth_service.py",
    "backend/app/services/user_service.py",
    "backend/app/services/config_service.py",
    "backend/app/services/log_service.py",
    "backend/app/services/operation_log_service.py",
    "backend/app/services/api_call_log_service.py",
    "backend/app/services/platform_api_key_service.py",
    "backend/app/services/parsed_pricing_service.py",
    "backend/app/services/consumption_sync_service.py",
    "backend/app/services/review_service.py",
    "backend/app/services/toolbox_service.py",
    "backend/app/services/oss_service.py",
    "backend/app/services/avatar_service.py",
    "backend/app/services/cache_service.py",
    "backend/app/services/rate_limiter.py",
    "backend/app/services/health_service.py",
    "backend/app/services/system_config_manager.py",
    "backend/app/services/crypto_utils.py",
    "backend/app/services/token_revocation_service.py",
    "backend/app/services/phrase_service.py",
    "backend/app/services/terminology_service.py",
    "backend/app/services/translation_memory_service.py",
    "backend/app/services/citation_formatter.py",
    "backend/app/services/search_options_validator.py",
    "backend/app/services/stream_search_state.py",
    "backend/app/services/vertical_domain_processor.py",
]


def check_pyarmor_installed():
    """检查PyArmor是否已安装"""
    try:
        result = subprocess.run(
            ["pyarmor", "--version"],
            capture_output=True,
            text=True,
            check=True
        )
        print(f"✓ PyArmor已安装: {result.stdout.strip()}")
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("✗ PyArmor未安装")
        print("\n请先安装PyArmor:")
        print("  pip install pyarmor")
        return False


def create_output_directory():
    """创建输出目录"""
    output_dir = Path("backend_obfuscated")
    if output_dir.exists():
        print(f"清理旧的输出目录: {output_dir}")
        shutil.rmtree(output_dir)
    
    output_dir.mkdir(parents=True, exist_ok=True)
    print(f"✓ 创建输出目录: {output_dir}")
    return output_dir


def copy_billing_files(output_dir: Path):
    """复制计费相关文件（不混淆）"""
    print("\n正在复制计费模块文件（保持可读）...")
    
    for file_path in BILLING_FILES:
        src = Path(file_path)
        if not src.exists():
            print(f"  ⚠ 文件不存在，跳过: {file_path}")
            continue
        
        # 计算目标路径
        dst = output_dir / src
        dst.parent.mkdir(parents=True, exist_ok=True)
        
        # 复制文件
        shutil.copy2(src, dst)
        print(f"  ✓ 复制: {file_path}")
    
    print(f"✓ 已复制 {len(BILLING_FILES)} 个计费模块文件")


def obfuscate_non_billing_files(output_dir: Path):
    """混淆非计费模块文件"""
    print("\n正在混淆非计费模块...")
    
    obfuscated_count = 0
    failed_files = []
    
    for file_path in NON_BILLING_MODULES:
        src = Path(file_path)
        if not src.exists():
            print(f"  ⚠ 文件不存在，跳过: {file_path}")
            continue
        
        # 计算目标路径
        dst = output_dir / src.parent
        dst.mkdir(parents=True, exist_ok=True)
        
        try:
            # 使用PyArmor混淆单个文件
            result = subprocess.run(
                [
                    "pyarmor",
                    "obfuscate",
                    "--output", str(dst),
                    "--exact",  # 只混淆指定文件
                    "--no-cross-protection",  # 不启用交叉保护
                    str(src)
                ],
                capture_output=True,
                text=True,
                check=True
            )
            print(f"  ✓ 混淆: {file_path}")
            obfuscated_count += 1
        except subprocess.CalledProcessError as e:
            print(f"  ✗ 混淆失败: {file_path}")
            print(f"    错误: {e.stderr}")
            failed_files.append(file_path)
    
    print(f"\n✓ 成功混淆 {obfuscated_count} 个文件")
    if failed_files:
        print(f"✗ 失败 {len(failed_files)} 个文件:")
        for f in failed_files:
            print(f"  - {f}")
    
    return obfuscated_count, failed_files


def copy_other_files(output_dir: Path):
    """复制其他必要文件（配置、依赖等）"""
    print("\n正在复制其他必要文件...")
    
    other_files = [
        "backend/requirements.txt",
        "backend/.env.example",
        "backend/README.md",
        "backend/app/__init__.py",
        "backend/app/main.py",
    ]
    
    for file_path in other_files:
        src = Path(file_path)
        if not src.exists():
            continue
        
        dst = output_dir / src
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)
        print(f"  ✓ 复制: {file_path}")


def copy_init_files(output_dir: Path):
    """复制所有__init__.py文件"""
    print("\n正在复制__init__.py文件...")
    
    backend_dir = Path("backend")
    for init_file in backend_dir.rglob("__init__.py"):
        relative_path = init_file.relative_to(backend_dir.parent)
        dst = output_dir / relative_path
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(init_file, dst)
        print(f"  ✓ 复制: {relative_path}")


def create_readme(output_dir: Path):
    """创建README说明文件"""
    readme_content = """# 计费模块交付包

## 说明

本包包含完整的计费系统代码，其中：

### 可读模块（未混淆）
- **计费核心模块**: `app/models/billing.py`, `app/models/consumption.py`, `app/models/invoice.py`
- **计费服务**: `app/services/billing_service.py`, `app/services/balance_service.py`, `app/services/recharge_service.py`
- **计费API**: `app/routers/billing_router.py`, `app/routers/invoice_router.py`
- **用户模型**: `app/models/user.py`
- **核心配置**: `app/core/`, `app/database.py`

### 混淆模块（已保护）
- **业务模块**: AI对话、AI生图、AI视频、语音合成、语音识别等
- **业务服务**: 各业务模块的服务层实现
- **管理后台**: 管理员相关功能

## 部署说明

1. 安装依赖：
   ```bash
   pip install -r requirements.txt
   ```

2. 配置环境变量：
   ```bash
   cp .env.example .env
   # 编辑.env文件，配置数据库、支付宝等参数
   ```

3. 运行数据库迁移：
   ```bash
   alembic upgrade head
   ```

4. 启动服务：
   ```bash
   uvicorn app.main:app --host 0.0.0.0 --port 8000
   ```

## 计费模块文档

详细的计费模块文档请参考：`.kiro/specs/billing-system-documentation/design.md`

## 技术支持

如有问题，请联系技术支持团队。

---
生成时间: {datetime}
混淆工具: PyArmor
"""
    
    from datetime import datetime
    readme_path = output_dir / "README.md"
    readme_path.write_text(
        readme_content.format(datetime=datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
        encoding="utf-8"
    )
    print(f"\n✓ 创建README: {readme_path}")


def main():
    """主函数"""
    print("=" * 60)
    print("非计费模块代码混淆工具")
    print("=" * 60)
    
    # 1. 检查PyArmor
    if not check_pyarmor_installed():
        return
    
    # 2. 创建输出目录
    output_dir = create_output_directory()
    
    # 3. 复制计费文件（不混淆）
    copy_billing_files(output_dir)
    
    # 4. 混淆非计费文件
    obfuscated_count, failed_files = obfuscate_non_billing_files(output_dir)
    
    # 5. 复制其他文件
    copy_other_files(output_dir)
    copy_init_files(output_dir)
    
    # 6. 创建README
    create_readme(output_dir)
    
    # 7. 总结
    print("\n" + "=" * 60)
    print("混淆完成！")
    print("=" * 60)
    print(f"输出目录: {output_dir.absolute()}")
    print(f"计费模块文件: {len(BILLING_FILES)} 个（未混淆）")
    print(f"非计费模块文件: {obfuscated_count} 个（已混淆）")
    
    if failed_files:
        print(f"\n⚠ 警告: {len(failed_files)} 个文件混淆失败")
    
    print("\n下一步:")
    print("1. 检查输出目录中的文件")
    print("2. 测试混淆后的代码是否正常运行")
    print("3. 将整个目录打包交付给客户")


if __name__ == "__main__":
    main()