import json
from typing import Dict
from pyspark.sql import SparkSession, DataFrame
from pyspark.sql import functions as F
from .config import dq_config
from .validate_checks import validate_checks


def run_dq_pipeline(
    spark: SparkSession,
    df_dict: Dict[str, DataFrame],
    config_path: str,
    checks: Dict[str, callable] = None,
    max_json_records: int = 100
) -> Dict[str, Dict[str, DataFrame]]:
    """
    Highly optimized Data Quality Pipeline

    - Computes summary on raw_df for each check in a single pass
    - Produces final clean_df
    - Produces error_df
    - JSON includes nulls and limited to max_json_records
    """

    dq_config_df, return_mode = dq_config(spark, config_path)
    validate_checks(dq_config_df, checks)

    result_dict = []
    summary_rows = []

    for table_name, df in df_dict.items():
        df = df.persist()
        try:
            raw_df = df
            total_rows = raw_df.count()
            clean_df = raw_df

            # Prepare a column to track failed checks per row
            failed_checks_col = F.array()  # start empty array
            row_invalid_flags = []

            # Wrap checks
            def _wrapped(func):
                def call(df_, row_):
                    params = json.loads(row_["params"]) if row_["params"] else {}
                    return func(df_, row_, **params)
                return call
            table_check_map = {name: _wrapped(func) for name, func in checks.items()}

            # Run all checks on raw_df in one pass
            for row in dq_config_df.filter(F.col("table") == table_name).collect():
                column = row["column"]
                check_name = row["check"]

                if check_name not in table_check_map:
                    raise ValueError(f"Check {check_name} not found in provided checks")

                # Run check on raw_df
                valid_df, invalid_df, _ = table_check_map[check_name](raw_df, row)

                # Tag invalid rows with this check
                if invalid_df.limit(1).count() > 0:
                    invalid_df_tagged = invalid_df.withColumn("__failed_check", F.lit(check_name))
                    row_invalid_flags.append(invalid_df_tagged)

                # Update clean_df sequentially
                clean_df = valid_df

            # Build error_df once by unioning all tagged invalid dfs
            if row_invalid_flags:
                error_df = row_invalid_flags[0]
                for df_ in row_invalid_flags[1:]:
                    error_df = error_df.unionByName(df_, allowMissingColumns=True)
            else:
                error_df = spark.createDataFrame([], df.schema)

            # Build summary
            for row in dq_config_df.filter(F.col("table") == table_name).collect():
                column = row["column"]
                check_name = row["check"]

                # Use filtered error_df for this check
                check_invalid_df = error_df.filter(F.col("__failed_check") == check_name)
                failed_count = check_invalid_df.count()
                failed_percentage = round((failed_count / total_rows) * 100, 2) if total_rows > 0 else 0.0

                failed_records_json = (
                    check_invalid_df.limit(max_json_records)
                    .select(F.to_json(F.struct(*check_invalid_df.columns), options={"ignoreNullFields": "false"}).alias("json"))
                    .agg(F.collect_list("json").alias("failed_records_json"))
                    .collect()[0]["failed_records_json"]
                ) or []

                summary_rows.append(
                    (
                        table_name,
                        column,
                        check_name,
                        total_rows,
                        failed_count,
                        failed_percentage,
                        json.dumps(failed_records_json)
                    )
                )

            # Build return dictionary
            table_result = {}
            if "summary+clean" in return_mode or "all" in return_mode:
                table_result["clean"] = clean_df
            if "summary+error" in return_mode or "all" in return_mode:
                table_result["error"] = error_df.drop("__failed_check")

            if table_result:
                result_dict.append((table_name.lower(), table_result))

        finally:
            df.unpersist()

    # Build summary dataframe
    summary_schema = """
        table_name STRING, 
        column_name STRING, 
        check_type STRING, 
        total_rows LONG, 
        failed_count LONG, 
        failed_percentage DOUBLE, 
        failed_records_json STRING
    """
    summary_df = spark.createDataFrame(summary_rows, schema=summary_schema)

    # Assemble final result dict
    final_result = {name: tbl for name, tbl in result_dict}
    final_result['summary'] = summary_df
    return final_result
