from typing import Any, Dict
import importlib.util
import numpy as np
import pandas as pd
import subprocess
import json
import os
import uuid
from app.main.engine.enum.StatusEnum import StatusEnum
from dpss import find_subset
from wrapt_timeout_decorator import timeout
from pyspark.sql.functions import sum, collect_list, round, when, col, lit, size, udf, array_except, array
from pyspark.sql.types import ArrayType, StringType, IntegerType, LongType, List
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 = 3
FACTOR = 1000


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 = False
        self.timeout_time = 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_time = 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_obs):
        response = {"status": StatusEnum.ERROR.name}
        try:
            @timeout(self.timeout_time)
            def __process(source_obj):
                try:
                    # 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, when, col, lit

                    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"])

                    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

                    # 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.limit(1)

                    # 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(self.pivot_params["columns-transaction"],
                                     self.ctp_params["columns-transaction"])]

                    total_merged = pivot_df2.join(ctp_df2, condition, 'left')
                    total_merged = total_merged.withColumn("DIFF",
                                                           when(col(self.ctp_params["columns-transaction"][0]).isNotNull(),
                                                                lit(0)).otherwise(lit(None)))
                    total_merged = total_merged.select(*pivot_df2.columns, "DIFF")

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

                    merged = merged.withColumn("DIFF", when(col("DIFF").isNull(),
                                                            total_merged[self.pivot_params["amount-column"]] - ctp_df2[
                                                                self.ctp_params["amount-column"]]).otherwise(col("DIFF")))

                    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)
                    columns = [col(column) for column in group_cols]
                    columns_amount = columns.copy()
                    columns_amount.append(col(amount_col))
                    custom = udf(custom_func_udf, ArrayType(IntegerType()))
                    resultado = df3.groupby(*columns).agg(
                        custom(collect_list(amount_col), collect_list(id_col), collect_list("DIFF"),collect_list(EXCLUDE_ROWS_FIELD),
                               lit(max_combinations)).alias("LISTA_DIFF"))

                    meged2 = resultado.join(merged_df, group_cols, 'left')
                    handle_array_udf = udf(handle_array, ArrayType(IntegerType()))

                    meged2 = meged2.withColumn("LISTA_DIFF", handle_array_udf("LISTA_DIFF"))

                    meged2 = meged2.filter((col("DIFF") == 0) | ((col("DIFF") != 0) & (size(col("LISTA_DIFF")) > 0)))
                    if self.exclude_pivot:
                        meged2 = meged2.withColumn("INTER_PIVOT_ID", array_except(meged2[self.pivot_params["id-column"]],
                                                                                  meged2["LISTA_DIFF"]))
                        meged2 = meged2.withColumnRenamed(self.ctp_params["id-column"], "INTER_CTP_ID")
                        if meged2.schema["INTER_CTP_ID"].dataType == LongType():
                            meged2 = meged2.withColumn("INTER_CTP_ID",
                                                       array(col("INTER_CTP_ID")).cast(ArrayType(LongType())))
                    else:
                        meged2 = meged2.withColumn("INTER_CTP_ID",
                                                   array_except(meged2[self.ctp_params["id-column"]], meged2["LISTA_DIFF"]))
                        meged2 = meged2.withColumnRenamed(self.pivot_params["id-column"], "INTER_PIVOT_ID")
                        if meged2.schema["INTER_PIVOT_ID"].dataType == LongType():
                            meged2 = meged2.withColumn("INTER_PIVOT_ID",
                                                       array(col("INTER_PIVOT_ID")).cast(ArrayType(LongType())))
                    meged2 = meged2.toPandas()
                    return meged2

                except Exception as e:
                    self.timeout = True
                    self.app.logger.error(f"Error de Timeout. Error: {e}")
                    raise TimeoutError("Tiempo de ejecución superado.")

            response["status"] = StatusEnum.OK.name
            response["result"] = __process(source_obs)
        except TimeoutError as e:
            response["status"] = StatusEnum.TIMEOUT.name
            response["message"] = e
        except Exception as e:
            response["status"] = StatusEnum.ERROR.name
            response["message"] = e
        finally:
            self.output = response

    def response(self):
        return self.output

    def timeout_response(self):
        return self.timeout


    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 handle_array(x):
    if isinstance(x, List):
        return x
    else:
        return []


def custom_func(group, amount_field, id_field, max_combinations):

    diff = int(group["DIFF"].values[0]*(10**ROUND_DECIMAL))

    if pd.isna(diff) or diff == 0:
        return None
    group = group[group[EXCLUDE_ROWS_FIELD] == 'S']
    group[amount_field] = group[amount_field].astype(float)
    group = group.reset_index(drop=True)
    values = group[amount_field].values
    values *= (10**ROUND_DECIMAL)
    values = values.astype(np.int64)
    ids = group[id_field].values

    result = find_subset(values, diff, max_combinations)
    if len(result) <1:
        return None
    result = result[0]
    indices = []
    for idx, val in zip(ids, values):
        if val in result:
            indices.append(idx)
            result.remove(val)
    return indices



# def custom_func_udf(amount_values, id_values, diffs, max_combinations):
#     diff = diffs[0]
#     if pd.isna(diff) or diff == 0:
#         return None
#     diff = int(diff * FACTOR)
#     amount_values = [int(value * FACTOR) for value in amount_values]
#     result = find_subset(amount_values, diff, max_combinations)
#     if len(result) <1:
#         return None
#     result = result[0]
#     indices = []
#     for idx, val in zip(id_values, amount_values) :
#         if val in result:
#             indices.append(idx)
#             result.remove(val)
#     return indices

def custom_func_udf(amount_values, id_values, diffs, excludes, max_combinations):
    diff = diffs[0]
    if pd.isna(diff) or diff == 0:
        return None
    diff = int(diff * FACTOR)
    amount_values = [int(value * FACTOR) for value, exclude in zip(amount_values, excludes) if exclude=="S"]
    dir_name = str(uuid.uuid4())
    prefix = "/tmp/" + dir_name + "_"

    tmp_file_arr1, tmp_file_arr2 = "values.txt", "target.txt"
    full_path_arr1, full_path_arr2 = prefix + tmp_file_arr1, prefix + tmp_file_arr2
    with open(full_path_arr1, 'w') as archivo:
        archivo.writelines([f'{entero}\n' for entero in amount_values])
    with open(full_path_arr2, 'w') as archivo:
        archivo.write(str(diff))
    executable_path = '/home/evillarroel/Descargas/Telegram Desktop/subset_sum_linux'
    indices = []
    for comb in range(1, max_combinations+1):
        argumentos = [full_path_arr1, full_path_arr2, str(comb), '1', '1', 'false', 'false']
        result = subprocess.run([executable_path] + argumentos, check=True, capture_output=True, text=True)
        result = str(result)
        if "keys:[" in result:
            match = result[result.index("keys:[") + 5:result.index("keys remainder") - 20]
            match = match.replace("targets:", "").replace("+", ",")
            match = match.split("==")[0].replace(" ", "")
            match = json.loads(match)
            for idx, val in zip(id_values, amount_values):
                if val in match:
                    indices.append(idx)
                    match.remove(val)
            break
    os.remove(full_path_arr1), os.remove(full_path_arr2)
    return indices