from pyspark.sql.functions import col, get_json_object

class Key:	
 
    def __init__(self, value):
        self.value = value        
        self.value_rep = None
        self.values = [value]
        self.lvalues = []
        self.rvalues = []

    def right_keys(self):
        return dict.fromkeys(self.rvalues).keys()

    def left_keys(self):
        return dict.fromkeys(self.lvalues).keys()

    def chain_key(self, df_alias, key_name):
        return get_json_object(col(f"{df_alias}.payload"), f"$.{key_name}")
    
    def propagate_values(self, other):
        # na primeira vez, precisa preencher as chaves a esquerda da operação
        if len(self.lvalues) == 0:
            self.lvalues.append(self.value)

        # na primeira vez, precisa preencher as chaves a direita da operação
        if len(self.rvalues) == 0:
            self.rvalues.append(other.value)

        # propaga as chaves da esquerda
        self.lvalues = self.lvalues + other.lvalues
        # propaga as chaves da direita
        self.rvalues = self.rvalues + other.rvalues
        # propaga todos os valores de chaves
        self.values = self.values + other.values  

    def init_values(self, other):        
        if self.value_rep is None:
            self.value_rep = self.chain_key("block", self.value)        
        if other.value_rep is None:
            other.value_rep = self.chain_key("block_prev", other.value)

    def __repr__(self):
        return str(self.value_rep)

    def __str__(self):
        return str(self.value_rep)

    def __or__(self, other):                
        self.init_values(other)
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep | other.value_rep   
        # propaga as chaves
        self.propagate_values(other)       
        return self

    def __and__(self, other):                
        self.init_values(other)
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep & other.value_rep    
        # propaga as chaves
        self.propagate_values(other)       
        return self
    
    def __lt__(self, other):   
        self.init_values(other)     
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep < other.value_rep
        # propaga as chaves
        self.propagate_values(other)
        return self
    
    def __le__(self, other):   
        self.init_values(other)     
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep <= other.value_rep
        # propaga as chaves
        self.propagate_values(other)
        return self
    
    def __eq__(self, other):   
        self.init_values(other)     
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep == other.value_rep
        # propaga as chaves
        self.propagate_values(other)
        return self
    
    def __ne__(self, other):   
        self.init_values(other)     
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep != other.value_rep
        # propaga as chaves
        self.propagate_values(other)
        return self
    
    def __gt__(self, other):   
        self.init_values(other)     
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep > other.value_rep
        # propaga as chaves
        self.propagate_values(other)
        return self
    
    def __ge__(self, other):   
        self.init_values(other)     
        # popula a estrutura de chaves usando o col do spark
        self.value_rep = self.value_rep >= other.value_rep
        # propaga as chaves
        self.propagate_values(other)
        return self