import hashlib
import time
import json


class Block:
    def __init__(self, previous_hash, trigger_set_huggingface_hash, trigger_set_client_hash, encrypted_watermarked_model_hash, counter,
                 timestamp=None):
        self.timestamp = timestamp if timestamp else time.time()
        self.previous_hash = previous_hash
        self.counter = counter
        self.trigger_set_huggingface = trigger_set_huggingface_hash
        self.trigger_set_client = trigger_set_client_hash
        self.encrypted_watermarked_model = encrypted_watermarked_model_hash
        self.hash = self.calculate_hash()

    def calculate_hash(self):
        hash_string = (
                f"{self.timestamp:.6f}" +
                str(self.previous_hash) +
                str(self.counter) +
                str(self.trigger_set_huggingface) +
                str(self.trigger_set_client) +
                str(self.encrypted_watermarked_model)
        )
        return hashlib.sha256(hash_string.encode()).hexdigest()

    @staticmethod
    def hash_data(data):
        return hashlib.sha256(str(data).encode()).hexdigest()

    def to_dict(self):
        return {
            "timestamp": self.timestamp,
            "previous_hash": self.previous_hash,
            "counter": self.counter,
            "trigger_set_huggingface": self.trigger_set_huggingface,
            "trigger_set_client": self.trigger_set_client,
            "encrypted_watermarked_model": self.encrypted_watermarked_model,
            "hash": self.hash
        }


class Blockchain:
    def __init__(self):
        self.chain = {}
        self.add_block("Genesis HuggingFace", "Genesis Client", "Genesis Model")

    def add_block(self, trigger_set_huggingface, trigger_set_client, encrypted_watermarked_model):
        counter = len(self.chain)
        previous_hash = self.chain[counter - 1].hash if counter > 0 else "0"
        new_block = Block(previous_hash, trigger_set_huggingface, trigger_set_client, encrypted_watermarked_model,
                          counter)
        self.chain[counter] = new_block
        return new_block

    def is_chain_valid(self):
        for i in range(1, len(self.chain)):
            current_block = self.chain[i]
            previous_block = self.chain[i - 1]

            if current_block.hash != current_block.calculate_hash():
                print(f"Invalid hash for block {i}")
                return False

            if current_block.previous_hash != previous_block.hash:
                print(f"Invalid previous hash for block {i}")
                return False

        return True

    def to_dict(self):
        return {str(counter): block.to_dict() for counter, block in self.chain.items()}

    def save_to_file(self, filename):
        with open(filename, 'w') as file:
            json.dump(self.to_dict(), file, indent=4)
        print(f"Blockchain saved to {filename}")

    @classmethod
    def load_from_file(cls, filename):
        with open(filename, 'r') as file:
            data = json.load(file)

        blockchain = cls()
        blockchain.chain.clear()  # Clear the genesis block
        for counter, block_data in data.items():
            block = Block(
                block_data["previous_hash"],
                block_data["trigger_set_huggingface"],
                block_data["trigger_set_client"],
                block_data["encrypted_watermarked_model"],
                int(counter),
                block_data["timestamp"]
            )
            blockchain.chain[int(counter)] = block

        print(f"Blockchain loaded from {filename}")
        return blockchain, data


def print_blockchain_details(blockchain):
    for counter, block in blockchain.chain.items():
        print(f"Block {counter}:")
        print(f"  Timestamp: {block.timestamp:.6f}")
        print(f"  Previous Hash: {block.previous_hash}")
        print(f"  Hash: {block.hash}")
        print(f"  Calculated Hash: {block.calculate_hash()}")
        print()


# # Exemple d'utilisation
# blockchain = Blockchain()
#
# # Ajouter quelques blocs
# blockchain.add_block("HF Trigger Set 1", "Client Trigger Set 1", "Encrypted Model 1")
# blockchain.add_block("HF Trigger Set 2", "Client Trigger Set 2", "Encrypted Model 2")
# blockchain.add_block("HF Trigger Set 3", "Client Trigger Set 3", "Encrypted Model 3")
#
# print("Original Blockchain:")
# print_blockchain_details(blockchain)
#
# # Sauvegarder la blockchain dans un fichier JSON
# blockchain.save_to_file("blockchain.json")
#
# # Charger la blockchain depuis le fichier JSON
# loaded_blockchain, _ = Blockchain.load_from_file("blockchain.json")
#
# print("\nLoaded Blockchain:")
# print_blockchain_details(loaded_blockchain)
#
# # Vérifier que la blockchain chargée est valide
# print(f"La blockchain chargée est valide : {loaded_blockchain.is_chain_valid()}")
#
# # Ajouter un nouveau bloc à la blockchain chargée
# loaded_blockchain.add_block("HF Trigger Set 4", "Client Trigger Set 4", "Encrypted Model 4")
#
# # Sauvegarder la blockchain mise à jour
# loaded_blockchain.save_to_file("blockchain.json")