# req_embed_encryption_example.py
import numpy as np
import json
import bytedance.jeddak_secure_channel as jsc
from bytedance.jeddak_crypto_rag.crypto_rag_client import CryptoRAGClient

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='Process some authentication information.')
    parser.add_argument('--account_id', type=str, help='火山主账号ID', required=True)
    parser.add_argument('--ak', type=str, help='火山访问密钥AK', required=True)
    parser.add_argument('--sk', type=str, help='火山访问密钥SK', required=True)
    parser.add_argument('--app_id', type=str, required=True, help='向量加密应用ID')
    parser.add_argument('--policy_id', type=str, help='机密传输的策略ID')
    parser.add_argument('--url', type=str, required=True, help='向量加密服务的访问域名')
    args = parser.parse_args()

    server_url = args.url
    rag_app_id = args.app_id
    account_id = args.account_id
    # 准备机密通信client的配置，其中ra_pods_info是向量加密服务的信息，用于验证向量加密服务的环境。
    secure_channel_config = jsc.ClientConfig.from_dict({
        "ra_url": "open.volcengineapi.com",
        "ra_service_name": "PCC.CryptoRag",
        "ra_policy_id": args.policy_id,
        "attest_interval": 1800,
        "ra_uid": account_id,
        "bytedance_top_info": json.dumps({
            "ak": args.ak,
            "sk": args.sk,
            "service": "pcc"
        })
    })
    # 初始化机密通信client，并进行远程证明
    secure_channel_client = jsc.Client(secure_channel_config)
    secure_channel_client.attest_server()
    # 生成随机待加密向量
    embeddings = np.random.uniform(-1, 1, size=(5, 1024)).tolist()
    print(f"account_id: {account_id}, rag_app_id: {rag_app_id}, server_url: {server_url}")
    # 初始化向量加密client
    crypto_rag_client = CryptoRAGClient(
        account_id,
        rag_app_id,
        server_url,
        secure_channel_client
    )
    # 请求向量加密，设置downgrad_by_plain_channel=False表示机密通信失败时不降级为明文通信
    encrypted_dense_vectors = crypto_rag_client.req_embed_encryption_secure_channel(
        embeddings=embeddings, 
        disable_ra_during_downgrade=False)
    
    print("========== post embedding encryption success ==========")
    print("明文向量（前10位）: ", embeddings[0][:10])
    print("密文向量（前10位）: ", encrypted_dense_vectors[0][:10])