import redis
import pymongo
import time
import json
import yaml

def sync_data(redis_client, mongo_db, redis_db, mongo_collection_name):
    mongo_collection = mongo_db[mongo_collection_name]
    all_keys = redis_client.keys('*')
    synced_count = 0
    
    for key in all_keys:
        try:
            poem_data = redis_client.get(key)
            
            if poem_data:
                poem_dict = json.loads(poem_data)
                
                # 检查MongoDB中是否已存在相同的文档
                existing_doc = mongo_collection.find_one({'_id': key.decode('utf-8')})
                
                if not existing_doc or existing_doc != poem_dict:
                    mongo_collection.update_one(
                        {'_id': key.decode('utf-8')},
                        {'$set': poem_dict},
                        upsert=True
                    )
                    print(f"Synced data with key: {key} from Redis DB {redis_db} to MongoDB collection '{mongo_collection_name}'")
                    synced_count += 1
        except json.JSONDecodeError:
            print(f"Error decoding JSON for key: {key} in Redis DB {redis_db}")
        except Exception as e:
            print(f"Error syncing data for key {key} in Redis DB {redis_db}: {str(e)}")
    
    return synced_count

def main(config):
    # Redis配置
    redis_host = config['redis']['host']
    redis_port = config['redis']['port']
    redis_password = config['redis']['password']

    # MongoDB配置
    mongo_uri = config['mongodb']['uri']
    mongo_db_name = config['mongodb']['db_name']

    # 连接到MongoDB
    mongo_client = pymongo.MongoClient(mongo_uri)
    mongo_db = mongo_client[mongo_db_name]

    # 固定的Redis数据库和MongoDB集合映射
    db_collection_map = {
        0: 'pose-result-db0',
        1: 'pose-result-db1',
        2: 'cpm-result-db2'
    }

    print("Selected databases and collections for syncing:")
    for db, collection in db_collection_map.items():
        print(f"  Redis DB {db} -> MongoDB collection '{collection}'")

    while True:
        print("Starting sync...")
        total_synced = 0
        for db, collection in db_collection_map.items():
            print(f"Syncing Redis DB {db} to MongoDB collection '{collection}'...")
            try:
                redis_client = redis.Redis(host=redis_host, port=redis_port, db=db, password=redis_password)
                synced_count = sync_data(redis_client, mongo_db, db, collection)
                total_synced += synced_count
            except redis.exceptions.AuthenticationError:
                print(f"Error: Authentication failed for Redis DB {db}. Skipping...")
            except redis.exceptions.ConnectionError:
                print(f"Error: Unable to connect to Redis DB {db}. Skipping...")
            except Exception as e:
                print(f"Error occurred while syncing Redis DB {db}: {str(e)}. Skipping...")
        
        if total_synced > 0:
            print(f"Sync completed. {total_synced} documents synced. Waiting for next update...")
        else:
            print("No new data to sync. Waiting for next update...")
        
        time.sleep(300)  # 等待5分钟后再次同步

if __name__ == "__main__":
    # 加载配置文件
    with open('worker_sys/function/config.yaml', 'r') as file:
        config = yaml.safe_load(file)

    # 运行主程序
    main(config)