#!/usr/bin/env python3

import os
import shutil
from test_framework.test_framework import TestFramework
from utility.utils import wait_until

class SnapshotTask(TestFramework):
    def setup_params(self):
        self.num_nodes = 2

        # Enable random auto sync only
        for i in range(self.num_nodes):
            self.zgs_node_configs[i] = {
                "sync": {
                    "auto_sync_enabled": True,
                    "max_sequential_workers": 3,
                    "max_random_workers": 3,
                    "neighbors_only": True,
                }
            }

    def run_test(self):
        # Submit and upload files on node 0
        data_root_1 = self.__upload_file__(0, 256 * 1024)
        wait_until(lambda: self.nodes[1].zgs_get_file_info(data_root_1) is not None)
        wait_until(lambda: self.nodes[1].zgs_get_file_info(data_root_1)["finalized"])

        # Stop the last node to remove entire db folder and copy the flow_db (snapshot) from the first node
        self.stop_storage_node(1, clean=True)
        shutil.copytree(os.path.join(self.nodes[0].data_dir, 'db/flow_db'), os.path.join(self.nodes[1].data_dir, 'db/flow_db'))
    
        # Start the last node to verify snapshot sync
        self.start_storage_node(1)
        self.nodes[1].wait_for_rpc_connection()
        
        wait_until(lambda: self.nodes[1].zgs_get_file_info(data_root_1) is not None)
        wait_until(lambda: self.nodes[1].zgs_get_file_info(data_root_1)["finalized"])


if __name__ == "__main__":
    SnapshotTask().main()