from pymongo import MongoClient, errors
from datetime import datetime
import time

# ============ 配置区 ============
# 连接整个副本集（推荐的生产方式）
REPLICA_SET_URI = "mongodb://192.168.124.10:37017,192.168.124.11:37017/?replicaSet=rs0"
# 或直接连接从库（仅用于测试）
SECONDARY_DIRECT_URI = "mongodb://192.168.124.11:37017/"

DATABASE_NAME = "replica_test"
COLLECTION_NAME = "python_test"
# ================================

def test_replica_set_connection():
    """测试1：连接整个副本集，并检查成员状态"""
    print("=== 测试1：连接整个副本集 ===")
    try:
        # readPreference=secondaryPreferred 表示优先读从库
        client = MongoClient(
            REPLICA_SET_URI,
            readPreference='secondaryPreferred',
            serverSelectionTimeoutMS=5000
        )
        
        # 强制发起一次网络请求来测试连接
        client.admin.command('ping')
        print("✓ 成功连接到副本集")
        
        # 获取副本集状态信息
        rs_status = client.admin.command('replSetGetStatus')
        print(f"  副本集名称: {rs_status['set']}")
        
        # 打印成员信息
        print("  成员状态:")
        for member in rs_status['members']:
            state_str = ['STARTUP', 'PRIMARY', 'SECONDARY', 'RECOVERING', 'STARTUP2', 'UNKNOWN', 'ARBITER', 'DOWN'][member['state']] if member['state'] <= 7 else str(member['state'])
            print(f"    - {member['name']}: {state_str} (健康: {member['health']})")
        
        # 检查当前主节点
        is_master = client.admin.command('isMaster')
        print(f"  当前主节点: {is_master.get('primary', '未知')}")
        
        client.close()
        return True
        
    except errors.ServerSelectionTimeoutError as e:
        print(f"✗ 连接超时: {e}")
        return False
    except errors.ConnectionFailure as e:
        print(f"✗ 连接失败: {e}")
        return False

def test_write_to_primary():
    """测试2：向主节点写入数据"""
    print("\n=== 测试2：向主节点写入数据 ===")
    
    # 注意：写操作必须连接到主节点，所以这里不使用readPreference
    client = MongoClient(REPLICA_SET_URI, serverSelectionTimeoutMS=5000)
    
    try:
        db = client[DATABASE_NAME]
        collection = db[COLLECTION_NAME]
        
        # 准备测试数据
        test_doc = {
            "message": "Python客户端写入的测试数据",
            "timestamp": datetime.now(),
            "counter": 1,
            "source": "PRIMARY"
        }
        
        # 插入数据
        result = collection.insert_one(test_doc)
        doc_id = result.inserted_id
        print(f"✓ 数据写入成功，文档ID: {doc_id}")
        print(f"  写入内容: {test_doc}")
        
        # 立即从主节点读取验证（强一致性读）
        immediate_read = collection.find_one({"_id": doc_id})
        print(f"✓ 主节点立即读取验证: {immediate_read['message'] if immediate_read else '未找到'}")
        
        client.close()
        return doc_id
        
    except Exception as e:
        print(f"✗ 写入失败: {e}")
        return None

def test_read_from_secondary(doc_id):
    """测试3：从从节点读取数据"""
    print("\n=== 测试3：从从节点读取数据 ===")
    
    if not doc_id:
        print("⚠ 跳过从节点读取测试（无有效文档ID）")
        return
    
    # 方法1：通过副本集连接，设置读偏好
    print("方法1: 通过副本集连接（推荐）")
    try:
        client = MongoClient(
            REPLICA_SET_URI,
            readPreference='secondaryPreferred',  # 优先从从库读
            serverSelectionTimeoutMS=5000
        )
        
        db = client[DATABASE_NAME]
        collection = db[COLLECTION_NAME]
        
        # 可能需要等待数据同步
        max_retries = 5
        for i in range(max_retries):
            doc = collection.find_one({"_id": doc_id})
            if doc:
                print(f"✓ 第{i+1}次尝试: 从从节点读取成功")
                print(f"  读取内容: {doc['message']}")
                print(f"  来源: {doc.get('source', '未知')}")
                break
            else:
                print(f"  ... 第{i+1}次尝试: 数据尚未同步，等待1秒")
                time.sleep(1)
        else:
            print("✗ 数据在从节点上未找到（可能同步延迟）")
        
        client.close()
        
    except Exception as e:
        print(f"✗ 方法1读取失败: {e}")
    
    # 方法2：直接连接从库
    print("\n方法2: 直接连接从库")
    try:
        # 注意：直接连接从库时，需要设置readPreference或者处理可能的连接错误
        client = MongoClient(
            SECONDARY_DIRECT_URI,
            readPreference='secondary',
            serverSelectionTimeoutMS=3000
        )
        
        # 测试从库连接
        client.admin.command('ping')
        print("✓ 直接连接到从库成功")
        
        db = client[DATABASE_NAME]
        # MongoDB 3.6需要在从库上执行rs.slaveOk()的等效操作
        # 对于PyMongo，设置readPreference='secondary'即可
        
        doc = db[COLLECTION_NAME].find_one({"_id": doc_id})
        if doc:
            print(f"✓ 直接读取从库数据成功")
            print(f"  内容: {doc['message']}")
        else:
            print("✗ 直接读取从库未找到数据")
        
        client.close()
        
    except errors.ServerSelectionTimeoutError:
        print("✗ 直接连接从库超时（可能从库不可用）")
    except Exception as e:
        print(f"✗ 直接读取失败: {e}")

def test_data_sync():
    """测试4：完整的数据同步测试"""
    print("\n=== 测试4：完整数据同步测试 ===")
    
    # 先写入一批数据
    client = MongoClient(REPLICA_SET_URI)
    db = client[DATABASE_NAME]
    collection = db[COLLECTION_NAME]
    
    # 清理旧测试数据（可选）
    # collection.delete_many({"source": "SYNC_TEST"})
    
    # 插入多条测试数据
    test_data = []
    for i in range(5):
        doc = {
            "batch_id": "SYNC_TEST",
            "number": i,
            "message": f"同步测试数据 #{i}",
            "timestamp": datetime.now(),
            "source": "SYNC_TEST"
        }
        test_data.append(doc)
    
    result = collection.insert_many(test_data)
    print(f"✓ 插入 {len(result.inserted_ids)} 条测试数据")
    
    # 等待2秒确保数据同步
    print("  等待2秒确保数据同步...")
    time.sleep(2)
    
    # 从从节点读取
    secondary_client = MongoClient(
        REPLICA_SET_URI,
        readPreference='secondaryPreferred'
    )
    secondary_collection = secondary_client[DATABASE_NAME][COLLECTION_NAME]
    
    # 统计从节点上的数据
    sync_count = secondary_collection.count_documents({"batch_id": "SYNC_TEST"})
    print(f"  从节点上的数据数量: {sync_count}/5")
    
    if sync_count == 5:
        print("✓ 数据同步完整")
    elif sync_count > 0:
        print(f"⚠ 部分数据同步: {sync_count}/5")
    else:
        print("✗ 数据未同步")
    
    # 显示同步的数据
    if sync_count > 0:
        print("  同步的数据样本:")
        for doc in secondary_collection.find({"batch_id": "SYNC_TEST"}).limit(2):
            print(f"    - {doc['message']}")
    
    client.close()
    secondary_client.close()

def main():
    print("=" * 50)
    print("MongoDB 副本集 Python 客户端测试")
    print("=" * 50)
    
    # 运行所有测试
    if test_replica_set_connection():
        # 测试写入和读取
        doc_id = test_write_to_primary()
        
        # 短暂等待，给复制一点时间
        if doc_id:
            print("\n等待1秒让数据复制...")
            time.sleep(1)
            test_read_from_secondary(doc_id)
        
        # 完整同步测试
        test_data_sync()
    
    print("\n" + "=" * 50)
    print("测试完成")
    print("=" * 50)

if __name__ == "__main__":
    main()