"""Load the CSVs, clean up the text, and pair each tweet with its target."""
import re

import pandas as pd
import torch
from torch.utils.data import Dataset

LABEL2ID = {"Against": 0, "Favor": 1, "None": 2}
ID2LABEL = {v: k for k, v in LABEL2ID.items()}

_DIACRITICS = re.compile(r"[ؐ-ًؚ-ٰٟـ]")
_NON_ARABIC = re.compile(r"[^؀-ۿ0-9\s]+")
_URL = re.compile(r"https?://\S+|www\.\S+")
_MENTION = re.compile(r"@\w+")
_MULTI_SPACE = re.compile(r"\s+")
_REPEAT = re.compile(r"(.)\1{2,}")

# Short target descriptions used when the target side is expanded.
TARGET_DESCRIPTIONS = {
    "Covid Vaccine": "لقاح فيروس كورونا كوفيد 19",
    "Digital Transformation": "التحول الرقمي في الخدمات والمجتمع",
    "Women empowerment": "تمكين المرأة وحقوقها في المجتمع",
    "Women Driving": "قيادة المرأة للسيارة",
    "Ecars": "السيارات الكهربائية",
    "Trimester": "نظام الفصول الدراسية الثلاثة في العام الدراسي",
}

# More informative, neutral glosses that spell out what the target is and
# the axis of debate around it (used with --rich_desc).
TARGET_DESCRIPTIONS_RICH = {
    "Covid Vaccine": "لقاح فيروس كورونا كوفيد 19 وأخذه للوقاية من المرض",
    "Digital Transformation": (
        "التحول الرقمي: رقمنة الخدمات والمعاملات الحكومية والمجتمعية"
    ),
    "Women empowerment": "تمكين المرأة وتوسيع دورها وحقوقها في المجتمع",
    "Women Driving": "قيادة المرأة للسيارة",
    "Ecars": "السيارات الكهربائية التي تعمل بالبطارية بدلاً من الوقود",
    "Trimester": (
        "نظام الترمات: تقسيم العام الدراسي إلى ثلاثة فصول بدلاً من فصلين"
    ),
}


def preprocess(text, mode="preserve"):
    """Normalize tweet text.

    mode='strip'     removes non-Arabic characters (emoji, latin).
    mode='preserve'  keeps emoji/latin, drops only urls, mentions,
                     diacritics, repeated characters, and hashtag marks.
    """
    text = str(text)
    text = _URL.sub(" ", text)
    text = _MENTION.sub(" ", text)
    text = _DIACRITICS.sub("", text)
    text = _REPEAT.sub(r"\1\1", text)
    if mode == "strip":
        text = _NON_ARABIC.sub(" ", text)
    else:
        text = text.replace("#", " ").replace("_", " ")
    return _MULTI_SPACE.sub(" ", text).strip()


def target_side(target, use_description=False):
    target = target.strip()
    if use_description:
        return TARGET_DESCRIPTIONS.get(target, target)
    return target


def load_split(csv_path, prep_mode="preserve", has_labels=True):
    """Load a CSV into a dataframe with text, target and optional stance."""
    df = pd.read_csv(csv_path, keep_default_na=False, encoding="utf-8-sig")
    df.columns = df.columns.astype(str).str.strip()
    if "text" not in df.columns and "tweet_text" in df.columns:
        df = df.rename(columns={"tweet_text": "text"})
    df["text"] = df["text"].astype(str).str.strip()
    df["target"] = df["target"].astype(str).str.strip()
    if has_labels:
        df["stance"] = df["stance"].astype(str).str.strip()
        df = df[
            (df["text"] != "")
            & (df["target"] != "")
            & (df["stance"] != "")
        ].copy()
        bad = sorted(set(df["stance"]) - set(LABEL2ID))
        if bad:
            raise ValueError(f"Unknown stance labels in {csv_path}: {bad}")
        df["label"] = df["stance"].map(LABEL2ID).astype(int)
    df["text_clean"] = df["text"].apply(lambda t: preprocess(t, prep_mode))
    return df.reset_index(drop=True)


class StanceDataset(Dataset):
    def __init__(self, df, tokenizer, max_len=128, use_description=False,
                 has_labels=True):
        self.df = df.reset_index(drop=True)
        self.tok = tokenizer
        self.max_len = max_len
        self.use_description = use_description
        self.has_labels = has_labels

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        enc = self.tok(
            target_side(row["target"], self.use_description),
            row["text_clean"],
            truncation=True,
            padding="max_length",
            max_length=self.max_len,
            return_tensors="pt",
        )
        item = {k: v.squeeze(0) for k, v in enc.items()}
        if self.has_labels:
            item["labels"] = torch.tensor(int(row["label"]), dtype=torch.long)
        return item