#!/usr/bin/env python3
"""
从 JWT Token 中提取 Workspace ID
"""

import json
import base64
import sys

def extract_workspace_id(token):
    """
    从 JWT Token 中提取 Workspace ID
    
    Args:
        token: JWT Token（带或不带 "Bearer " 前缀）
    
    Returns:
        workspace_id: 工作区 ID
    """
    try:
        # 移除 "Bearer " 前缀
        if token.startswith("Bearer "):
            token = token[7:]
        
        # 分割 JWT
        parts = token.split(".")
        if len(parts) != 3:
            print("❌ Token 格式错误：应该是 JWT 格式（header.payload.signature）")
            return None
        
        # 解码 payload
        payload_b64 = parts[1]
        # 添加 padding
        padding = 4 - len(payload_b64) % 4
        if padding != 4:
            payload_b64 += "=" * padding
        
        payload = json.loads(base64.b64decode(payload_b64))
        
        # 显示 Token 信息
        print("\n" + "=" * 80)
        print("JWT Token 信息")
        print("=" * 80)
        print(f"Issuer (iss): {payload.get('iss')}")
        print(f"Audience (aud): {payload.get('aud')}")
        print(f"Subject (sub): {payload.get('sub')}")
        print(f"Expires (exp): {payload.get('exp')}")
        print(f"Issued At (iat): {payload.get('iat')}")
        print(f"Source (src): {payload.get('src')}")
        print()
        
        # 提取 workspace_id
        sub = payload.get('sub', '')
        
        # 尝试不同的格式
        if 'workload_identity/id:' in sub:
            workspace_id = sub.split('workload_identity/id:')[-1]
            print(f"✅ 格式 1: workload_identity/id:{workspace_id}")
        elif 'workspace/id:' in sub:
            workspace_id = sub.split('workspace/id:')[-1]
            print(f"✅ 格式 2: workspace/id:{workspace_id}")
        else:
            # 尝试直接使用 sub
            workspace_id = sub
            print(f"⚠️  格式 3: 直接使用 sub: {workspace_id}")
        
        print()
        print("=" * 80)
        print(f"🔑 提取的 Workspace ID: {workspace_id}")
        print("=" * 80)
        print()
        
        return workspace_id
        
    except Exception as e:
        print(f"❌ 解析 Token 失败: {str(e)}")
        return None

def main():
    print()
    print("=" * 80)
    print("Workspace ID 提取工具")
    print("=" * 80)
    print()
    
    # 方式 1：从环境变量读取
    token = None
    
    # 1. 尝试从命令行参数读取
    if len(sys.argv) > 1:
        token = sys.argv[1]
    # 2. 尝试从环境变量读取
    else:
        token = __import__('os').environ.get('COZE_API_KEY')
    
    if not token:
        print("❌ 未找到 Token")
        print()
        print("使用方法：")
        print("  方式 1: python extract_workspace_id.py 'Bearer eyJhbGci...'")
        print("  方式 2: COZE_API_KEY='Bearer eyJhbGci...' python extract_workspace_id.py")
        print()
        return 1
    
    # 提取 workspace_id
    workspace_id = extract_workspace_id(token)
    
    if workspace_id:
        print()
        print("✅ 提取成功！")
        print()
        print("下一步：")
        print(f"1. 在 docker-compose.yml 中设置：")
        print(f"   COZE_WORKSPACE_ID={workspace_id}")
        print()
        return 0
    else:
        print()
        print("❌ 提取失败，请检查 Token 格式")
        print()
        return 1

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