#!/usr/bin/env python3
"""
测试环境变量配置是否正确
"""

import os
import sys
from datetime import datetime

def check_env_var(var_name, required=True):
    """检查环境变量"""
    value = os.environ.get(var_name)
    if required and not value:
        print(f"❌ {var_name}: 未设置（必须）")
        return False
    elif value:
        # 隐藏敏感信息
        if "KEY" in var_name or "TOKEN" in var_name:
            display_value = f"{value[:20]}...{value[-10:]}" if len(value) > 30 else "***"
        else:
            display_value = value
        print(f"✅ {var_name}: {display_value}")
        return True
    else:
        print(f"⚠️  {var_name}: 未设置（可选）")
        return True

def check_jwt_token(token):
    """检查 JWT token 是否包含必要字段"""
    try:
        import base64
        import json

        if token.startswith("Bearer "):
            token = token[7:]

        parts = token.split(".")
        if len(parts) != 3:
            print(f"❌ Token 格式错误：应该是 JWT 格式")
            return False

        # 解码 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))

        print("\nJWT Token 信息:")
        print(f"  Issuer (iss): {payload.get('iss')}")
        print(f"  Subject (sub): {payload.get('sub')}")
        print(f"  Expires (exp): {datetime.fromtimestamp(payload.get('exp', 0))}")
        print(f"  Issued At (iat): {datetime.fromtimestamp(payload.get('iat', 0))}")

        # 检查是否过期
        import time
        if payload.get('exp', 0) < time.time():
            print(f"❌ Token 已过期")
            return False
        else:
            print(f"✅ Token 未过期")

        # 提取 workspace_id
        sub = payload.get('sub', '')
        if 'workload_identity/id:' in sub:
            workspace_id = sub.split('workload_identity/id:')[-1]
            print(f"\n✅ 从 Token 提取的 Workspace ID: {workspace_id}")
            return workspace_id
        else:
            print(f"⚠️  无法从 Token 提取 Workspace ID")

        return True

    except Exception as e:
        print(f"❌ 解析 Token 失败: {str(e)}")
        return False

def main():
    print("=" * 80)
    print("环境变量配置检查")
    print("=" * 80)
    print()

    # 检查关键环境变量
    print("1. 检查环境变量:")
    print("-" * 80)
    all_ok = True
    all_ok &= check_env_var("COZE_API_KEY", required=True)
    all_ok &= check_env_var("COZE_WORKSPACE_ID", required=True)
    all_ok &= check_env_var("COZE_INTEGRATION_BASE_URL", required=True)
    all_ok &= check_env_var("LLM_MODEL_NAME", required=True)
    print()

    # 检查 JWT token
    print("2. 检查 JWT Token:")
    print("-" * 80)
    token = os.environ.get("COZE_API_KEY", "")
    if token:
        check_jwt_token(token)
    print()

    # 检查端点配置
    print("3. 检查 API 端点配置:")
    print("-" * 80)
    base_url = os.environ.get("COZE_INTEGRATION_BASE_URL", "")
    if base_url:
        if "/v1" in base_url:
            print(f"✅ API 端点包含 /v1 前缀: {base_url}")
        else:
            print(f"❌ API 端点缺少 /v1 前缀: {base_url}")
            all_ok = False
    else:
        print(f"❌ API 端点未设置")
        all_ok = False
    print()

    # 总结
    print("=" * 80)
    print("检查总结")
    print("=" * 80)
    if all_ok:
        print("✅ 所有必需的环境变量都已正确配置")
        print("\n下一步:")
        print("1. 重新构建 Docker 镜像: docker-compose build")
        print("2. 启动服务: docker-compose up -d")
        print("3. 检查日志: docker-compose logs -f")
    else:
        print("❌ 存在配置问题，请修复后再试")
        print("\n修复方法:")
        print("1. 设置 COZE_API_KEY 环境变量")
        print("2. 设置 COZE_WORKSPACE_ID 环境变量")
        print("3. 确保 API 端点包含 /v1 前缀")

    return 0 if all_ok else 1

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