from typing import Any, Dict

import pandas as pd

# from model_integration import TransformedOutput, execute_model_step

from RULE.model_integration import execute_model_step, TransformedOutput

def run_model_step(
    current_df: pd.DataFrame,
    operation: Dict[str, Any]
) -> dict:
    """
    Integrates the model step logic directly into the rule execution framework.
    Takes the current running DataFrame and the workflow step operation config,
    executes the remote model, and returns a dictionary compatible with 
    existing execution traces.
    """
    model_name = operation.get("model_name")
    if not model_name:
        raise ValueError("Model step requires 'model_name'")
        
    compliance_type = operation.get("compliance_type", "firco")
    user_id = operation.get("user_id", "system") 
    version = operation.get("version", "latest")
    number_of_reasonings = operation.get("number_of_reasonings", 1)
    op_name = operation.get("name", "Unnamed Model Step")
    
    transformed_output: TransformedOutput = execute_model_step(
        current_df=current_df,
        model_name=model_name,
        compliance_type=compliance_type,
        user_id=user_id,
        version=version,
        number_of_reasonings=number_of_reasonings
    )
    
    # Return enriched dataframe for the next step, plus trace info 
    return {
        "status": "success",
        "type": "model",
        "name": op_name,
        "prediction_id": getattr(transformed_output, "prediction_id", None),
        "input_rows": len(transformed_output.input_data),
        "output_rows": len(transformed_output.enriched_output),
        "enriched_df": transformed_output.enriched_df 
    }