"""
Run batch predictions for all active customers.
Hybrid model loading: tries app-specific model first, falls back to global.
Writes results to churn_predictions table, including SHAP explanations.
"""
import sys
import os
import argparse
import joblib
import pandas as pd
import numpy as np
import json
from datetime import datetime
from config import (
    CLICKHOUSE_CONFIG, CHURN_FEATURE_STORE_TABLE, CHURN_MODEL_REGISTRY_TABLE,
    CHURN_PREDICTIONS_TABLE, DEFAULT_APP_ID, MODEL_DIR,
    GLOBAL_MODEL_FILENAME, app_model_filename,
    FEATURE_COLUMNS, PHASE2_FEATURE_COLUMNS,
)
from clickhouse_client import ClickHouseClient

USE_SHAP_DEFAULT = str(os.getenv('AI_CHURN_USE_SHAP', '0')).lower() in {'1', 'true', 'yes', 'on'}


def get_clickhouse_client():
    return ClickHouseClient(CLICKHOUSE_CONFIG)


def _load_model(app_id):
    """
    Try to load the app-specific model for app_id.
    If not found, fall back to the global model.
    Returns (model_data, model_version, model_type).
    """
    # 1. Try app-specific model file
    app_model_path = MODEL_DIR / app_model_filename(app_id)
    if app_model_path.exists():
        print(f"  ✅ Loaded APP-SPECIFIC model for app_id={app_id}")
        model_data = joblib.load(app_model_path)
        # Get version from registry
        model_version = _get_registry_version(app_id, 'app_specific')
        return model_data, model_version or 'LGBM-Pipeline-v1', 'app_specific'

    # 2. Fall back to global model
    global_model_path = MODEL_DIR / GLOBAL_MODEL_FILENAME
    if global_model_path.exists():
        print(f"  ⚠  No app-specific model for app_id={app_id}. Using GLOBAL fallback.")
        model_data = joblib.load(global_model_path)
        model_version = _get_registry_version(None, 'global')
        return model_data, model_version or 'LGBM-Pipeline-v1', 'global'

    # 3. Legacy: try old model file or registry
    legacy_path = MODEL_DIR / 'lgbm_model_v1.pkl'
    if legacy_path.exists():
        print(f"  ⚠  Using LEGACY model file (lgbm_model_v1.pkl)")
        model_data = joblib.load(legacy_path)
        return model_data, 'LGBM-Pipeline-v1', 'global'

    return None, None, None


def _get_registry_version(app_id, model_type):
    """Fetch model_version from registry for given app_id and model_type."""
    try:
        client = get_clickhouse_client()
        if app_id is not None:
            rows = client.select(f"""
                SELECT toString(`model_version`) AS model_version
                FROM {client.qualified_table(CHURN_MODEL_REGISTRY_TABLE)}
                WHERE `app_id` = {int(app_id)}
                  AND `model_type` = '{client.escape(model_type)}'
                  AND `is_active` = 1
                ORDER BY `trained_at` DESC
                LIMIT 1
            """)
        else:
            rows = client.select(f"""
                SELECT toString(`model_version`) AS model_version
                FROM {client.qualified_table(CHURN_MODEL_REGISTRY_TABLE)}
                WHERE `app_id` = 0
                  AND `model_type` = '{client.escape(model_type)}'
                  AND `is_active` = 1
                ORDER BY `trained_at` DESC
                LIMIT 1
            """)
        return rows[0]['model_version'] if rows else None
    except Exception:
        return None


def _extract_model_payload(model_data):
    """
    Normalize model payload across new and legacy artifact formats.
    New format: dict with model + feature metadata.
    Legacy format: raw estimator object (no metadata).
    """
    if isinstance(model_data, dict) and 'model' in model_data:
        model = model_data['model']
        expected_features = model_data.get('features') or (FEATURE_COLUMNS + PHASE2_FEATURE_COLUMNS)
        numerical_features = model_data.get('numerical_features') or expected_features
        categorical_features = model_data.get('categorical_features') or []
        return model, expected_features, numerical_features, categorical_features

    # Legacy artifact fallback: model object only.
    legacy_features = []
    if hasattr(model_data, "feature_name_"):
        try:
            legacy_features = [str(name) for name in list(model_data.feature_name_) if str(name)]
        except Exception:
            legacy_features = []
    if not legacy_features and hasattr(model_data, "booster_"):
        try:
            legacy_features = [str(name) for name in model_data.booster_.feature_name() if str(name)]
        except Exception:
            legacy_features = []

    features = legacy_features or (FEATURE_COLUMNS + PHASE2_FEATURE_COLUMNS)
    return model_data, features, features, []


def run_predictions(app_id):
    """Score all active customers with the best available model."""

    print(f"\n{'='*60}")
    print(f"  Churn Prediction — app_id={app_id}")
    print(f"{'='*60}")

    try:
        client = get_clickhouse_client()
    except Exception as e:
        print(f"ClickHouse connection error: {e}", file=sys.stderr)
        return

    # 1. Load model (hybrid: app-specific → global fallback)
    model_data, model_version, model_type = _load_model(app_id)
    if model_data is None:
        print("❌ No trained model found. Run train_model.py first.")
        return

    model, expected_features, numerical_features, categorical_features = _extract_model_payload(model_data)

    print(f"  Model version : {model_version}")
    print(f"  Model type    : {model_type}")

    # 2. Load features for all active customers
    query = """
    SELECT *
    FROM {table}
    WHERE `app_id` = {app_id}
      AND trimBoth(toString(ifNull(`sms_cust_id`, ''))) != ''
    """.format(
        table=client.qualified_table(CHURN_FEATURE_STORE_TABLE),
        app_id=int(app_id),
    )
    feature_rows = client.select(query)
    df = pd.DataFrame(feature_rows)

    if df.empty:
        print("No features found. Run feature extraction first.")
        return

    # Extract ID columns as strings to avoid aborting an entire app on bad casts.
    ids = df[['sms_cust_id', 'app_id']].copy()
    ids['sms_cust_id'] = ids['sms_cust_id'].astype(str).str.strip()

    # Process features to match training expectations
    X = df.reindex(columns=numerical_features, fill_value=0).fillna(0)

    if categorical_features:
        X = pd.get_dummies(X, columns=categorical_features)
        for col in expected_features:
            if col not in X.columns:
                X[col] = 0
        X = X[expected_features]

    print(f"  Scoring {len(X)} customers...")

    # 3. Predict probabilities
    probabilities = model.predict_proba(X)[:, 1]

    shap_values = None
    feature_importance = None

    # 4. Generate explanations (optional SHAP; fallback to feature-importance scoring)
    if USE_SHAP_DEFAULT:
        try:
            import shap  # Lazy import: keeps inference usable in minimal environments.
            print("  Computing SHAP explanations...")
            explainer = shap.TreeExplainer(model)
            shap_values = explainer.shap_values(X)
            if isinstance(shap_values, list):
                shap_values = shap_values[1]
        except Exception as exc:
            print(f"  SHAP disabled due to runtime error: {exc}")

    if shap_values is None:
        if hasattr(model, 'feature_importances_'):
            fi = np.asarray(model.feature_importances_, dtype=float)
            if fi.shape[0] == len(expected_features):
                feature_importance = fi / (fi.max() if fi.max() > 0 else 1.0)

    # 5. Build prediction records
    predictions = []

    for idx in range(len(X)):
        prob = float(probabilities[idx])
        risk_level = get_risk_level(prob)

        raw_impacts = {}
        row_values = {}
        shap_row = shap_values[idx] if shap_values is not None else None
        x_row = X.iloc[idx]

        for i, col in enumerate(expected_features):
            val = x_row[col]
            if shap_row is not None:
                impact = shap_row[i]
            else:
                # Fallback impact estimate: normalized feature importance * value scale.
                val_num = 0.0
                try:
                    val_num = float(val) if val is not None else 0.0
                except Exception:
                    val_num = 0.0
                if feature_importance is not None:
                    impact = float(feature_importance[i] * np.tanh(val_num / 10.0))
                else:
                    impact = 0.0

            is_dummy = False
            for cat_feat in categorical_features:
                if col.startswith(f"{cat_feat}_"):
                    is_dummy = True
                    if val == 1:
                        cat_val = col.replace(f"{cat_feat}_", "")
                        raw_impacts[cat_feat] = raw_impacts.get(cat_feat, 0) + impact
                        row_values[cat_feat] = cat_val
                    break

            if not is_dummy:
                raw_impacts[col] = impact
                row_values[col] = val

        feature_impacts = sorted(
            [(k, v) for k, v in raw_impacts.items()],
            key=lambda item: abs(item[1]),
            reverse=True
        )[:5]

        top_factors = []
        for feat, impact in feature_impacts:
            val = row_values.get(feat, 0)
            top_factors.append({
                "feature": feat,
                "impact": round(float(impact), 4),
                "value": round(float(val), 2) if isinstance(val, (int, float)) and not np.isnan(val) else val,
                "direction": "increases_risk" if impact > 0 else "decreases_risk"
            })

        actions = generate_recommendations(top_factors, prob, risk_level)

        # Store which model was used for this prediction
        version_tag = f"{model_version} ({model_type})"

        sms_cust_id = ids.iloc[idx]['sms_cust_id']
        if not sms_cust_id:
            # Defensive guard for any residual empty id row.
            continue

        predictions.append((
            sms_cust_id,
            int(ids.iloc[idx]['app_id']),
            prob,
            risk_level,
            None,  # predicted_churn_date placeholder
            float(min(1.0, abs(prob - 0.5) * 2)),
            json.dumps(top_factors),
            json.dumps(actions),
            version_tag,
        ))

    # 6. Insert predictions
    print(f"  Saving {len(predictions)} predictions to database...")

    if not predictions:
        print("No valid prediction rows to save after ID sanitization.")
        return

    try:
        delete_ids = [sms_cust_id for sms_cust_id, _, *_ in predictions]
        for chunk_start in range(0, len(delete_ids), 1000):
            chunk = delete_ids[chunk_start:chunk_start + 1000]
            if not chunk:
                continue
            quoted_ids = ", ".join(f"'{client.escape(value)}'" for value in chunk)
            client.command(f"""
                ALTER TABLE {client.qualified_table(CHURN_PREDICTIONS_TABLE)}
                DELETE WHERE `app_id` = {int(app_id)}
                  AND `sms_cust_id` IN ({quoted_ids})
                SETTINGS mutations_sync = 2
            """)

        now_string = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        payload_rows = []
        for prediction in predictions:
            payload_rows.append({
                'sms_cust_id': prediction[0],
                'app_id': prediction[1],
                'churn_probability': prediction[2],
                'risk_level': prediction[3],
                'predicted_churn_date': prediction[4],
                'confidence': prediction[5],
                'top_factors': prediction[6],
                'recommended_actions': prediction[7],
                'model_version': prediction[8],
                'predicted_at': now_string,
            })

        for chunk_start in range(0, len(payload_rows), 1000):
            chunk = payload_rows[chunk_start:chunk_start + 1000]
            client.insert_json_rows(
                CHURN_PREDICTIONS_TABLE,
                chunk,
                columns=[
                    'sms_cust_id', 'app_id', 'churn_probability', 'risk_level',
                    'predicted_churn_date', 'confidence', 'top_factors',
                    'recommended_actions', 'model_version', 'predicted_at',
                ],
            )

        critical = sum(1 for p in predictions if p[3] == 'critical')
        high = sum(1 for p in predictions if p[3] == 'high')
        print(f"  ✅ Done! {critical} Critical Risk, {high} High Risk. Model used: {model_type}")
    except Exception as e:
        print(f"Database error during insert: {e}")


def get_risk_level(probability):
    if probability >= 0.80:
        return 'high'
    elif probability >= 0.50:
        return 'medium'
    else:
        return 'low'


def generate_recommendations(top_factors, probability, risk_level):
    actions = []
    factor_names = {f['feature'] for f in top_factors if f['direction'] == 'increases_risk'}

    if 'failed_payment_count_30d' in factor_names or 'failed_payment_count_90d' in factor_names:
        actions.append("Trigger payment method update sequence")
    if 'days_until_expiry' in factor_names:
        actions.append("Send early renewal email with promotional code")
    if 'auto_renew_enabled' in factor_names:
        actions.append("Offer 1 month free to enable auto-renewal")
    if 'has_downgraded' in factor_names:
        actions.append("Promote features they miss from higher tier")
    if 'days_since_last_login' in factor_names:
        actions.append("Send re-engagement campaign ('We Miss You')")
    if 'customer_tenure_days' in factor_names:
        actions.append("Send 'Anniversary' or 'Milestone' appreciation gift")

    if risk_level == 'high':
        actions.insert(0, "Offer 10-20% proactive retention discount")

    if not actions:
        actions.append("Monitor account for further degradation")

    return actions[:4]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate batch predictions (hybrid)")
    parser.add_argument("--app-id", type=int, default=DEFAULT_APP_ID, help="App ID")
    args = parser.parse_args()

    if args.app_id <= 0:
        print("Error: --app-id must be positive.", file=sys.stderr)
        sys.exit(1)

    run_predictions(args.app_id)