import os
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import classification_report, accuracy_score
import joblib

from src.preprocess import clean_text

def train_and_evaluate(data_path, text_col='text', label_col='label', models_dir='models'):
    """
    Trains models for Hate Speech Detection and saves the best one.
    Assumes data_path points to a CSV with a text column and a label column.
    """
    print(f"Loading data from {data_path}...")
    df = pd.read_csv(data_path)
    
    # Preprocessing
    print("Cleaning text data... this may take a moment.")
    df['clean_text'] = df[text_col].apply(clean_text)
    
    X = df['clean_text']
    y = df[label_col]
    
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # Feature extraction using TF-IDF
    print("Vectorizing text using TF-IDF...")
    vectorizer = TfidfVectorizer(max_features=5000, ngram_range=(1, 2))
    X_train_vec = vectorizer.fit_transform(X_train)
    X_test_vec = vectorizer.transform(X_test)
    
    # Train Logistic Regression
    print("Training Logistic Regression...")
    lr_model = LogisticRegression(max_iter=1000)
    lr_model.fit(X_train_vec, y_train)
    lr_preds = lr_model.predict(X_test_vec)
    lr_acc = accuracy_score(y_test, lr_preds)
    print(f"Logistic Regression Accuracy: {lr_acc:.4f}")
    print(classification_report(y_test, lr_preds))
    
    # Train Naive Bayes
    print("Training Naive Bayes...")
    nb_model = MultinomialNB()
    nb_model.fit(X_train_vec, y_train)
    nb_preds = nb_model.predict(X_test_vec)
    nb_acc = accuracy_score(y_test, nb_preds)
    print(f"Naive Bayes Accuracy: {nb_acc:.4f}")
    print(classification_report(y_test, nb_preds))
    
    # Save the best model
    os.makedirs(models_dir, exist_ok=True)
    
    # In this case, we prefer Logistic Regression for saving as it usually performs well and is robust
    best_model = lr_model if lr_acc >= nb_acc else nb_model
    model_name = 'logistic_regression.pkl' if lr_acc >= nb_acc else 'naive_bayes.pkl'
    
    joblib.dump(best_model, os.path.join(models_dir, 'model.pkl'))
    joblib.dump(vectorizer, os.path.join(models_dir, 'vectorizer.pkl'))
    
    print(f"Best model ({model_name}) and vectorizer saved to '{models_dir}/'.")

if __name__ == "__main__":
    # Example usage (will need dummy data to run directly)
    default_data_path = os.path.join('data', 'sample.csv')
    if os.path.exists(default_data_path):
        train_and_evaluate(default_data_path)
    else:
        print(f"Data not found at {default_data_path}. Please create a sample dataset first.")