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

# ============ 配置区 ============
# 使用副本集连接字符串，这是实现读写分离和高可用的关键[citation:2]
REPLICA_SET_URI = "mongodb://192.168.124.10:37017,192.168.124.11:37017/?replicaSet=rs0"
DATABASE_NAME = "rw_split_test"
COLLECTION_NAME = "orders"
# ================================

def setup_connection():
    """建立具备读写分离能力的连接"""
    print("=== 1. 建立智能连接 ===")
    # read_preference=SECONDARY_PREFERRED 是关键[citation:2]：
    # 写操作会自动路由到主节点（PRIMARY）
    # 读操作会优先路由到从节点（SECONDARY），若从节点不可用则回退到主节点
    client = MongoClient(
        REPLICA_SET_URI,
        read_preference=ReadPreference.SECONDARY_PREFERRED,
        serverSelectionTimeoutMS=5000
    )
    
    try:
        # 测试连接并获取副本集信息
        client.admin.command('ping')
        print("✅ 已成功连接到副本集")
        
        # 显示当前主节点（写操作将发往此处）
        is_master = client.admin.command('isMaster')
        print(f"   当前主节点: {is_master.get('primary', '未知')}")
        print(f"   读取偏好模式: SECONDARY_PREFERRED (优先从库读)")
        
        db = client[DATABASE_NAME]
        collection = db[COLLECTION_NAME]
        # 可选：清空旧测试数据
        # collection.delete_many({})
        
        return client, db, collection
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        return None, None, None

def test_write_operations(collection):
    """测试写操作：验证是否全部发往主节点"""
    print("\n=== 2. 测试写操作 (强制发往主节点) ===")
    
    test_docs = [
        {"item": "apple", "price": 5, "quantity": 10, "timestamp": datetime.now()},
        {"item": "banana", "price": 2, "quantity": 20, "timestamp": datetime.now()},
    ]
    
    try:
        # 插入多条数据
        result = collection.insert_many(test_docs)
        print(f"✅ 插入 {len(result.inserted_ids)} 条订单数据")
        print(f"   文档ID: {result.inserted_ids}")
        
        # 更新一条数据
        update_result = collection.update_one(
            {"item": "apple"},
            {"$set": {"price": 6, "updated_at": datetime.now()}}
        )
        print(f"✅ 更新了 {update_result.modified_count} 条数据 (苹果涨价)")
        
        return result.inserted_ids
    except Exception as e:
        print(f"❌ 写操作失败: {e}")
        return []

def test_read_operations(collection, doc_ids):
    """测试读操作：展示如何在主从节点间分配"""
    print("\n=== 3. 测试读操作 (在主从节点间分配) ===")
    
    if not doc_ids:
        print("⚠ 无有效数据，跳过读测试")
        return
    
    try:
        # 测试1：普通查询（根据连接设置的 SECONDARY_PREFERRED，优先从从节点读）[citation:2]
        print("\n[测试1] 普通查询 (策略: SECONDARY_PREFERRED)")
        all_orders = list(collection.find())
        print(f"   查询到 {len(all_orders)} 条订单")
        for order in all_orders[:2]:  # 显示前两条
            print(f"     - {order['item']}: 单价{order.get('price')}元, 数量{order.get('quantity')}")
        
        # 测试2：强制从主节点读（用于需要读取最新写入数据的场景）[citation:2]
        print("\n[测试2] 强制从主节点读 (策略: PRIMARY)")
        # 通过 with_options 临时覆盖读取偏好[citation:2]
        primary_read = collection.with_options(read_preference=ReadPreference.PRIMARY)
        latest_order = primary_read.find_one(sort=[("timestamp", -1)])
        if latest_order:
            print(f"   从主节点读取的最新订单: {latest_order.get('item')}, 时间: {latest_order.get('timestamp')}")
        
        # 测试3：模拟多次读请求，观察负载分配
        print("\n[测试3] 模拟多次查询 (观察读请求分配)")
        print("   进行5次快速查询...")
        for i in range(5):
            count = collection.count_documents({})
            # 注意：由于驱动内部优化，连续快速请求可能命中同一节点
            time.sleep(0.1)  # 短暂间隔
        
        print(f"   当前总订单数: {count}")
        
    except Exception as e:
        print(f"❌ 读操作失败: {e}")

def monitor_replication_delay(client, db_name, coll_name):
    """监控主从复制延迟"""
    print("\n=== 4. 检查复制状态 ===")
    try:
        admin_db = client.admin
        # 获取副本集状态
        status = admin_db.command('replSetGetStatus')
        
        print("   成员状态:")
        for member in status['members']:
            state_str = member['stateStr']
            health = '健康' if member['health'] == 1 else '异常'
            lag = member.get('optimeDate', 'N/A')
            print(f"     - {member['name']}: {state_str} ({health})")
        
        # 检查从节点是否可读
        # 在MongoDB Shell中需要使用 rs.slaveOk()，在PyMongo中设置read_preference即可[citation:7]
        print(f"   说明: 驱动已通过 read_preference 设置处理从节点读取，无需手动执行 rs.slaveOk()")
        
    except Exception as e:
        print(f"   复制状态检查受限: {e}")

def main():
    print("=" * 60)
    print("MongoDB 副本集读写分离测试")
    print("=" * 60)
    
    client, db, collection = setup_connection()
    if not client:
        return
    
    try:
        # 测试写入
        inserted_ids = test_write_operations(collection)
        
        # 等待2秒，确保数据有足够时间复制到从节点
        print("\n⏳ 等待2秒，让数据同步到从节点...")
        time.sleep(2)
        
        # 测试读取
        test_read_operations(collection, inserted_ids)
        
        # 检查复制状态
        monitor_replication_delay(client, DATABASE_NAME, COLLECTION_NAME)
        
        # 最终验证：从从节点查询写入的数据
        print("\n=== 5. 最终验证 ===")
        # 创建一个明确只从从节点读的集合句柄（SECONDARY模式）[citation:2]
        secondary_collection = collection.with_options(read_preference=ReadPreference.SECONDARY)
        try:
            # 此查询应该成功，因为数据已复制
            verify_data = secondary_collection.find_one({"item": "banana"})
            if verify_data:
                print(f"✅ 验证成功！已从从节点读取到数据: {verify_data['item']}")
            else:
                print("⚠ 警告：暂未从从节点读取到最新数据，可能复制延迟稍高")
        except Exception as e:
            print(f"❌ 从节点读取验证失败: {e} (可能从节点尚未同步完成或不可用)")
        
    finally:
        client.close()
        print("\n" + "=" * 60)
        print("测试完成，连接已关闭")
        print("=" * 60)

if __name__ == "__main__":
    main()