from pymongo import MongoClient, ReadPreference
import time
from datetime import datetime
from collections import Counter

# ============ 配置区 ============
# 最佳实践：在连接字符串中包含所有已知节点，包括新节点。
# 这样即使部分种子节点宕机，驱动仍能发现集群。
REPLICA_SET_URIS = [
    "mongodb://192.168.124.10:37017,192.168.124.11:37017,192.168.124.15:37017/?replicaSet=rs0&socketTimeoutMS=5000",
    # 备选简洁连接，驱动可通过任一节点发现整个集群
    "mongodb://192.168.124.10:37017/?replicaSet=rs0&socketTimeoutMS=5000"
]
DATABASE_NAME = "load_test"
COLLECTION_NAME = "requests"
# ================================

def connect_and_show_topology(uri):
    """连接并显示集群拓扑，确认新节点被识别"""
    print("="*60)
    print("步骤1: 连接集群并检查拓扑状态")
    print("="*60)
    try:
        client = MongoClient(uri, read_preference=ReadPreference.SECONDARY_PREFERRED)
        # 强制发起一次网络请求以获取最新状态
        client.admin.command('ping')
        
        print("✅ 连接成功!")
        
        # 显示副本集详细信息
        rs_status = client.admin.command('replSetGetStatus')
        print(f"\n📊 副本集成员状态:")
        print("-" * 40)
        member_info = []
        for member in rs_status['members']:
            state_map = {1: 'PRIMARY', 2: 'SECONDARY', 3: 'RECOVERING', 7: 'ARBITER', 0: 'STARTUP'}
            state_str = state_map.get(member['state'], f'UNKNOWN({member["state"]})')
            health = '✅' if member['health'] == 1 else '❌'
            lag_sec = (member.get('optimeDate', datetime.now()) - datetime.now()).total_seconds()
            lag_info = f"延迟约{abs(lag_sec):.0f}秒" if member['state'] == 2 and lag_sec < 0 else ""
            info = f"{health} {member['name']:25} -> {state_str:12} {lag_info}"
            member_info.append(info)
        
        # 按角色排序显示：主节点在前
        for role_order in ['PRIMARY', 'SECONDARY', 'RECOVERING', 'STARTUP', 'ARBITER']:
            for info in member_info:
                if role_order in info:
                    print(info)
        
        # 确认新节点是否存在
        all_hosts = [m['name'] for m in rs_status['members']]
        if "192.168.124.15:37017" in all_hosts:
            print(f"\n🎉 新节点 '192.168.124.15:37017' 已在副本集成员列表中。")
        else:
            print(f"\n⚠ 警告：新节点 '192.168.124.15:37017' 未在成员列表中发现。")
            print("  请确认已执行 rs.add('192.168.124.15:37017') 且节点状态正常。")
        
        return client
    except Exception as e:
        print(f"❌ 连接或拓扑检查失败: {e}")
        return None

def test_read_distribution(client, num_requests=20):
    """测试读请求在从节点间的分布情况"""
    print("\n" + "="*60)
    print("步骤2: 测试读请求在从节点间的分布")
    print("="*60)
    
    db = client[DATABASE_NAME]
    coll = db[COLLECTION_NAME]
    
    # 准备少量测试数据
    test_doc = {"test": "read_distribution", "created_at": datetime.now()}
    coll.insert_one(test_doc)  # 此写入会到主节点
    time.sleep(2)  # 等待数据同步到所有从节点
    
    print(f"🚀 开始模拟 {num_requests} 次读请求...")
    print("   每次请求后暂停0.1秒以模拟真实间隔")
    
    # 用于统计请求落到哪个节点的计数器
    request_distribution = Counter()
    
    for i in range(num_requests):
        try:
            # 关键：通过 explain() 获取本次查询实际执行的服务器
            explain_result = coll.find({"test": "read_distribution"}).limit(1).explain()
            # MongoDB 3.6的explain格式中，服务器信息通常在 `server` 字段
            server = explain_result.get('server', 'unknown')
            # 也可能在 `serverInfo` 中，格式可能为 `host:port`
            if server == 'unknown' and 'serverInfo' in explain_result:
                server = f"{explain_result['serverInfo'].get('host')}:{explain_result['serverInfo'].get('port')}"
            
            request_distribution[server] += 1
            print(f"   请求 {i+1:2d}: 由服务器 {server} 处理")
            time.sleep(0.1)  # 短暂间隔，让分布更明显
        except Exception as e:
            print(f"   请求 {i+1:2d}: 失败 - {e}")
            request_distribution['FAILED'] += 1
    
    print("\n📈 读请求分布统计结果:")
    print("-" * 40)
    total_success = sum([count for server, count in request_distribution.items() if server != 'FAILED'])
    for server, count in request_distribution.items():
        percentage = (count / num_requests) * 100 if num_requests > 0 else 0
        if server == 'FAILED':
            print(f"   ❌ 失败请求: {count:3d} 次")
        else:
            # 判断是否为从节点 (包含 SECONDARY 字样或非主节点的IP)
            role = "从节点" if ("SECONDARY" in server or 
                              ("192.168" in server and server != get_primary_host(client))) else "主节点"
            print(f"   {server:30} ({role}): {count:3d} 次 ({percentage:5.1f}%)")
    
    if total_success > 0:
        print(f"\n💡 解读：共 {num_requests} 次读请求，已成功分发到 {len([k for k in request_distribution.keys() if k != 'FAILED'])} 个不同节点。")
        print("     在 SECONDARY_PREFERRED 模式下，驱动会自动选择它认为'最快'的健康从节点。")
        print("     短时间内连续请求可能命中同一节点，这是正常现象。")

def get_primary_host(client):
    """获取当前主节点地址"""
    try:
        is_master = client.admin.command('isMaster')
        return is_master.get('primary', 'unknown')
    except:
        return 'unknown'

def main():
    print("🧪 MongoDB 副本集扩展测试：验证新节点与读请求分配")
    print("📌 目标：确认新节点 192.168.124.15 已加入，并观察读流量分配")
    
    # 尝试多个连接字符串
    client = None
    for uri in REPLICA_SET_URIS:
        print(f"\n尝试连接: {uri[:80]}...")
        client = connect_and_show_topology(uri)
        if client:
            break
    
    if not client:
        print("\n❌ 所有连接尝试均失败，请检查网络和配置。")
        return
    
    try:
        # 测试读请求分布
        test_read_distribution(client, num_requests=100)
        
        print("\n" + "="*60)
        print("测试完成摘要")
        print("="*60)
        print("1. 请检查上方拓扑状态，确认所有节点（特别是新节点）健康。")
        print("2. 观察读请求分布，它们应至少出现在多个从节点上。")
        print("3. 这不是严格的'平分'，而是驱动的智能选择。")
        print("4. 要获得更均匀的分布，可能需要:")
        print("   - 确保各从节点网络延迟相近")
        print("   - 在应用层实现自定义路由（如按会话ID哈希）")
        print("="*60)
        
    finally:
        client.close()
        print("连接已关闭。")

if __name__ == "__main__":
    main()