from __future__ import annotations

import json
from pathlib import Path

import joblib
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.impute import SimpleImputer
from sklearn.metrics import accuracy_score, classification_report
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline

DATA_PATH = Path("data/data.csv")
MODEL_DIR = Path("model")
MODEL_PATH = MODEL_DIR / "model.joblib"
METRICS_PATH = MODEL_DIR / "metrics.json"

FEATURE_COLUMNS = [
    "customer_tenure",
    "total_spent",
    "last_purchase_days",
    "visit_count",
    "email_open_rate",
    "discount_usage",
    "support_tickets",
    "satisfaction_score",
]
TARGET_COLUMN = "lead_status"


def train() -> None:
    if not DATA_PATH.exists():
        raise FileNotFoundError(f"Dataset not found: {DATA_PATH}")

    df = pd.read_csv(DATA_PATH)
    X = df[FEATURE_COLUMNS].copy()
    y = df[TARGET_COLUMN].copy()

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

    preprocess = ColumnTransformer(
        transformers=[
            (
                "num",
                Pipeline([
                    ("imputer", SimpleImputer(strategy="median")),
                ]),
                FEATURE_COLUMNS,
            )
        ]
    )

    model = Pipeline([
        ("preprocess", preprocess),
        (
            "classifier",
            GradientBoostingClassifier(
                n_estimators=260,
                learning_rate=0.05,
                max_depth=3,
                random_state=42,
            ),
        ),
    ])

    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)

    accuracy = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, output_dict=True)

    classifier = model.named_steps["classifier"]
    feature_importance = {
        feature: float(importance)
        for feature, importance in zip(FEATURE_COLUMNS, classifier.feature_importances_)
    }

    metadata = {
        "feature_columns": FEATURE_COLUMNS,
        "target_column": TARGET_COLUMN,
        "class_names": sorted(df[TARGET_COLUMN].unique().tolist()),
        "feature_importance": feature_importance,
        "medians": {col: float(df[col].median()) for col in FEATURE_COLUMNS},
    }

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump({"pipeline": model, "metadata": metadata}, MODEL_PATH)

    with METRICS_PATH.open("w", encoding="utf-8") as f:
        json.dump(
            {
                "accuracy": accuracy,
                "classification_report": report,
                "train_rows": len(X_train),
                "test_rows": len(X_test),
            },
            f,
            indent=2,
        )

    print(f"Saved model to {MODEL_PATH}")
    print(f"Saved metrics to {METRICS_PATH}")
    print(f"Accuracy: {accuracy:.4f}")


if __name__ == "__main__":
    train()