"""
download_datasets.py
====================
Downloads all public datasets from HuggingFace Hub and saves them
as CSV files in data/raw/ for downstream processing.

Datasets:
  - Civil Comments (google/civil_comments)
  - HateXplain (Hate-speech-CNERG/hatexplain)
  - OLID (christophsonntag/OLID)
  - Jigsaw Toxic Comment Classification

Usage:
    python preprocessing/download_datasets.py
"""

import json
import os
import sys
import logging
import urllib.request
from collections import Counter
from pathlib import Path

import pandas as pd
from datasets import load_dataset
import threading
import multiprocess.resource_tracker as rt

# Silence the multiprocess cleanup exception during interpreter shutdown
def _noop_del(self):
    pass

rt.ResourceTracker.__del__ = _noop_del
# Patch Python background thread cleanup bug in `multiprocess`
if not hasattr(threading.RLock, "_recursion_count"):
    threading.RLock._recursion_count = lambda self: 0
# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
os.environ["HF_DATASETS_NUM_PROC"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Add this line:
os.environ["MULTIPROCESS_NO_RESOURCE_TRACKER"] = "1"
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger(__name__)

PROJECT_ROOT = Path(__file__).resolve().parent.parent
RAW_DIR = PROJECT_ROOT / "data" / "raw"
RAW_DIR.mkdir(parents=True, exist_ok=True)


# ---------------------------------------------------------------------------
# Download helpers
# ---------------------------------------------------------------------------

def download_civil_comments():
    """
    Civil Comments — ~1.8M comments with continuous toxicity scores.
    We sample a manageable subset (100k) to keep training feasible.
    """
    output_path = RAW_DIR / "civil_comments.csv"
    if output_path.exists():
        logger.info("civil_comments.csv already exists, skipping download.")
        return

    logger.info("Downloading Civil Comments dataset...")
    ds = load_dataset("google/civil_comments", split="train")
    df = ds.to_pandas()

    # Keep relevant columns
    keep_cols = [
        "text", "toxicity", "severe_toxicity",
        "obscene", "threat", "insult", "identity_attack", "sexual_explicit",
    ]
    df = df[[c for c in keep_cols if c in df.columns]]

    # Sample 100k to keep dataset manageable (stratified by toxic/non-toxic)
    df["is_toxic"] = (df["toxicity"] >= 0.5).astype(int)
    toxic = df[df["is_toxic"] == 1]
    safe = df[df["is_toxic"] == 0].sample(n=min(70000, len(df[df["is_toxic"] == 0])), random_state=42)
    df = pd.concat([toxic.head(30000), safe], ignore_index=True)
    df = df.drop(columns=["is_toxic"])

    df.to_csv(output_path, index=False)
    logger.info(f"Saved {len(df)} rows to {output_path}")


def download_hatexplain():
    """
    HateXplain — ~20K posts with hatespeech/offensive/normal labels.
    Downloads raw JSON dataset directly from GitHub.
    """
    output_path = RAW_DIR / "hatexplain.csv"
    if output_path.exists():
        logger.info("hatexplain.csv already exists, skipping download.")
        return

    logger.info("Downloading HateXplain dataset from source...")
    url = "https://raw.githubusercontent.com/punyajoy/HateXplain/master/Data/dataset.json"
    try:
        with urllib.request.urlopen(url, timeout=120) as response:
            data = json.loads(response.read().decode())
    except Exception as e:
        logger.error(f"  Failed to download HateXplain: {e}")
        logger.warning("  Skipping HateXplain — continuing with other datasets.")
        return

    rows = []
    for pid, entry in data.items():
        tokens = entry.get("post_tokens", [])
        text = " ".join(tokens) if isinstance(tokens, list) else str(tokens)
        if not text.strip():
            continue

        # Annotators: list of dicts with "label" key
        # Label values: "hatespeech", "normal", "offensive"
        ann_labels = []
        annotators = entry.get("annotators", [])
        if isinstance(annotators, list):
            ann_labels = [a.get("label", "normal") for a in annotators if isinstance(a, dict)]
        elif isinstance(annotators, dict):
            # Some versions store as {"label": [...], "annotator_id": [...]}
            ann_labels = annotators.get("label", [])

        if not ann_labels:
            continue

        majority_label = Counter(ann_labels).most_common(1)[0][0]
        # Normalise to string
        if isinstance(majority_label, int):
            majority_label = {0: "hatespeech", 1: "normal", 2: "offensive"}.get(majority_label, "normal")
        rows.append({"text": text, "label": str(majority_label).lower()})

    df = pd.DataFrame(rows)
    df.to_csv(output_path, index=False)
    logger.info(f"Saved {len(df)} rows to {output_path}")


def download_olid():
    """
    OLID — 14K tweets with offensive language identification.
    Subtask A: OFF/NOT
    """
    output_path = RAW_DIR / "olid.csv"
    if output_path.exists():
        logger.info("olid.csv already exists, skipping download.")
        return

    logger.info("Downloading OLID dataset...")
    ds = load_dataset("christophsonntag/OLID", split="train")
    df = ds.to_pandas()

    # Rename columns to standard format
    if "tweet" in df.columns:
        df = df.rename(columns={"tweet": "text"})
    if "subtask_a" in df.columns:
        df = df.rename(columns={"subtask_a": "label_a"})
    if "subtask_b" in df.columns:
        df = df.rename(columns={"subtask_b": "label_b"})
    if "subtask_c" in df.columns:
        df = df.rename(columns={"subtask_c": "label_c"})

    df.to_csv(output_path, index=False)
    logger.info(f"Saved {len(df)} rows to {output_path}")


def download_jigsaw_toxic():
    """
    Jigsaw Toxic Comment Classification — ~160K comments with
    multi-label toxicity annotations.
    """
    output_path = RAW_DIR / "jigsaw_toxic.csv"
    if output_path.exists():
        logger.info("jigsaw_toxic.csv already exists, skipping download.")
        return

    logger.info("Downloading Jigsaw Toxic Comment dataset...")
    try:
        ds = load_dataset(
            "thesofakillers/jigsaw-toxic-comment-classification-challenge",
            split="train",
        )
        df = ds.to_pandas()

        keep_cols = [
            "comment_text", "toxic", "severe_toxic",
            "obscene", "threat", "insult", "identity_hate",
        ]
        df = df[[c for c in keep_cols if c in df.columns]]
        df = df.rename(columns={"comment_text": "text"})

        df.to_csv(output_path, index=False)
        logger.info(f"Saved {len(df)} rows to {output_path}")
    except Exception as e:
        logger.warning(f"Could not download Jigsaw dataset: {e}")
        logger.warning("You may need to download it manually from Kaggle.")
        logger.warning("Place the train.csv in data/raw/jigsaw_toxic.csv")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    logger.info("=" * 60)
    logger.info("Starting dataset downloads...")
    logger.info("=" * 60)

    download_civil_comments()
    download_hatexplain()
    download_olid()
    download_jigsaw_toxic()

    logger.info("=" * 60)
    logger.info("All downloads complete!")
    logger.info(f"Raw data saved to: {RAW_DIR}")
    logger.info("=" * 60)

    # Show summary
    for f in RAW_DIR.glob("*.csv"):
        df = pd.read_csv(f, nrows=0)
        row_count = sum(1 for _ in open(f, encoding="utf-8")) - 1
        logger.info(f"  {f.name}: {row_count:,} rows, columns: {list(df.columns)}")


if __name__ == "__main__":
    try:
        main()
    finally:
        # Force resource_tracker cleanup before garbage collection tears down threading objects
        try:
            import multiprocess.resource_tracker as rt
            rt.get_tracker()._stop()
        except Exception:
            pass