import os
import pandas as pd
import kagglehub
from kagglehub import KaggleDatasetAdapter
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import re

# --- Configurations ---
DATASET_NAME = "arslandoula/ah-and-aitd-arslans-human-and-ai-text-database"
FILE_PATH = "Dataset.xlsx"
RAW_DATA_PATH = os.path.join(os.path.dirname(__file__), "data/raw_kaggle_dataset.csv")
CLEAN_DATA_PATH = os.path.join(os.path.dirname(__file__), "data/clean_dataset.csv")
FORCE_DOWNLOAD = False  # 🔁 Set True if you want to re-download even if file exists


# --- Step 1: Smart Dataset Loader ---
def load_or_download_dataset():
    os.makedirs(os.path.dirname(RAW_DATA_PATH), exist_ok=True)

    if os.path.exists(RAW_DATA_PATH) and not FORCE_DOWNLOAD:
        print(f"✅ Found existing dataset at {RAW_DATA_PATH}")
        df = pd.read_csv(RAW_DATA_PATH)
    else:
        print(f"⬇️ Downloading dataset from Kaggle: {DATASET_NAME}")
        df = kagglehub.dataset_load(
            KaggleDatasetAdapter.PANDAS,
            DATASET_NAME,
            FILE_PATH,
        )
        df.to_csv(RAW_DATA_PATH, index=False)
        print(f"📁 Dataset saved to {RAW_DATA_PATH}")
    return df


def clean_text(text):
    text = str(text)
    # Preservation of structural signatures (punctuation, numbers, spacing)
    # which are critical for AI vs Human detection.
    text = re.sub(r"http\S+|www\S+", "", text)       # remove URLs
    text = re.sub(r"\s+", " ", text).strip()         # collapse extra spaces
    return text


# --- Step 3: Preprocess Pipeline ---
def preprocess_dataset(df):
    print("🧹 Cleaning and preprocessing dataset...")
    df = df.rename(columns={
        "text": "content",
        "label_name": "label"
    })

    # Apply cleaning
    df["content"] = df["content"].apply(clean_text)

    # Encode labels (Human-written = 0, AI-generated = 1)
    label_encoder = LabelEncoder()
    df["label_id"] = label_encoder.fit_transform(df["label"])

    # Drop unnecessary columns
    df = df[["content", "label", "label_id"]]

    print(f"✅ Preprocessed {len(df)} samples.")
    df.to_csv(CLEAN_DATA_PATH, index=False)
    print(f"📁 Clean dataset saved to {CLEAN_DATA_PATH}")

    return df


# --- Step 4: Split dataset (optional) ---
def split_dataset(df):
    train, test = train_test_split(df, test_size=0.2, random_state=42, stratify=df["label_id"])
    print(f"📊 Training samples: {len(train)} | Testing samples: {len(test)}")
    return train, test


# --- Main Execution ---
if __name__ == "__main__":
    df_raw = load_or_download_dataset()
    df_clean = preprocess_dataset(df_raw)
    train_df, test_df = split_dataset(df_clean)