from typing import Any, Dict

import pandas as pd

from .schema import TransformedOutput


def parse_model_response(
    original_df: pd.DataFrame,
    api_response: Dict[str, Any]
) -> TransformedOutput:
    """
    Parses the completed API response and aligns it strictly 
    with the original input data.
    """
    
    # 1. Extract predictions block from the API response
    api_preds = []
    
    # The results are typically nested under "data"
    data_block = api_response.get("data", api_response)
    
    if "predictions" in data_block:
        outer_preds = data_block["predictions"]
        if isinstance(outer_preds, dict) and "predictions" in outer_preds:
            api_preds = outer_preds["predictions"]
        elif isinstance(outer_preds, list):
            api_preds = outer_preds
            
    if not api_preds:
        raise ValueError("No predictions returned in API response payload")
        
    if len(api_preds) != len(original_df):
        raise ValueError(
            f"Row mismatch: Input has {len(original_df)} rows, "
            f"but predictions have {len(api_preds)} rows."
        )
        
    pred_df = pd.DataFrame(api_preds)
        
    # 2. Strict row alignment logic
    # If a specific merge column 'id' is present in both, merge by ID to guarantee order and alignment
    # Otherwise, fallback to row-index concatenation.
    merge_col = "id"
    if merge_col in original_df.columns and merge_col in pred_df.columns:
        enriched_df = pd.merge(original_df, pred_df, on=merge_col, how="left")
    else:
        # Strict row alignment without ID: ensure lengths match
        if len(original_df) != len(pred_df):
            raise ValueError(
                f"Row mismatch: Input has {len(original_df)} rows, "
                f"but predictions have {len(pred_df)} rows."
            )
        
        # Reset index to guarantee correct horizontal concat
        df_left = original_df.reset_index(drop=True)
        df_right = pred_df.reset_index(drop=True)
        enriched_df = pd.concat([df_left, df_right], axis=1)

    # Convert all NaN values to None for clean JSON serialization
    input_data = original_df.where(pd.notnull(original_df), None).to_dict(orient="records")
    
    # Add row_index to predictions if not present, to fulfill structured output req
    structured_predictions = []
    for idx, row in pred_df.iterrows():
        structured_predictions.append({
            "row_index": idx,
            "predicted_label": row.get("predicted_label") or row.get("label"), 
            "probability": row.get("probability") or row.get("score")
        })

    enriched_output = enriched_df.where(pd.notnull(enriched_df), None).to_dict(orient="records")
    
    return TransformedOutput(
        input_data=input_data,
        predictions=structured_predictions,
        enriched_output=enriched_output,
        enriched_df=enriched_df
    )