import os
import json
import argparse
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score

from src.data import generate_synthetic_data, prepare_data
from src.model import ChurnClassifier

def train_pipeline(n_samples: int = 1500, test_size: float = 0.2, model_path: str = "models/churn_model.pkl", metrics_path: str = "models/metrics.json"):
    print("Generating synthetic dataset...")
    df = generate_synthetic_data(n_samples=n_samples, random_state=42)
    
    print("Preparing train and test splits...")
    X_train, X_test, y_train, y_test, preprocessor = prepare_data(df, test_size=test_size, random_state=42)
    
    print(f"Train set size: {X_train.shape[0]} samples, Test set size: {X_test.shape[0]} samples.")
    print("Fitting Random Forest Churn Classifier...")
    
    classifier = ChurnClassifier(n_estimators=100, max_depth=8, random_state=42)
    classifier.fit(X_train, y_train, preprocessor=preprocessor)
    
    print("Evaluating trained model on testing dataset...")
    y_pred = classifier.predict(X_test)
    y_prob = classifier.predict_proba(X_test)[:, 1]
    
    metrics = {
        "accuracy": float(accuracy_score(y_test, y_pred)),
        "precision": float(precision_score(y_test, y_pred, zero_division=0)),
        "recall": float(recall_score(y_test, y_pred, zero_division=0)),
        "f1_score": float(f1_score(y_test, y_pred, zero_division=0)),
        "roc_auc": float(roc_auc_score(y_test, y_prob))
    }
    
    print(f"Metrics: {json.dumps(metrics, indent=2)}")
    
    print(f"Saving model and preprocessor to {model_path}...")
    classifier.save(model_path)
    
    print(f"Saving evaluation metrics to {metrics_path}...")
    os.makedirs(os.path.dirname(os.path.abspath(metrics_path)), exist_ok=True)
    with open(metrics_path, 'w') as f:
        json.dump(metrics, f, indent=2)
        
    print("Training pipeline completed successfully.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train customer churn model.")
    parser.add_argument("--samples", type=int, default=1500, help="Number of synthetic samples to generate")
    parser.add_argument("--test_size", type=float, default=0.2, help="Proportion of dataset to use for test split")
    parser.add_argument("--model_path", type=str, default="models/churn_model.pkl", help="Filepath to save the trained model pickle")
    parser.add_argument("--metrics_path", type=str, default="models/metrics.json", help="Filepath to save the JSON metrics")
    
    args = parser.parse_args()
    train_pipeline(
        n_samples=args.samples,
        test_size=args.test_size,
        model_path=args.model_path,
        metrics_path=args.metrics_path
    )