from pyspark.sql.types import StructType, StringType, IntegerType, DateType, StructField, TimestampType
from pyspark.sql.functions import col, lit
from dnalib.utils import TableUtils, Utils
from dnalib.log import log
from dnalib.tables import DynamicTable
from delta.tables import DeltaTable

class BlockChain(DynamicTable):

    fields_comment = {
        "block_name":  {'comment':"Nome do bloco, utilizado para otimizar as conexões do blockchain."},
        "block_level":  {'comment':"Nível do bloco, usado para identificar a sua profundidade na blockchain."},
        "block_business_keys":  {'comment':"Campo com as chaves de negócio concatenadas."},
        "id_block_gen":  {'comment':"Identificador gerado no primeiro bloco e propagado por toda a cadeia."},
        "id_block":  {'comment':"Identificador único do bloco, gerado a partir de um hash."},
        "id_block_prev":  {'comment':"Identificador id_block do bloco anterior."},
        "is_gen":  {'comment':"Flag que indica se o bloco é gerador (1 se sim e 0 se não)."},
        "block_watermark":  {'comment':"Timestamp da inserção do bloco na blockchain."},
        "payload":  {'comment':"Metadados em formato JSON, são as colunas do dataframe que representa o bloco."},
        "last_block":  {'comment':"Flag que indica que o bloco é o último na cadeia."},
    }

    schema = StructType([
        StructField("block_name", StringType(), True, fields_comment["block_name"]),
        StructField("block_level", IntegerType(), True, fields_comment["block_level"]),
        StructField("block_business_keys", StringType(), True, fields_comment["block_business_keys"]),
        StructField("id_block_gen", StringType(), True, fields_comment["id_block_gen"]),
        StructField("id_block", StringType(), True, fields_comment["id_block"]),
        StructField("id_block_prev", StringType(), True, fields_comment["id_block_prev"]),
        StructField("is_gen", IntegerType(), True, fields_comment["is_gen"]),
        StructField("block_watermark", TimestampType(), True, fields_comment["block_watermark"]),
        StructField("payload", StringType(), True, fields_comment["payload"]),
        StructField("last_block", IntegerType(), True, fields_comment["last_block"]),  
    ])
    
    def __init__(self, table_name, layer="gold"):    
        super().__init__(layer, table_name, json_fields=["payload"])                       
        if not TableUtils.table_exists(self.layer, self.table_name):
            self.df_block_chain = (
                Utils.spark_instance().catalog.createTable(f"{self.layer}.{table_name}", 
                                        path=Utils.lakehouse_path(self.layer, table_name),
                                        schema=BlockChain.schema, 
                                        source='delta'))
        else:
            self.df_block_chain = Utils.spark_instance().table(f"""{self.layer}.{self.table_name}""") 
        # store the current computation of blockchain
        self.df_block_chain_current = Utils.spark_instance().createDataFrame([], self.schema)        

    def add_block(self, block, prev_block_name=None, prev_block_join_rules=None):         
        # carrega o bloco
        block.load()

        # usa a blockchain como o "bloco anterior"
        df_block_chain = self.df_block_chain.alias("block_prev")        

        #
        #   É preciso determinar primeiro quais transações são novas
        #   e isso pode ser feito a partir do id block (hash)
        #   toda transação nova não pode estar na blockchain
        #   
        #   Dado que A é o bloco novo e B a blockchain, 
        #   então A \ B são as novas transações (garantia de unicidade de registro na blockchain).
        #
        df_block = (
            block.df_block.join(
                df_block_chain, 
                on=[col("block.id_block") == col("block_prev.id_block")], 
                how="leftanti"
        )).alias("block")        

        #
        # Estrutura dos blocos geradores
        #
        gen_block_fields = (
            col("block.block_name"),
            col("block.block_level"),
            col("block.block_business_keys"),
            col("block.id_block").alias("id_block_gen"),
            col("block.id_block"),
            lit(None).alias("id_block_prev"),
            lit(1).alias("is_gen"),            
            col("block.block_watermark"),
            col("block.payload"),
            lit(1).alias("block.last_block")
        )

        # blocos que iniciam a cadeia não tem estrutura de prev (são raízes do grafo) 
        if prev_block_join_rules is not None:     
            
            # otimizando a consulta filtrando pelo tipo do block anterior (usando o nome), já que todo bloco deve ter apenas um parent
            if prev_block_name is not None:                
                df_block_chain = df_block_chain.filter(col("block_name") == lit(prev_block_name))
                
            #
            # Primeira operação é verificar se existe conexão com bloco anterior (usando left anti join)
            # nesse caso os blocos são geradores (ou seja não possuem conexão anterior na cadeia).
            #
            df_next_block = (
                df_block.join(
                    df_block_chain, 
                    on=prev_block_join_rules.value_rep, # chaves de negócio
                    how="leftanti"
            ))                        

            # o bloco é inserido na blockchain       
            df_block_union = df_next_block.select(*gen_block_fields)     
            self.df_block_chain_current = self.df_block_chain_current.union(df_block_union)

            # agora é verificado os blocks que tem conexão na cadeia (via inner join)
            df_next_block = (
                df_block.join(
                    df_block_chain, 
                    on=prev_block_join_rules.value_rep, # chaves de negócio
                    how="inner"
            ))      

            #
            # As colunas do blockchain são agora propagadas para o próximo bloco:
            #   o id_block_gen vem do bloco anterior
            #   o id_block_prev é o id do bloco anterior 
            #   como conectou, esse bloco não é gerador
            #
            propagate_block_fields = (
                col("block.block_name"),
                (col("block_prev.block_level") + lit(1)).alias("block_level"),
                col("block.block_business_keys"),
                col("block_prev.id_block_gen"),  
                col("block.id_block"),
                col("block_prev.id_block").alias("id_block_prev"),    
                lit(0).alias("is_gen"),
                col("block.block_watermark"),
                col("block.payload"),
                lit(1).alias("block.last_block")
            )
            
            # o bloco é inserido na blockchain                         
            df_block_union = df_next_block.select(*propagate_block_fields)
            self.df_block_chain_current = self.df_block_chain_current.union(df_block_union)                                  
        else:
          # quando o bloco é a "raíz" ele não possui bloco prev, portanto ele é inserido direto na blockchain
          df_block_union = df_block.select(*gen_block_fields)          
          self.df_block_chain_current = self.df_block_chain_current.union(df_block_union)           
        
        # logando o número de novos blocos
        total_new_blocks = self.df_block_chain_current.count()        

        # atualizando a blockchain
        if total_new_blocks > 0:
            log(__name__).info(f"Total of {total_new_blocks} blocks of type {block.block_name} will be added to blockchain.")
            self.update()
        else:
            log(__name__).info(f"No new blocks of type {block.block_name} to be added to blockchain.")

        # cleaning the current blockchain to allow new computations
        self.df_block_chain_current = Utils.spark_instance().createDataFrame([], self.schema)        
        return self

    def add_last_block(self):              
        # para todo bloco A na cadeia, se existe B tal que B \neq A e id_block_prev(B) = id_block(A), então last_block(A) = 0
        log(__name__).info(f"Updating last_block column to previously blockchain blocks.")
        df_non_last_blocks = self.df_block_chain_current.select(col("id_block_prev").alias("_id_block_prev")).filter("_id_block_prev is not null").distinct()

        # como a computação é incremental, então existem blocos onde o last_block(A) = 0, porém last_block_{-1}(A) = 1 (ou seja, estão salvos na tabela delta como 1)
        df_prev_last_blocks = df_non_last_blocks.alias("block").join(
            self.df_block_chain_current.alias("block_prev"),
            on = [col("block._id_block_prev") == col("block_prev.id_block")],
            how = "leftanti",
        )

        # atualiza a tabela delta
        df_source_block_chain = DeltaTable.forName(Utils.spark_instance(), f"{self.layer}.{self.table_name}")
        (df_source_block_chain.alias("block_prev")
            .merge(
                df_prev_last_blocks.alias("block"),
                "block_prev.id_block = block._id_block_prev"
            )
            .whenMatchedUpdate(set={"last_block":lit(0)})
            .execute())

    # atualiza a tabela dos blocos
    def update(self):                        
        self.add_last_block()                        
        # as novas transações são inseridas na tabela do blockchain
        log(__name__).info(f"Appeding new blocks to blockchain table {self.layer}.{self.table_name}.")
        self.df_block_chain_current.write.mode("append").saveAsTable(f"{self.layer}.{self.table_name}")

    # métodos para consultar os blocos de maneira friendly
    """def find_blocks_by_single_field(self, **kwargs):           
        return self.find_all_by_single_parameter(layer=self.layer, 
                                                table_name=self.table_name,     
                                                **kwargs)

    def find_blocks_by_single_payload_parameter(self, **kwargs):   
        json_value_key, json_value = kwargs.popitem()
        return self.find_all_by_json_parameter(layer=self.layer, 
                                               table_name=self.table_name, 
                                               json_field="payload", 
                                               json_value_key=json_value_key, 
                                               json_value=json_value)

    def find_chains_by_single_payload_parameter(self, **kwargs):
        json_value_key, json_value = kwargs.popitem()
        return self.find_all_by_json_parameter_in_field(layer=self.layer, 
                                                        table_name=self.table_name, 
                                                        json_field="payload", 
                                                        in_field="id_block_gen",
                                                        json_value_key=json_value_key,
                                                        json_value=json_value)"""
        