import base64
import pickle
import pandas as pd
import numpy as np
import json
from typing import List, Dict, Any
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score, roc_auc_score, precision_score, recall_score, confusion_matrix
import xgboost as xgb


def _is_id_column(series: pd.Series) -> bool:
    if series.dtype == object:
        n_unique = series.nunique()
        n_rows = len(series)
        if n_unique / n_rows > 0.8 or n_unique > 100:
            return True
        if series.str.match(r"^[A-Z]+-\d+$").all():
            return True
    return False


def _is_categorical(series: pd.Series) -> bool:
    if series.dtype == object or series.dtype.name == "category":
        n_unique = series.nunique()
        if 2 <= n_unique <= 20:
            return True
    return False


def train_from_data(data: List[Dict[str, Any]], target_col: str, model_type: str = "churn") -> dict:
    df = pd.DataFrame(data)
    y = df[target_col]
    X = df.drop(columns=[target_col])

    skipped = []
    encoders = {}
    X_processed = pd.DataFrame(index=X.index)

    for col in X.columns:
        if _is_id_column(X[col]):
            skipped.append(col)
            continue
        if _is_categorical(X[col]):
            le = LabelEncoder()
            values = X[col].astype(str).fillna("__missing__")
            le.fit(values)
            X_processed[col] = le.transform(values)
            encoders[col] = {"classes_": le.classes_.tolist()}
        elif pd.api.types.is_numeric_dtype(X[col]):
            X_processed[col] = X[col].fillna(X[col].median() if not X[col].isna().all() else 0)
        else:
            skipped.append(col)

    if len(X_processed.columns) == 0:
        raise ValueError("No usable features found.")

    if y.nunique() < 2:
        raise ValueError(f"Target column needs both positive and negative examples.")

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

    pos = int(y_train.sum()) if hasattr(y_train, 'sum') else int((y_train == 1).sum())
    neg = len(y_train) - pos
    scale_weight = neg / max(pos, 1)

    model = xgb.XGBClassifier(
        n_estimators=100, max_depth=5, learning_rate=0.08,
        subsample=0.8, colsample_bytree=0.8,
        scale_pos_weight=scale_weight,
        random_state=42, eval_metric="logloss"
    )
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]

    model_b64 = base64.b64encode(pickle.dumps(model)).decode()

    metrics = {
        "accuracy": round(float(accuracy_score(y_test, y_pred)), 4),
        "roc_auc": round(float(roc_auc_score(y_test, y_proba)), 4),
        "precision": round(float(precision_score(y_test, y_pred, zero_division=0)), 4),
        "recall": round(float(recall_score(y_test, y_pred, zero_division=0)), 4),
        "confusion_matrix": confusion_matrix(y_test, y_pred).tolist(),
        "feature_importance": dict(
            zip(X_processed.columns.tolist(), [float(v) for v in model.feature_importances_])
        ),
    }

    f1_denom = metrics["precision"] + metrics["recall"]
    metrics["f1"] = round(2 * metrics["precision"] * metrics["recall"] / max(f1_denom, 1e-9), 4)

    summary = {
        "n_rows": len(df),
        "n_features": len(X_processed.columns),
        "target_col": target_col,
        "imbalance_ratio": round(neg / max(pos, 1), 1),
        "skipped_columns": skipped,
    }

    return {
        "model_binary": model_b64,
        "feature_names": X_processed.columns.tolist(),
        "encoders_json": json.dumps(encoders),
        "metrics": metrics,
        "metrics_json": json.dumps(metrics),
        "summary": summary,
        "summary_json": json.dumps(summary),
        "n_features": len(X_processed.columns),
        "n_rows": len(X),
    }


def predict_with_model(model_b64: str, feature_names: list, encoders_json: str,
                       data: List[Dict[str, Any]]) -> list:
    model = pickle.loads(base64.b64decode(model_b64))
    encoders_data = json.loads(encoders_json)
    df = pd.DataFrame(data)

    for col, enc_info in encoders_data.items():
        if col in df.columns:
            classes = enc_info["classes_"]
            mapping = {c: i for i, c in enumerate(classes)}
            df[col] = df[col].astype(str).map(lambda x: mapping.get(x, -1))

    for col in feature_names:
        if col not in df.columns:
            df[col] = 0

    X = df[feature_names].fillna(0)
    probs = model.predict_proba(X)[:, 1]
    scores = [round(float(p * 100), 1) for p in probs]

    return scores