from typing import Any, Dict, List
import importlib.util
from itertools import combinations
import time
import numpy as np
import pandas as pd
import dask.dataframe as dd
from concurrent.futures import ThreadPoolExecutor
import concurrent.futures
# import multiprocessing as mp
from functools import partial
# import asyncio

from app.main.engine.action.ActionInterface import ActionInterface

# RELACION DE IDENTIFICADOR DE ACCION Y NOMBRE DE CLASE
relation_classname_identifier = {
    "match-and-exclude-records-actions": "MatchAndExcludeRecordsAction"
}

# CONFIGURACION DE SESION DE SPARK
MASTER = "local[*]"
DRIVER_MEMORY = "8g"
EXECUTOR_MEMORY = "8g"
MYSQL_JAR_PATH = "jars/mysql-connector-java-8.0.30.jar"

# EXCLUDE VALIDATION FIELD
EXCLUDE_ROWS_FIELD = "EXCLUDE_VALID"
# REDONDEO DE DECIMALES
ROUND_DECIMAL = 2


class MatchAndExcludeRecordsAction(ActionInterface):

    library_required = "pyspark"
    version_required = "3.4.0"

    def __init__(self, app) -> None:
        super().__init__(app)
        self.max_combinations = None
        self.timeout = None
        self.exclude_pivot = None
        self.pivot_params = None
        self.ctp_params = None
        self.output = None
        self.config_params = ["max-records-per-combinations", "max-timeout-per-combinations", "exclude-entity-pivot"]

    def parser(self, descriptor: Dict[str, Any]):
        # Validar si pyspark y su versión está instalada
        pyspark_lib = importlib.util.find_spec(self.library_required)
        if pyspark_lib is None:
            raise ImportError(f"Librería requerida *{self.library_required}* no instalada")
        import pyspark
        version = pyspark.__version__
        if version != self.version_required:
            raise ImportError(f"Versión requerida no instalada. Requerida: {self.version_required}. Instalada: {version}")

        # Validación de parámetros de entrada
        entity_config_params = ["tablename", "id-column", "amount-column", "columns-group", "columns-transaction"]

        if "config-params" not in descriptor.keys():
            raise ReferenceError(f"Parámetro *params-input* no definido en descriptor")

        configs = descriptor["config-params"]
        for param in self.config_params:
            if param not in configs.keys():
                raise ReferenceError(f"Parámetro *{param}* no definido en configuracion")

        if "params-input" not in descriptor.keys():
            raise ReferenceError(f"Parámetro *params-input* no definido en descriptor")

        params_input = descriptor["params-input"]
        if "pivot-config" not in params_input.keys() or "counterpart-config" not in params_input.keys():
            raise ReferenceError(f"Pivot o Contraparte no definidos en descriptor")

        pivot_params, ctp_params = params_input["pivot-config"], params_input["counterpart-config"]
        for param in entity_config_params:
            if param not in pivot_params.keys() or param not in ctp_params.keys():
                raise ReferenceError(f"Parámetro *{param}* no encontrado en pivot o contraparte")

        self.max_combinations = configs["max-records-per-combinations"]
        self.timeout = configs["max-timeout-per-combinations"]
        self.exclude_pivot = configs["exclude-entity-pivot"]
        self.pivot_params = pivot_params
        self.ctp_params = ctp_params

    def process(self, source_obj):

        # Inicializar la sesion de Spark
        session = self.createSession()

        # Traer la data desde BD tanto pivot como contraparte
        pivot_table, ctp_table = self.pivot_params["tablename"], self.ctp_params["tablename"]
        jdbc_conn = source_obj.create_spark_connection()
        jdbc_url = jdbc_conn["url"]
        jdbc_properties = jdbc_conn["properties"]
        pivot_df = session.read.jdbc(url=jdbc_url, table=pivot_table, properties=jdbc_properties)
        ctp_df = session.read.jdbc(url=jdbc_url, table=ctp_table, properties=jdbc_properties)

        # Agregar un prefijo a cada columna, tanto del pivot como contraparte. Actualizar campos del input
        # pivot: 'PIVOT_', contraparte: 'COUNTERPART_'
        for column in pivot_df.columns:
            if column == EXCLUDE_ROWS_FIELD:
                continue
            pivot_df = pivot_df.withColumnRenamed(column, "PIVOT_"+column)

        for column in ctp_df.columns:
            if column == EXCLUDE_ROWS_FIELD:
                continue
            ctp_df = ctp_df.withColumnRenamed(column, "COUNTERPART_"+column)

        for key_p, key_c in zip(self.pivot_params.keys(), self.ctp_params.keys()):
            if isinstance(self.pivot_params[key_p], str):
                self.pivot_params[key_p] = "PIVOT_"+self.pivot_params[key_p]
                self.ctp_params[key_c] = "COUNTERPART_"+self.ctp_params[key_c]
            else:
                self.pivot_params[key_p] = ["PIVOT_"+column for column in self.pivot_params[key_p]]
                self.ctp_params[key_c] = ["COUNTERPART_" + column for column in self.ctp_params[key_c]]

        from pyspark.sql.functions import sum, collect_list, round

        pivot_cols = self.pivot_params["columns-transaction"].copy()
        if self.pivot_params["amount-column"] in pivot_cols:
            pivot_cols.remove(self.pivot_params["amount-column"])

        ctp_cols = self.ctp_params["columns-transaction"].copy()
        if self.ctp_params["amount-column"] in ctp_cols:
            ctp_cols.remove(self.ctp_params["amount-column"])

        timeout = self.timeout
        max_combinations = self.max_combinations

        # Ejecutamos lógica de excluir registros
        if len(self.pivot_params["columns-group"]) == 0 and len(self.ctp_params["columns-group"]) == 0:
            raise RuntimeError(f"Debe haber al menos pivot o contraparte agrupado")

        # Caso: 1 - Muchos
        elif len(self.pivot_params["columns-group"]) == 0 and len(self.ctp_params["columns-group"]) > 0:
            ctp_df2 = ctp_df.groupby(self.ctp_params["columns-group"]). \
                agg(round(sum(self.ctp_params["amount-column"]), ROUND_DECIMAL).alias(self.ctp_params["amount-column"]),
                    collect_list(self.ctp_params["id-column"]).alias(self.ctp_params["id-column"]))

            pivot_df2 = pivot_df.dropDuplicates(subset=pivot_cols)

        # Caso: Muchos - 1
        elif len(self.pivot_params["columns-group"]) > 0 and len(self.ctp_params["columns-group"]) == 0:
            pivot_df2 = pivot_df.groupby(self.pivot_params["columns-group"]).\
                agg(round(sum(self.pivot_params["amount-column"]), ROUND_DECIMAL).alias(self.pivot_params["amount-column"]),
                    collect_list(self.pivot_params["id-column"]).alias(self.pivot_params["id-column"]))

            ctp_df2 = ctp_df.dropDuplicates(subset=ctp_cols)

        # Caso: Muchos - Muchos
        elif len(self.pivot_params["columns-group"]) > 0 and len(self.ctp_params["columns-group"]) > 0:
            pivot_df2 = pivot_df.groupby(self.pivot_params["columns-group"]). \
                agg(round(sum(self.pivot_params["amount-column"]), ROUND_DECIMAL).alias(self.pivot_params["amount-column"]),
                    collect_list(self.pivot_params["id-column"]).alias(self.pivot_params["id-column"]))
            ctp_df2 = ctp_df.groupby(self.ctp_params["columns-group"]). \
                agg(round(sum(self.ctp_params["amount-column"]), ROUND_DECIMAL).alias(self.ctp_params["amount-column"]),
                    collect_list(self.ctp_params["id-column"]).alias(self.ctp_params["id-column"]))

        condition = [pivot_df2[col1] == ctp_df2[col2] for col1, col2 in zip(pivot_cols, ctp_cols)]
        merged = pivot_df2.join(ctp_df2, condition)

        merged = merged.withColumn("DIFF", pivot_df2[self.pivot_params["amount-column"]] - ctp_df2[self.ctp_params["amount-column"]])

        merged_df = merged.withColumn("DIFF", round(merged["DIFF"], ROUND_DECIMAL))
        if self.exclude_pivot:
            df = pivot_df
            group_cols = self.pivot_params["columns-group"]
            amount_col = self.pivot_params["amount-column"]
            id_col = self.pivot_params["id-column"]
        else:
            df = ctp_df
            group_cols = self.ctp_params["columns-group"]
            amount_col = self.ctp_params["amount-column"]
            id_col = self.ctp_params["id-column"]

        total_tmp_cols = group_cols + ["DIFF"]
        df3 = df.join(merged_df.select(*total_tmp_cols), group_cols)

        df3 = df3.toPandas()
        total_cols = group_cols + [amount_col, id_col, EXCLUDE_ROWS_FIELD,  "DIFF"]
        num_chunks = 4

        resultado = df3[total_cols].groupby(group_cols).apply(lambda x: custom_func(x, amount_col, id_col, timeout, max_combinations))
        resultado = resultado.reset_index()
        if len(resultado.columns) == 1:
            resultado = pd.DataFrame([], columns=group_cols + ["LISTA_DIFF"])
        else:
            resultado.columns = group_cols + ["LISTA_DIFF"]
        meged2 = resultado.merge(merged_df.toPandas(), 'left', group_cols)

        meged2["LISTA_DIFF"] = meged2["LISTA_DIFF"].apply(self.handle_array)
        meged2 = meged2[(meged2['DIFF'] == 0) | ((meged2['DIFF'] != 0) & (meged2['LISTA_DIFF'].apply(len) > 0))]
        if meged2.empty:
            pass
        elif self.exclude_pivot:
            meged2['INTER_PIVOT_ID'] = meged2.apply(lambda row: self.array_except(row[self.pivot_params["id-column"]], row['LISTA_DIFF']), axis=1)
            meged2 = meged2.rename(columns={self.ctp_params["id-column"]: "INTER_CTP_ID"})
            if meged2['INTER_CTP_ID'].dtype == 'int64':
                merged_df['INTER_CTP_ID'] = merged_df['INTER_CTP_ID'].apply(lambda x: [x]).astype('object')
        else:
            meged2['INTER_CTP_ID'] = meged2.apply(lambda row: self.array_except(row[self.ctp_params["id-column"]], row['LISTA_DIFF']), axis=1)
            meged2 = meged2.rename(columns={self.pivot_params["id-column"]: "INTER_PIVOT_ID"})
            if meged2['INTER_PIVOT_ID'].dtype == 'int64':
                merged_df['INTER_PIVOT_ID'] = merged_df['INTER_PIVOT_ID'].apply(lambda x: [x]).astype('object')

        self.output = meged2

    def response(self):
        return self.output

    def handle_array(self, x):
        if isinstance(x, List):
            return x
        else:
            return []

    def array_except(self, arr1, arr2):
        if arr2 is None:
            return arr1
        else:
            return [item for item in arr1 if item not in arr2]

    def createSession(self, name: str = "app_engine_spark"):
        try:
            from pyspark.sql import SparkSession
            session = SparkSession.builder.master(MASTER) \
                .appName(name) \
                .config("spark.jars", MYSQL_JAR_PATH) \
                .config("spark.executor.extraClassPath", MYSQL_JAR_PATH) \
                .config("spark.driver.extraClassPath", MYSQL_JAR_PATH) \
                .config("spark.driver.memory", DRIVER_MEMORY) \
                .config("spark.executor.memory", EXECUTOR_MEMORY) \
                .getOrCreate()
            self.app.logger.info(f"Sesión creada exitosamente")
            return session
        except Exception as e:
            raise Exception(f"Error creando sesion Spark. {e}")


def custom_func(group, amount_field, id_field, timeout, max_combinations):
    diff = group["DIFF"].values[0]
    if diff == 0:
        return None
    group[amount_field] = group[amount_field].astype(float)
    group = group.reset_index(drop=True)
    values = group[amount_field].values
    indexs = group.index.values
    tam = len(values)
    rang = range(1, tam + 1) if tam <= max_combinations else range(1, max_combinations + 1)

    final = None
    stop_event = False

    def buscar_combinacion(i):
        nonlocal final, stop_event

        if not stop_event:  # Solo continuar si no se ha encontrado una combinación aún
            for comb in combinations(indexs, i):
                if stop_event:  # Verificar si se ha establecido la señal de detención
                    break
                if np.isclose(np.sum(values[list(comb)]), diff):
                    # Validamos si la combinación son de registros validados 'EXCLUDE_VALID'
                    result = group[group.index.isin(comb)]
                    if (result[EXCLUDE_ROWS_FIELD] == 'S').all():
                        final = result[id_field].tolist()
                        stop_event = True  # Establecer la señal de detención
                        break

        return None

    start_time = time.time()

    with ThreadPoolExecutor() as executor:
        futures = [executor.submit(buscar_combinacion, i) for i in rang]

        for future in futures:
            try:
                future.result()
            except TimeoutError:
                stop_event = True  # Establecer la señal de detención si se alcanza el tiempo límite
                break

            if stop_event or final is not None:
                break

    if time.time() - start_time >= timeout:
        return None  # Devolver None si se alcanzó el tiempo límite
    return final












