"""
Train LightGBM churn model — Hybrid per-app + global fallback.

Usage:
  python train_model.py --app-id=1004          # Train for one app (if >=200 samples)
  python train_model.py --global               # Train the global fallback model
  python train_model.py --all                   # Train global + every eligible app
"""

import os
import sys
import argparse
import psycopg2
import pandas as pd
import lightgbm as lgb
import joblib
import json
from datetime import datetime
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score

from config import (
    CLICKHOUSE_CONFIG, CHURN_FEATURE_STORE_TABLE, CHURN_MODEL_REGISTRY_TABLE,
    DB_CONFIG, FEATURE_COLUMNS, PHASE2_FEATURE_COLUMNS,
    MODEL_DIR, MIN_TRAINING_SAMPLES, GLOBAL_MODEL_FILENAME,
    app_model_filename, DEFAULT_APP_ID,
)
from clickhouse_client import ClickHouseClient

EXPIRED_STATUS_CODE = 7


def get_db_connection():
    return psycopg2.connect(**DB_CONFIG)


def get_clickhouse_client():
    return ClickHouseClient(CLICKHOUSE_CONFIG)


def _load_training_data(conn, features, app_id=None):
    """Load joined feature + subscription data, optionally filtered by app_id."""
    if app_id is not None:
        query = """
        SELECT
            c.sms_cust_id,
            c.app_id,
            c.status
        FROM tbl_customer_subscriptions c
        WHERE c.app_id = %s
        """
        subscription_df = pd.read_sql_query(query, conn, params=(app_id,))
    else:
        query = """
        SELECT
            c.sms_cust_id,
            c.app_id,
            c.status
        FROM tbl_customer_subscriptions c
        """
        subscription_df = pd.read_sql_query(query, conn)

    clickhouse = get_clickhouse_client()
    feature_columns_sql = ", ".join(f"`{column}`" for column in ['sms_cust_id', 'app_id', *features])
    feature_query = f"SELECT {feature_columns_sql} FROM {clickhouse.qualified_table(CHURN_FEATURE_STORE_TABLE)}"
    if app_id is not None:
        feature_query += f" WHERE `app_id` = {int(app_id)}"

    feature_rows = clickhouse.select(feature_query)
    feature_df = pd.DataFrame(feature_rows)

    if subscription_df.empty or feature_df.empty:
        return pd.DataFrame(columns=['sms_cust_id', 'app_id', 'status', *features])

    subscription_df['sms_cust_id'] = subscription_df['sms_cust_id'].astype(str).str.strip()
    subscription_df['app_id'] = subscription_df['app_id'].astype(int)
    feature_df['sms_cust_id'] = feature_df['sms_cust_id'].astype(str).str.strip()
    feature_df['app_id'] = feature_df['app_id'].astype(int)

    df = subscription_df.merge(feature_df, how='inner', on=['sms_cust_id', 'app_id'])

    return df


def train_model(app_id=None):
    """
    Train a churn model.
    - app_id=None  → train a GLOBAL model (all apps combined)
    - app_id=int   → train an APP-SPECIFIC model (if enough data)
    """
    is_global = app_id is None
    label = "GLOBAL" if is_global else f"app_id={app_id}"
    print(f"\n{'='*60}")
    print(f"  Training {'GLOBAL' if is_global else 'APP-SPECIFIC'} model — {label}")
    print(f"{'='*60}")

    conn = get_db_connection()
    features = FEATURE_COLUMNS + PHASE2_FEATURE_COLUMNS

    df = _load_training_data(conn, features, app_id)

    # ── Check minimum sample threshold ───────────────────────────
    if not is_global and len(df) < MIN_TRAINING_SAMPLES:
        print(f"⚠  Only {len(df)} samples for {label} (minimum: {MIN_TRAINING_SAMPLES}).")
        print(f"   SKIPPING app-specific training — this app will use the global fallback.")
        conn.close()
        return False

    if len(df) < 50:
        print(f"Not enough data to train ({len(df)} rows). Need at least 50.")
        conn.close()
        return False

    # ── Prepare labels ───────────────────────────────────────────
    df['is_churned'] = (df['status'] == EXPIRED_STATUS_CODE).astype(int)
    df = df.fillna(0)

    X = df[features]
    y = df['is_churned']

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    print(f"Training LightGBM on {len(X_train)} samples with {len(features)} features...")
    model = lgb.LGBMClassifier(
        n_estimators=200,
        learning_rate=0.03,
        max_depth=7,
        is_unbalance=True,
        random_state=42,
        verbose=-1,
    )
    model.fit(X_train, y_train)

    # ── Evaluate ─────────────────────────────────────────────────
    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    accuracy  = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred, zero_division=0)
    recall    = recall_score(y_test, y_pred, zero_division=0)
    f1        = f1_score(y_test, y_pred, zero_division=0)
    auc_roc   = roc_auc_score(y_test, y_prob) if len(set(y_test)) > 1 else 0.5

    print(f"Results — Acc: {accuracy:.4f}, Prec: {precision:.4f}, "
          f"Rec: {recall:.4f}, F1: {f1:.4f}, AUC: {auc_roc:.4f}")

    # ── Save model ───────────────────────────────────────────────
    if is_global:
        pkl_name = GLOBAL_MODEL_FILENAME
        model_version = 'LGBM-Pipeline-v1'
        model_type = 'global'
    else:
        pkl_name = app_model_filename(app_id)
        model_version = 'LGBM-Pipeline-v1'
        model_type = 'app_specific'

    model_path = str(MODEL_DIR / pkl_name)
    model_data = {
        'model': model,
        'features': features,
        'numerical_features': features,
        'categorical_features': [],
    }
    joblib.dump(model_data, model_path)
    print(f"Model saved → {model_path}")

    # ── Update registry ──────────────────────────────────────────
    training_samples  = len(X_train)
    positive_samples  = int(y_train.sum())
    negative_samples  = training_samples - positive_samples
    feature_count     = len(features)
    feature_importance = json.dumps(
        dict(zip(features, model.feature_importances_.tolist()))
    )

    clickhouse = get_clickhouse_client()

    if is_global:
        clickhouse.command(f"""
            ALTER TABLE {clickhouse.qualified_table(CHURN_MODEL_REGISTRY_TABLE)}
            UPDATE `is_active` = 0
            WHERE `app_id` = 0 AND `model_type` = 'global'
            SETTINGS mutations_sync = 2
        """)
        target_app_id = 0
    else:
        clickhouse.command(f"""
            ALTER TABLE {clickhouse.qualified_table(CHURN_MODEL_REGISTRY_TABLE)}
            UPDATE `is_active` = 0
            WHERE `app_id` = {int(app_id)}
            SETTINGS mutations_sync = 2
        """)
        target_app_id = int(app_id)

    clickhouse.command(f"""
        ALTER TABLE {clickhouse.qualified_table(CHURN_MODEL_REGISTRY_TABLE)}
        DELETE WHERE `app_id` = {target_app_id}
          AND `model_version` = '{clickhouse.escape(model_version)}'
        SETTINGS mutations_sync = 2
    """)

    clickhouse.insert_json_rows(
        CHURN_MODEL_REGISTRY_TABLE,
        [{
            'app_id': target_app_id,
            'model_version': model_version,
            'model_type': model_type,
            'model_path': model_path,
            'accuracy': float(accuracy),
            'precision_score': float(precision),
            'recall_score': float(recall),
            'f1_score': float(f1),
            'auc_roc': float(auc_roc),
            'training_samples': training_samples,
            'positive_samples': positive_samples,
            'negative_samples': negative_samples,
            'feature_count': feature_count,
            'feature_importance': feature_importance,
            'is_active': 1,
            'trained_at': datetime.now(),
        }],
        columns=[
            'app_id', 'model_version', 'model_type', 'model_path',
            'accuracy', 'precision_score', 'recall_score', 'f1_score', 'auc_roc',
            'training_samples', 'positive_samples', 'negative_samples',
            'feature_count', 'feature_importance', 'is_active', 'trained_at',
        ],
    )
    print(f"✅ Registry updated — {model_type} model for {label}.")

    conn.close()
    return True


def discover_app_ids():
    """Get all distinct app_ids from the feature store."""
    clickhouse = get_clickhouse_client()
    rows = clickhouse.select(f"""
        SELECT DISTINCT toInt32(`app_id`) AS app_id
        FROM {clickhouse.qualified_table(CHURN_FEATURE_STORE_TABLE)}
        ORDER BY `app_id`
    """)
    return [int(row['app_id']) for row in rows if int(row.get('app_id', 0)) > 0]


def train_all():
    """Train global model first, then per-app models for eligible apps."""
    print("\n" + "▓" * 60)
    print("  HYBRID TRAINING — Global + All Eligible Apps")
    print("▓" * 60)

    # Step 1: Train global model
    train_model(app_id=None)

    # Step 2: Train per-app models
    app_ids = discover_app_ids()
    print(f"\nDiscovered {len(app_ids)} app_ids: {app_ids}")

    trained = []
    skipped = []
    for aid in app_ids:
        success = train_model(app_id=aid)
        if success:
            trained.append(aid)
        else:
            skipped.append(aid)

    print(f"\n{'='*60}")
    print(f"  SUMMARY")
    print(f"  Global model : ✅ Trained")
    print(f"  App-specific : {len(trained)} trained {trained}")
    print(f"  Skipped (<{MIN_TRAINING_SAMPLES} samples) : {len(skipped)} {skipped}")
    print(f"{'='*60}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train churn model (hybrid)")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--app-id", type=int, help="Train model for a specific app_id")
    group.add_argument("--global", dest="train_global", action="store_true",
                       help="Train the global fallback model only")
    group.add_argument("--all", dest="train_all", action="store_true",
                       help="Train global + all eligible app-specific models")
    args = parser.parse_args()

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