# app/pipelines/risk_model.py
from __future__ import annotations

from typing import Dict, Any
from pathlib import Path
from io import BytesIO

import numpy as np
import pandas as pd
import joblib
from google.cloud import storage

from config import (
    MODEL_PATH,
    GCP_PROJECT_ID,
    GCS_BUCKET_NAME,
    VALIDATION_RESULTS_PREFIX,
)

_model = None
_preprocess = None


def _validation_blob_name(applicant_id: str) -> str:
    return f"{VALIDATION_RESULTS_PREFIX}applicant_{applicant_id}_validated.csv"


def _load_validated_df(applicant_id: str) -> pd.DataFrame:
    """Load validated results CSV for a single applicant from GCS."""
    blob_name = _validation_blob_name(applicant_id)
    bucket = storage_client.bucket(GCS_BUCKET_NAME)
    blob = bucket.blob(blob_name)
    data = blob.download_as_bytes()
    return pd.read_csv(BytesIO(data))
storage_client = storage.Client(project=GCP_PROJECT_ID) if GCP_PROJECT_ID else storage.Client()


def load_model():
    global _model
    if _model is None:
        _model = joblib.load(MODEL_PATH)
    return _model


def load_preprocess():
    """Load preprocessing pipeline used during training."""
    global _preprocess
    if _preprocess is None:
        preprocess_path = Path(__file__).resolve().parent.parent / "artifacts" / "preprocess_with_target_passthrough.joblib"
        _preprocess = joblib.load(preprocess_path)
    return _preprocess


_FIELD_MAP = {
    "Annual Income": "Annual_Income",
    "Applicant Age": "Applicant_Age",
    "Marital Status": "Marital_Status",
    "House Ownership": "House_Ownership",
    "Occupation": "Occupation",
    "Residence City": "Residence_City",
    "Residence State": "Residence_State",
    "Years in Current Employment": "Years_in_Current_Employment",
}
NUM_FEATURES = ["Annual_Income", "Applicant_Age", "Years_in_Current_Employment"]
STR_FEATURES = ["Marital_Status", "House_Ownership", "Occupation", "Residence_City", "Residence_State"]


def _fields_to_feature_df(validated_df: pd.DataFrame) -> pd.DataFrame:
    """Pivot validated fields into one row with expected feature names."""
    filtered = validated_df[validated_df["field_name"].isin(_FIELD_MAP.keys())].copy()
    if filtered.empty:
        raise ValueError("No required fields found in validated data.")

    # Drop duplicate field_names, keeping the first occurrence per applicant
    filtered = filtered.sort_index()
    filtered = filtered.drop_duplicates(subset=["applicant_id", "field_name"], keep="first")

    # Normalize field names to improve mapping (trim/case-insensitive) and fallback to originals if unmapped
    normalized_map = {k.lower(): v for k, v in _FIELD_MAP.items()}
    field_series = filtered["field_name"].astype(str).str.strip()
    mapped = field_series.str.lower().map(normalized_map)
    filtered["field_name"] = mapped.fillna(field_series)

    pivot = filtered.pivot(
        index="applicant_id",
        columns="field_name",
        values="value",
    ).reset_index()

    # Ensure column order and include applicant_id
    desired_cols = ["applicant_id"] + list(_FIELD_MAP.values())
    pivot = pivot.reindex(columns=desired_cols)

    num_cols = [c for c in NUM_FEATURES if c in pivot.columns]
    str_cols = [c for c in STR_FEATURES if c in pivot.columns]

    if num_cols:
        for col in num_cols:
            pivot[col] = (
                pd.to_numeric(pivot[col], errors="coerce")
                .fillna(-1)
                .astype(int)
            )

    if str_cols:
        for col in str_cols:
            pivot[col] = pivot[col].fillna("missing").astype(str)

    print("Pivoted dataframe before converting column names")
    print(pivot.head())
    print("=" * 50)
    print(pivot.info())

    return pivot


def predict_risk_for_applicant(applicant_id: str, validated_df: pd.DataFrame | None = None) -> Dict[str, Any]:
    """Load validated data for applicant, pivot features, and return default probability."""
    if validated_df is None:
        validated_df = _load_validated_df(applicant_id)

    # ensure we only use this applicant's rows
    validated_df = validated_df[validated_df["applicant_id"].astype(str) == str(applicant_id)]
    if validated_df.empty:
        raise ValueError(f"No validated data found for applicant_id={applicant_id}")

    model = load_model()
    preprocess = load_preprocess()
    features_df = _fields_to_feature_df(validated_df)

    try:
        # Drop applicant_id before preprocessing
        feature_input = features_df.drop(columns=["applicant_id"], errors="ignore")

        # Determine expected columns and target passthrough cols from the preprocessor
        expected_cols = list(getattr(preprocess, "feature_names_in_", feature_input.columns))
        target_cols: list[str] = []
        for name, _, cols in getattr(preprocess, "transformers_", []):
            if name == "target_passthrough":
                target_cols.extend(cols)

        # Add any missing columns with sensible defaults
        for col in expected_cols:
            if col not in feature_input.columns:
                feature_input[col] = 0 if col in target_cols else np.nan

        feature_input = feature_input.reindex(columns=expected_cols)
    except Exception as exc:
        raise RuntimeError(f"Feature alignment failed: {exc}")

    try:
        X_pre = preprocess.transform(feature_input)
    except Exception as exc:
        raise RuntimeError(f"Preprocess transform failed: {exc}")

    # If the preprocessor includes a target passthrough, drop it before model prediction
    try:
        feature_names_out = list(preprocess.get_feature_names_out())
        target_indices = []
        for name, _, cols in getattr(preprocess, "transformers_", []):
            if name == "target_passthrough":
                for col in cols:
                    candidate = f"target_passthrough__{col}"
                    if candidate in feature_names_out:
                        target_indices.append(feature_names_out.index(candidate))

        X_for_model = X_pre
        if target_indices:
            X_for_model = np.delete(X_pre, target_indices, axis=1)

        # Ensure numeric matrix for the model; if this fails, surface the offending values
        try:
            X_for_model = np.asarray(X_for_model, dtype=float)
        except Exception as exc_inner:
            raise RuntimeError(f"Preparing model input failed during dtype cast: {exc_inner}; sample values={X_for_model[:5]}")
    except Exception as exc:
        raise RuntimeError(f"Preparing model input failed: {exc}")

    try:
        proba = model.predict_proba(X_for_model)[0, 1]
    except Exception as exc:
        raise RuntimeError(f"Model prediction failed: {exc}; X_for_model shape={getattr(X_for_model, 'shape', None)}")

    return {
        "default_probability": float(proba),
    }
