from web3 import Web3
from src.utils import load_abi

class Scheduler():
    def __init__(self, w3: Web3, address: str, abi_file: str):
        self.w3 = w3
        abi = load_abi(abi_file)
        self.contract = self.w3.eth.contract(address=address, abi=abi)
        print(self.contract.functions)

    def signed_and_send_transaction(self, account, function):
        tx = function.build_transaction({
            'from': account.address,
            'nonce': self.w3.eth.get_transaction_count(account.address)
        })
        signed_tx = self.w3.eth.account.sign_transaction(tx, private_key=account.key)

        tx_hash = self.w3.eth.send_raw_transaction(signed_tx.rawTransaction)
        tx_receipt = self.w3.eth.wait_for_transaction_receipt(tx_hash)
        return tx_receipt

    def register_cluster(self, account, gpu_id: int, gpu_num: int):
        function = self.contract.functions.registerCluster(gpu_id, gpu_num)
        return self.signed_and_send_transaction(account, function)

    def register_task_with_conditions(self, account, data_image: str, train_image: str, gpu_id: int, cluster_size: int):
        function = self.contract.functions.registerTaskWithConditions(data_image, train_image, gpu_id, cluster_size)
        return self.signed_and_send_transaction(account, function)

    def get_clusters(self):
        return self.contract.functions.getClusters().call()

    def get_tasks(self):
        return self.contract.functions.getTasks().call()

    '''
    def get_start_run_event(self):
        event_filter = self.contract.events.StartRun.create_filter(fromBlock=0)
        events = event_filter.get_new_entries()
        print(events)

    def getRegisterClusterEvent(self):
        event_filter = self.contract.events.RegisterCluster.create_filter(fromBlock=0)
        events = event_filter.get_new_entries()
        print(events)
    '''

    def get_cluster_event(self, receipt):
        event = self.contract.events.RegisterCluster().process_receipt(receipt)
        return event