import os

import joblib
import pandas as pd

from .config import Config, ensure_dirs
from .io_utils import read_any
from .preprocess_bd_structured import build_profile_matrix, preprocess_bd
from .preprocess_text import build_text_splits, preprocess_bangla, preprocess_english


def main():
    cfg = Config()
    ensure_dirs(cfg)

    print("\n================= LOADING DATASETS =================")
    bd_df = read_any(cfg.bd_file)
    en_df = read_any(cfg.en_file)
    bn_df = read_any(cfg.bn_file)

    print("\n================= TEXT PREPROCESS =================")
    en_clean = preprocess_english(en_df)
    bn_clean = preprocess_bangla(bn_df)

    text_all = pd.concat([en_clean, bn_clean], ignore_index=True)
    print("Combined text shape:", text_all.shape)
    print(text_all["lang"].value_counts())
    print(text_all["label"].value_counts())

    train_df, val_df, test_df = build_text_splits(
        text_all, test_size=cfg.test_size, val_size=cfg.val_size, seed=cfg.random_seed
    )
    print("Train/Val/Test:", train_df.shape, val_df.shape, test_df.shape)

    print("\n================= BD STRUCTURED PREPROCESS =================")
    bd_clean = preprocess_bd(bd_df, drop_cols=cfg.drop_cols, defaults=cfg.defaults)
    print("BD structured cleaned shape:", bd_clean.shape)

    print("\n================= PROFILE MATRIX (X_profile) =================")
    X_profile, preprocessor, num_cols, cat_cols = build_profile_matrix(bd_clean)
    print("Numeric cols:", num_cols)
    print("Categorical cols:", cat_cols)
    print("X_profile shape:", X_profile.shape)

    print("\n================= SAVING OUTPUTS =================")
    # CSV outputs (same as colab)
    text_all_path = os.path.join(cfg.out_dir, "text_all_clean.csv")
    train_path = os.path.join(cfg.out_dir, "text_train.csv")
    val_path = os.path.join(cfg.out_dir, "text_val.csv")
    test_path = os.path.join(cfg.out_dir, "text_test.csv")
    bd_path = os.path.join(cfg.out_dir, "bd_suicide_clean_structured.csv")

    text_all.to_csv(text_all_path, index=False, encoding="utf-8")
    train_df.to_csv(train_path, index=False, encoding="utf-8")
    val_df.to_csv(val_path, index=False, encoding="utf-8")
    test_df.to_csv(test_path, index=False, encoding="utf-8")
    bd_clean.to_csv(bd_path, index=False, encoding="utf-8")

    # Artifacts for later reuse (recommended)
    joblib.dump(
        preprocessor, os.path.join(cfg.artifact_dir, "profile_preprocessor.joblib")
    )
    joblib.dump(X_profile, os.path.join(cfg.artifact_dir, "X_profile.joblib"))

    print("✅ Saved:")
    print("-", text_all_path)
    print("-", train_path)
    print("-", val_path)
    print("-", test_path)
    print("-", bd_path)
    print("-", os.path.join(cfg.artifact_dir, "profile_preprocessor.joblib"))
    print("-", os.path.join(cfg.artifact_dir, "X_profile.joblib"))


if __name__ == "__main__":
    main()