import pandas as pd
import pickle
import json
from pathlib import Path

try:
    from xgboost import XGBClassifier  # type: ignore
except Exception:  # pragma: no cover
    XGBClassifier = None  # type: ignore
import logging

# Setup logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def train_model():
    """Train XGBoost model for CVE severity classification"""
    logger.info("Loading training dataset...")

    # Path to training dataset
    data_path = "data/balanced-training-dataset-through-smote.csv"

    df = pd.read_csv(data_path)
    logger.info(f"Loaded {len(df)} CVEs from {data_path}")

    # Check required columns
    required_columns = ["cve_id", "severity_label", "severity_numeric"]
    missing_columns = [col for col in required_columns if col not in df.columns]

    if missing_columns:
        logger.error(f"Missing required columns: {missing_columns}")
        logger.info(f"Available columns: {list(df.columns)}")
        return

    # Convert boolean features to integers (0/1) if needed
    feature_columns = [
        col
        for col in df.columns
        if col not in ["cve_id", "severity_label", "severity_numeric"]
    ]

    logger.info(f"Found {len(feature_columns)} feature columns")

    # Handle any non-numeric features
    for col in feature_columns:
        if df[col].dtype == "object":
            logger.info(f"Converting column '{col}' from object to numeric")
            df[col] = pd.to_numeric(df[col], errors="coerce").fillna(0)
        df[col] = df[col].astype(int)

    # Remove any rows with missing target values
    initial_count = len(df)
    df = df.dropna(subset=["severity_numeric"])
    if len(df) < initial_count:
        logger.info(
            f"Removed {initial_count - len(df)} rows with missing severity_numeric values"
        )

    # Prepare features and target
    X = df[feature_columns]
    y = df["severity_numeric"].astype(int)

    logger.info(f"Features: {len(feature_columns)} columns")
    logger.info(f"Training samples: {len(X)}")
    logger.info(f"Target distribution:\n{y.value_counts().sort_index()}")

    # Check if we have enough samples for each class
    class_counts = y.value_counts()
    if len(class_counts) < 2:
        logger.warning("Only one class found in target variable!")

    # Train XGBoost model
    logger.info("Training XGBoost model...")

    custom_class_weight = {0: 1.0, 1: 1.0, 2: 1.0}

    model = XGBClassifier(
        n_estimators=500,
        max_depth=6,
        learning_rate=0.5,
        min_child_weight=1,
        subsample=0.80,
        colsample_bytree=0.8,
        reg_alpha=0.1,
        reg_lambda=1.0,
        random_state=42,
        eval_metric="mlogloss",
        verbosity=1,
        class_weight=custom_class_weight,
    )

    try:
        model.fit(X, y)
        logger.info("Training complete!")
    except Exception as e:
        logger.error(f"Training failed: {e}")
        return

    # Create models directory
    models_dir = Path("models")
    models_dir.mkdir(exist_ok=True)
    logger.info(f"Created models directory: {models_dir.absolute()}")

    # Save the model
    model_path = models_dir / "cve_severity_model.pkl"
    with open(model_path, "wb") as f:
        pickle.dump(model, f)
    logger.info(f"Model saved to: {model_path.absolute()}")

    # Save metadata
    label_mapping = {0: "IMPORTANT", 1: "MODERATE", 2: "LOW"}
    metadata = {
        "feature_columns": feature_columns,
        "label_mapping": label_mapping,
        "model_type": "XGBClassifier",
        "num_features": len(feature_columns),
        "training_dataset": str(data_path),
        "trained_on_samples": len(df),
        "class_distribution": y.value_counts().to_dict(),
        "model_params": {
            "n_estimators": 200,
            "max_depth": 6,
            "learning_rate": 0.1,
            "random_state": 42,
        },
    }

    metadata_path = models_dir / "model_metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)
    logger.info(f"Metadata saved to: {metadata_path.absolute()}")

    logger.info("✅ Model training complete!")
    logger.info(f"Trained on {len(df)} CVEs with {len(feature_columns)} features")

    # Show feature importance
    if hasattr(model, "feature_importances_"):
        feature_importance = pd.DataFrame(
            {"feature": feature_columns, "importance": model.feature_importances_}
        ).sort_values("importance", ascending=False)

        logger.info("Top 10 most important features:")
        for i, row in feature_importance.head(10).iterrows():
            logger.info(f"  {row['feature']}: {row['importance']:.4f}")

        # Save feature importance
        importance_path = models_dir / "feature_importance.csv"
        feature_importance.to_csv(importance_path, index=False)
        logger.info(f"Feature importance saved to: {importance_path.absolute()}")


if __name__ == "__main__":
    train_model()
