"""ScamShield NLP — Model training script.
This script retrains the calibrated Linear SVM model and TF-IDF vectorizer
using scikit-learn 1.4, compatible with Python 3.10. It uses publicly available
spam datasets to replicate the original model's performance.

Usage:
    python train_model.py
"""

import os
import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import LinearSVC
from sklearn.calibration import CalibratedClassifierCV
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score

# Import preprocessing function to ensure consistency
from preprocessing import clean_text

# Create models directory if it doesn't exist
MODELS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
os.makedirs(MODELS_DIR, exist_ok=True)

def load_and_prepare_datasets():
    """Load and combine publicly available spam datasets.
    Downloads and processes the SMS Spam Collection and other public datasets
    to create a training corpus similar to the original ScamShield dataset.
    """
    print("Loading datasets...")
    
    # First, load the UCI SMS Spam Collection (most widely used spam dataset)
    # We'll download it and process it
    try:
        # Try to fetch from UCI repository using ucimlrepo
        from ucimlrepo import fetch_ucirepo
        sms_spam = fetch_ucirepo(id=228)
        X = sms_spam.data.features
        y = sms_spam.data.targets
        
        # Convert to DataFrame
        df = pd.DataFrame({
            'text': X['v2'],
            'label': y['v1'].map({'ham': 0, 'spam': 1})
        })
        print(f"Loaded UCI SMS Spam Collection: {len(df)} messages")
        
    except Exception as e:
        print(f"Could not load from ucimlrepo, falling back to downloading directly: {e}")
        # Alternative: download the dataset directly
        import requests
        url = "https://archive.ics.uci.edu/ml/machine-learning-databases/00228/smsspamcollection.zip"
        import zipfile
        import io
        
        r = requests.get(url)
        z = zipfile.ZipFile(io.BytesIO(r.content))
        with z.open('SMSSpamCollection') as f:
            lines = f.readlines()
            data = []
            for line in lines:
                line = line.decode('utf-8').strip()
                if line:
                    label, text = line.split('\t', 1)
                    data.append({'text': text, 'label': 1 if label == 'spam' else 0})
            df = pd.DataFrame(data)
            print(f"Loaded UCI SMS Spam Collection (direct download): {len(df)} messages")

    # Let's also add some additional phishing data to enrich the dataset
    # Try to load the phishing dataset from Hugging Face if available, else continue with what we have
    try:
        from datasets import load_dataset
        hf_dataset = load_dataset("ealvaradob/phishing-dataset", "combined", split="train")
        hf_df = pd.DataFrame(hf_dataset)
        # Only keep text and label columns, filter where label is 0 or 1
        hf_df = hf_df[hf_df['label'].isin([0, 1])][['text', 'label']]
        # Combine with our existing dataframe
        df = pd.concat([df, hf_df], ignore_index=True)
        print(f"Added Hugging Face phishing dataset: {len(hf_df)} messages, total now: {len(df)}")
    except Exception as e:
        print(f"Could not load Hugging Face dataset, continuing with SMS dataset only: {e}")

    # Remove duplicates
    df = df.drop_duplicates(subset=['text']).reset_index(drop=True)
    print(f"After removing duplicates: {len(df)} messages")
    
    # Apply the same preprocessing used in inference
    print("Applying text preprocessing...")
    df['cleaned_text'] = df['text'].apply(clean_text)
    
    # Remove any empty texts after preprocessing
    df = df[df['cleaned_text'].str.strip() != ''].reset_index(drop=True)
    print(f"Final dataset size after cleaning: {len(df)} messages")
    print(f"Legitimate messages (0): {len(df[df['label'] == 0])}, Scam messages (1): {len(df[df['label'] == 1])}")
    
    return df

def train_and_evaluate_model(df):
    """Train the TF-IDF vectorizer and calibrated LinearSVC model, then evaluate."""
    print("\nStarting model training...")
    
    # Split into train and test sets
    X_train, X_test, y_train, y_test = train_test_split(
        df['cleaned_text'], 
        df['label'], 
        test_size=0.2, 
        random_state=42, 
        stratify=df['label']
    )
    print(f"Train set size: {len(X_train)}, Test set size: {len(X_test)}")
    
    # Initialize and fit TF-IDF Vectorizer with the same parameters as original
    print("\nFitting TF-IDF vectorizer...")
    vectorizer = TfidfVectorizer(
        ngram_range=(1, 2),
        max_features=10000,
        min_df=2,
        sublinear_tf=True,
        lowercase=True
    )
    X_train_tfidf = vectorizer.fit_transform(X_train)
    X_test_tfidf = vectorizer.transform(X_test)
    print(f"TF-IDF vectorizer fitted. Vocabulary size: {len(vectorizer.vocabulary_)}")
    
    # Train LinearSVC and calibrate with CalibratedClassifierCV (same as original)
    print("\nTraining LinearSVC with calibration...")
    base_svm = LinearSVC(C=1.0, random_state=42, max_iter=10000)
    calibrated_svm = CalibratedClassifierCV(base_svm, cv=5)
    calibrated_svm.fit(X_train_tfidf, y_train)
    print("Model training complete.")
    
    # Evaluate on test set
    print("\nEvaluating model on test set...")
    y_pred = calibrated_svm.predict(X_test_tfidf)
    y_pred_proba = calibrated_svm.predict_proba(X_test_tfidf)[:, 1]
    
    # Calculate metrics
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)
    roc_auc = roc_auc_score(y_test, y_pred_proba)
    
    print(f"Model Performance Metrics:")
    print(f"  Accuracy:  {accuracy:.4f} ({accuracy*100:.2f}%)")
    print(f"  Precision: {precision:.4f} ({precision*100:.2f}%)")
    print(f"  Recall:    {recall:.4f} ({recall*100:.2f}%)")
    print(f"  F1-score:  {f1:.4f} ({f1*100:.2f}%)")
    print(f"  ROC-AUC:   {roc_auc:.4f} ({roc_auc*100:.2f}%)")
    
    # Calculate number of missed scams (false negatives)
    fn_count = np.sum((y_test == 1) & (y_pred == 0))
    print(f"  Missed scams (FN): {fn_count} / {len(y_test[y_test == 1])} ({fn_count/len(y_test[y_test == 1])*100:.2f}%)")
    
    return vectorizer, calibrated_svm

def save_models(vectorizer, model):
    """Save the trained vectorizer and model to the models directory."""
    print("\nSaving trained models...")
    
    vectorizer_path = os.path.join(MODELS_DIR, "tfidf_vectorizer.joblib")
    model_path = os.path.join(MODELS_DIR, "svm.joblib")
    
    # Save with joblib
    joblib.dump(vectorizer, vectorizer_path)
    joblib.dump(model, model_path)
    
    print(f"Vectorizer saved to: {vectorizer_path}")
    print(f"Model saved to: {model_path}")
    print("Models saved successfully!")

def main():
    """Main training workflow."""
    print("="*60)
    print("ScamShield NLP - Model Retraining Script")
    print("Using Python 3.10+ and scikit-learn 1.4")
    print("="*60)
    
    # Load and prepare datasets
    df = load_and_prepare_datasets()
    
    # Train and evaluate the model
    vectorizer, model = train_and_evaluate_model(df)
    
    # Save the models
    save_models(vectorizer, model)
    
    print("\n" + "="*60)
    print("Training process completed successfully!")
    print("The new models are now ready for use with the Gradio app.")
    print("="*60)

if __name__ == "__main__":
    main()