import sys
import os

# 确保可以导入 core 包
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from core.sp_trust_sharding import SentinelKeyEngine
from core.sp_vault_aes import SentinelVault
from core.sp_gateway_rsa import SentinelSystemProvider

def main():
    print("=== Sentinel 协议 Demo 数据流全景演示 ===\n")

    # ---------------------------------------------------------
    # 1. 密钥拆解流：身份的碎裂化 (Initialization)
    # ---------------------------------------------------------
    print("## 1. 密钥拆解流 (Initialization)")
    key_engine = SentinelKeyEngine()
    
    # 1.1 生成助记词 (BIP-39)
    master_words, entropy = key_engine.generate_vault_keys()
    print(f"   [生成] 原始助记词: {master_words}")
    
    # 1.2 SSS 分片 (3-of-2)
    shares = key_engine.split_to_shares(entropy)
    share_a = shares[0] # Device (手机)
    share_b = shares[1] # Cloud (云端)
    share_c = shares[2] # Physical (传承卡)
    
    print(f"   [分片] Share A (Device)  : {share_a}")
    print(f"   [分片] Share B (Cloud)   : {share_b}")
    print(f"   [分片] Share C (Physical): {share_c}")
    print("-" * 60)

    # ---------------------------------------------------------
    # 2. 用户内层加密流：建立私密金库 (Vault Layer)
    # ---------------------------------------------------------
    print("\n## 2. 用户内层加密流 (Vault Layer)")
    user_data = "我的瑞士银行账号是：CH123456789，密码是：Sentinel2026"
    print(f"   [输入] 用户隐私数据: {user_data}")
    
    vault = SentinelVault()
    # 2.1 派生 AES 密钥
    aes_key = vault.derive_key(master_words)
    # 2.2 加密数据
    ciphertext_1 = vault.encrypt_data(aes_key, user_data)
    print(f"   [加密] 密文 1 (AES-GCM): {ciphertext_1.hex()[:40]}... (Total len: {len(ciphertext_1)})")
    print("-" * 60)

    # ---------------------------------------------------------
    # 3. 系统外层加壳流：双重包封 (Gateway Layer)
    # ---------------------------------------------------------
    print("\n## 3. 系统外层加壳流 (Gateway Layer)")
    sys_provider = SentinelSystemProvider()
    
    # 3.1 生成系统级 RSA 密钥
    sys_private_pem, sys_public_pem = sys_provider.generate_system_keys()
    print("   [系统] 生成独立 RSA 公私钥对 (4096-bit)")
    
    # 3.2 使用系统公钥进行二次加密
    ciphertext_2 = sys_provider.encrypt_with_system_public(sys_public_pem, ciphertext_1)
    print(f"   [加密] 密文 2 (RSA-OAEP): {ciphertext_2.hex()[:40]}... (Total len: {len(ciphertext_2)})")
    print("-" * 60)

    # ---------------------------------------------------------
    # 4. 判定触发流：剥离系统外壳 (Trigger/Unlock Layer)
    # ---------------------------------------------------------
    print("\n## 4. 判定触发流 (Trigger/Unlock Layer)")
    print("   [事件] 模拟判定死亡或订阅失效，系统释放私钥权限...")
    
    # 4.1 系统私钥解密
    try:
        ciphertext_1_restored = sys_provider.decrypt_with_system_private(sys_private_pem, ciphertext_2)
        print(f"   [解密] 还原回 密文 1: {ciphertext_1_restored.hex()[:40]}...")
        
        if ciphertext_1_restored == ciphertext_1:
            print("   [验证] 密文 1 完整性校验通过 (Hash Match)")
        else:
            print("   [错误] 密文 1 校验失败")
            return
    except Exception as e:
        print(f"   [错误] 系统解密失败: {e}")
        return
    print("-" * 60)

    # ---------------------------------------------------------
    # 5. 多场景还原流：最终提取 (Restoration Scenarios)
    # ---------------------------------------------------------
    print("\n## 5. 多场景还原流 (Restoration Scenarios)")

    def run_scenario(scenario_name, share_1, share_2):
        print(f"\n   >>> 场景: {scenario_name}")
        print(f"       使用分片: {share_1[0]} 和 {share_2[0]}")
        
        # 5.1 恢复助记词
        recovered_words = key_engine.recover_from_shares(share_1, share_2)
        # print(f"       恢复助记词: {recovered_words}") 
        
        # 5.2 重新派生密钥
        restored_key = vault.derive_key(recovered_words)
        
        # 5.3 解密数据
        decrypted_text = vault.decrypt_data(restored_key, ciphertext_1_restored)
        print(f"       解密结果: {decrypted_text}")
        return decrypted_text

    # 场景 1：生前正常访问
    res1 = run_scenario("生前正常访问 (Device + Cloud)", share_a, share_b)
    assert res1 == user_data

    # 场景 2：死后标准传承
    res2 = run_scenario("死后标准传承 (Cloud + Physical)", share_b, share_c)
    assert res2 == user_data

    # 场景 3：测试验证
    res3 = run_scenario("测试验证 (Device + Physical)", share_a, share_c)
    assert res3 == user_data

    print("\n=== 演示结束: 所有流程验证通过 ===")

if __name__ == "__main__":
    main()