import pandas as pd

from .engines import QSRuleBasedFallbackEngine
from .model import QSSelectiveCalibratedModel


def predict_with_cascade(
    df_raw: pd.DataFrame,
    ml_model: QSSelectiveCalibratedModel,
    fallback_engine: QSRuleBasedFallbackEngine,
) -> pd.DataFrame:
    """
    Executes the ultimate MLOps Cascade Architecture.

    Workflow:
    1. Run the highly-precise ML Model.
    2. Identify rejected items (routed to 'Others').
    3. Pass ONLY the rejected items to the deterministic Rule Engine.
    4. Merge the results, tag the provenance/source, and re-infer NRM hierarchies.

    Parameters
    ----------
    df_raw : pd.DataFrame
        The raw incoming BoQ data.
    ml_model : QSSelectiveCalibratedModel
        The loaded, hermetic ML production object.
    fallback_engine : QSRuleBasedFallbackEngine
        The initialized rule-based regex matcher.

    Returns
    -------
    pd.DataFrame
        The final DataFrame with completely resolved predictions and provenance tagging.
    """
    # -------------------------------------------------------------------------
    # 1. Primary ML Inference
    # -------------------------------------------------------------------------
    df_results = ml_model.predict_full(df_raw)

    # -------------------------------------------------------------------------
    # 2. Identify the ML Rejects
    # -------------------------------------------------------------------------
    others_mask = df_results["Predicted_Category"] == ml_model.others_label

    # Initialize the provenance column (Auditability requirement)
    df_results["Prediction_Source"] = "ML_Model"

    # If the ML model handled everything perfectly, return early
    if not others_mask.any():
        return df_results

    ml_model._logger.info(f"Fallback triggered for {others_mask.sum()} items.")

    # -------------------------------------------------------------------------
    # 3. Execute Fallback Rules on Rejects
    # -------------------------------------------------------------------------
    # The production object normalizes descriptions to the 'description' column
    rejected_descriptions = df_results.loc[others_mask, "description"]
    rule_predictions = fallback_engine.predict(rejected_descriptions)

    # -------------------------------------------------------------------------
    # 4. Merge Results & Update Provenance
    # -------------------------------------------------------------------------
    # Overwrite 'Others' with the Rule Engine's decisions
    df_results.loc[others_mask, "Predicted_Category"] = rule_predictions
    df_results.loc[others_mask, "Prediction_Source"] = "Rule_Engine"

    # Deterministic rules have 1.0 confidence. Failed items get 0.0 confidence.
    df_results.loc[
        others_mask & (rule_predictions != "Unclassified"), "Prediction_Confidence"
    ] = 1.0
    df_results.loc[
        others_mask & (rule_predictions == "Unclassified"), "Prediction_Confidence"
    ] = 0.0

    # -------------------------------------------------------------------------
    # 5. Re-infer NRM Hierarchy for Rescued Items
    # -------------------------------------------------------------------------
    # Find the rows that were 'Others' but successfully classified by rules
    rescued_mask = others_mask & (df_results["Predicted_Category"] != "Unclassified")

    if rescued_mask.any():
        # Map the new categories back to the NRM hierarchy using the ML model's dictionary
        rescued_hierarchies = df_results.loc[rescued_mask, "Predicted_Category"].apply(
            lambda cat: ml_model.infer_nrm_hierarchy(cat)
        )
        rescued_df = pd.DataFrame(
            list(rescued_hierarchies), index=rescued_hierarchies.index
        )

        # Dynamically overwrite the output columns (Handling both 'Predicted_ge' and 'ge' formats)
        for level in ["ge", "group_element", "e", "element"]:
            pred_col = f"Predicted_{level}"

            # If the model created a 'Predicted_' column, update that
            if pred_col in df_results.columns:
                df_results.loc[rescued_mask, pred_col] = rescued_df[level]
            # Otherwise, update the base column
            elif level in df_results.columns:
                df_results.loc[rescued_mask, level] = rescued_df[level]

    return df_results