"""
prepare_data.py
---------------
Đọc ViClickbait-2025 CSV, tạo cột text kết hợp title + lead_paragraph,
encode label thành số, chia stratified 80/10/10 và lưu ra 3 file CSV.

Dùng:
    python scripts/prepare_data.py \
        --csv /tmp/dataset_nlp/clickbait_dataset_vietnamese.csv \
        --out_dir data/splits
"""

import argparse
import os

import pandas as pd
from sklearn.model_selection import train_test_split


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--csv",
        default="/tmp/dataset_nlp/clickbait_dataset_vietnamese.csv",
        help="Đường dẫn tới file CSV gốc",
    )
    parser.add_argument(
        "--out_dir",
        default="data/splits",
        help="Thư mục lưu train/val/test CSV",
    )
    parser.add_argument(
        "--text_mode",
        choices=["title", "title_lead", "lead"],
        default="title_lead",
        help="Trường văn bản: chỉ title, title+lead_paragraph, hoặc chỉ lead",
    )
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    # ── 1. Đọc dữ liệu ────────────────────────────────────────────────────────
    df = pd.read_csv(args.csv)
    print(f"Tổng số mẫu: {len(df)}")
    print("Phân phối nhãn gốc:")
    print(df["label"].value_counts(), "\n")

    # ── 2. Điền giá trị thiếu ─────────────────────────────────────────────────
    df["title"] = df["title"].fillna("")
    df["lead_paragraph"] = df["lead_paragraph"].fillna("")

    # ── 3. Tạo cột text ───────────────────────────────────────────────────────
    if args.text_mode == "title":
        df["text"] = df["title"]
    elif args.text_mode == "lead":
        df["text"] = df["lead_paragraph"]
    else:  # title_lead (mặc định — cung cấp nhiều ngữ cảnh nhất)
        df["text"] = df["title"] + " [SEP] " + df["lead_paragraph"]

    # ── 4. Encode nhãn ────────────────────────────────────────────────────────
    label_map = {"non-clickbait": 0, "clickbait": 1}
    df["label_id"] = df["label"].map(label_map)
    assert df["label_id"].notna().all(), "Có nhãn không xác định!"
    df["label_id"] = df["label_id"].astype(int)

    print("Label map:", label_map)
    print("Phân phối sau encode:")
    print(df["label_id"].value_counts(), "\n")

    # ── 5. Chỉ giữ cột cần thiết ─────────────────────────────────────────────
    df = df[["id", "text", "label_id"]].copy()

    # ── 6. Stratified split 80 / 10 / 10 ─────────────────────────────────────
    train_df, temp_df = train_test_split(
        df,
        test_size=0.20,
        random_state=args.seed,
        stratify=df["label_id"],
    )
    val_df, test_df = train_test_split(
        temp_df,
        test_size=0.50,
        random_state=args.seed,
        stratify=temp_df["label_id"],
    )

    print(f"Train : {len(train_df):>5} mẫu  | label dist: {dict(train_df['label_id'].value_counts().sort_index())}")
    print(f"Val   : {len(val_df):>5} mẫu  | label dist: {dict(val_df['label_id'].value_counts().sort_index())}")
    print(f"Test  : {len(test_df):>5} mẫu  | label dist: {dict(test_df['label_id'].value_counts().sort_index())}")

    # ── 7. Lưu ra file ────────────────────────────────────────────────────────
    os.makedirs(args.out_dir, exist_ok=True)
    train_df.to_csv(os.path.join(args.out_dir, "train.csv"), index=False)
    val_df.to_csv(os.path.join(args.out_dir, "val.csv"), index=False)
    test_df.to_csv(os.path.join(args.out_dir, "test.csv"), index=False)
    print(f"\nĐã lưu vào: {args.out_dir}/  (train.csv, val.csv, test.csv)")


if __name__ == "__main__":
    main()