import bz2
import codecs
import os
import re
import subprocess
import urllib.parse as up
import xml.etree.cElementTree as etree

import mwparserfromhell
from datasets import Dataset

from data_juicer.download.downloader import (DocumentDownloader,
                                             DocumentExtractor,
                                             DocumentIterator,
                                             download_and_extract,
                                             get_wikipedia_urls)
from data_juicer.utils.file_utils import expand_outdir_and_mkdir

# The majority of this code is taken from the HuggingFace
# implementation of the Wikipedia dataset preparation:
# https://github.com/huggingface/datasets/blob/7e30308f49f8c85dc7a2ab5aafbff04b5d2f38e2/datasets/wikipedia/wikipedia.py

MEDIA_ALIASES = {
    'ab': ['Медиа', 'Файл', 'Афаил', 'Амедиа', 'Изображение'],
    'ace': ['Beureukaih', 'Gambar', 'Alat', 'Berkas'],
    'ady': ['Медиа'],
    'af': ['Lêer', 'Beeld'],
    'als': ['Medium', 'Datei', 'Bild'],
    'am': ['ፋይል', 'ስዕል'],
    'an': ['Imachen', 'Imagen'],
    'ang': ['Ymele', 'Biliþ'],
    'ar': ['ميديا', 'صورة', 'وسائط', 'ملف'],
    'arc': ['ܠܦܦܐ', 'ܡܝܕܝܐ'],
    'arz': ['ميديا', 'صورة', 'وسائط', 'ملف'],
    'as': ['চিত্ৰ', 'चित्र', 'চিত্র', 'মাধ্যম'],
    'ast': ['Imaxen', 'Ficheru', 'Imaxe', 'Archivu', 'Imagen', 'Medios'],
    'atj': ['Tipatcimoctakewin', 'Natisinahikaniwoc'],
    'av': ['Медиа', 'Файл', 'Изображение'],
    'ay': ['Medio', 'Archivo', 'Imagen'],
    'az': ['Mediya', 'Şəkil', 'Fayl'],
    'azb': ['رسانه', 'تصویر', 'مدیا', 'فایل', 'رسانه‌ای'],
    'ba': ['Медиа', 'Рәсем', 'Файл', 'Изображение'],
    'bar': ['Medium', 'Datei', 'Bild'],
    'bat-smg': ['Vaizdas', 'Medėjė', 'Abruozdielis'],
    'bcl': ['Medio', 'Ladawan'],
    'be': ['Мультымедыя', 'Файл', 'Выява'],
    'be-x-old': ['Мэдыя', 'Файл', 'Выява'],
    'bg': ['Медия', 'Файл', 'Картинка'],
    'bh': ['मीडिया', 'चित्र'],
    'bjn': ['Barakas', 'Gambar', 'Berkas'],
    'bm': ['Média', 'Fichier'],
    'bn': ['চিত্র', 'মিডিয়া'],
    'bpy': ['ছবি', 'মিডিয়া'],
    'br': ['Skeudenn', 'Restr'],
    'bs': ['Mediji', 'Slika', 'Datoteka', 'Medija'],
    'bug': ['Gambar', 'Berkas'],
    'bxr': ['Файл', 'Меди', 'Изображение'],
    'ca': ['Fitxer', 'Imatge'],
    'cbk-zam': ['Medio', 'Archivo', 'Imagen'],
    'cdo': ['文件', '媒體', '圖像', '檔案'],
    'ce': ['Хlум', 'Медиа', 'Сурт', 'Файл', 'Медйа', 'Изображение'],
    'ceb': ['Payl', 'Medya', 'Imahen'],
    'ch': ['Litratu'],
    'ckb': ['میدیا', 'پەڕگە'],
    'co': ['Immagine'],
    'crh': ['Медиа', 'Resim', 'Файл', 'Fayl', 'Ресим'],
    'cs': ['Soubor', 'Média', 'Obrázok'],
    'csb': ['Òbrôzk', 'Grafika'],
    'cu': ['Видъ', 'Ви́дъ', 'Дѣло', 'Срѣдьства'],
    'cv': ['Медиа', 'Ӳкерчĕк', 'Изображение'],
    'cy': ['Delwedd'],
    'da': ['Billede', 'Fil'],
    'de': ['Medium', 'Datei', 'Bild'],
    'din': ['Ciɛl', 'Apamduööt'],
    'diq': ['Medya', 'Dosya'],
    'dsb': ['Wobraz', 'Dataja', 'Bild', 'Medija'],
    'dty': ['चित्र', 'मिडिया'],
    'dv': ['ފައިލު', 'މީޑިއާ', 'ފައިލް'],
    'el': ['Εικόνα', 'Αρχείο', 'Μέσο', 'Μέσον'],
    'eml': ['Immagine'],
    'eo': ['Dosiero', 'Aŭdvidaĵo'],
    'es': ['Medio', 'Archivo', 'Imagen'],
    'et': ['Pilt', 'Fail', 'Meedia'],
    'eu': ['Irudi', 'Fitxategi'],
    'ext': ['Archivu', 'Imagen', 'Mediu'],
    'fa': ['رسانه', 'تصویر', 'مدیا', 'پرونده', 'رسانه‌ای'],
    'ff': ['Média', 'Fichier'],
    'fi': ['Kuva', 'Tiedosto'],
    'fiu-vro': ['Pilt', 'Meediä'],
    'fo': ['Miðil', 'Mynd'],
    'fr': ['Média', 'Fichier'],
    'frp': ['Émâge', 'Fichiér', 'Mèdia'],
    'frr': ['Medium', 'Datei', 'Bild'],
    'fur': ['Immagine', 'Figure'],
    'fy': ['Ofbyld'],
    'ga': ['Íomhá', 'Meán'],
    'gag': ['Mediya', 'Medya', 'Resim', 'Dosya', 'Dosye'],
    'gan': ['媒体文件', '文件', '文檔', '档案', '媒體', '图像', '圖像', '媒体', '檔案'],
    'gd': ['Faidhle', 'Meadhan'],
    'gl': ['Imaxe', 'Ficheiro', 'Arquivo', 'Imagem'],
    'glk': ['رسانه', 'تصویر', 'پرونده', 'فاىل', 'رسانه‌ای', 'مديا'],
    'gn': ['Medio', 'Imagen', "Ta'ãnga"],
    'gom': ['माध्यम', 'मिडिया', 'फायल'],
    'gor': ['Gambar', 'Berkas'],
    'got': ['𐍆𐌴𐌹𐌻𐌰'],
    'gu': ['દ્રશ્ય-શ્રાવ્ય (મિડિયા)', 'દ્રશ્ય-શ્રાવ્ય_(મિડિયા)', 'ચિત્ર'],
    'gv': ['Coadan', 'Meanyn'],
    'hak': ['文件', '媒體', '圖像', '檔案'],
    'haw': ['Kiʻi', 'Waihona', 'Pāpaho'],
    'he': ['תמונה', 'קו', 'מדיה', 'קובץ'],
    'hi': ['मीडिया', 'चित्र'],
    'hif': ['file', 'saadhan'],
    'hr': ['Mediji', 'DT', 'Slika', 'F', 'Datoteka'],
    'hsb': ['Wobraz', 'Dataja', 'Bild'],
    'ht': ['Imaj', 'Fichye', 'Medya'],
    'hu': ['Kép', 'Fájl', 'Média'],
    'hy': ['Պատկեր', 'Մեդիա'],
    'ia': ['Imagine', 'Multimedia'],
    'id': ['Gambar', 'Berkas'],
    'ig': ['Nká', 'Midia', 'Usòrò', 'Ákwúkwó orünotu', 'Ákwúkwó_orünotu'],
    'ii': ['媒体文件', '文件', '档案', '图像', '媒体'],
    'ilo': ['Midia', 'Papeles'],
    'inh': ['Медиа', 'Файл', 'Изображение'],
    'io': ['Imajo', 'Arkivo'],
    'is': ['Miðill', 'Mynd'],
    'it': ['Immagine'],
    'ja': ['メディア', 'ファイル', '画像'],
    'jbo': ['velsku', 'datnyvei'],
    'jv': ['Barkas', 'Medhia', 'Gambar', 'Médhia'],
    'ka': ['მედია', 'სურათი', 'ფაილი'],
    'kaa': ['Swret', 'Таспа', 'سۋرەت', 'Taspa', "Su'wret", 'Сурет', 'تاسپا'],
    'kab': ['Tugna'],
    'kbd': ['Медиа', 'Файл'],
    'kbp': ['Média', 'Fichier'],
    'kg': ['Fisye'],
    'kk': ['Swret', 'سۋرەت', 'Таспа', 'Taspa', 'Сурет', 'تاسپا'],
    'kl': ['Billede', 'Fiileq', 'Fil'],
    'km': ['ឯកសារ', 'រូបភាព', 'មេឌា', 'មីឌា'],
    'kn': ['ಚಿತ್ರ', 'ಮೀಡಿಯ'],
    'ko': ['미디어', '파일', '그림'],
    'koi': ['Медиа', 'Файл', 'Изображение'],
    'krc': ['Медиа', 'Файл', 'Изображение'],
    'ks': ['میڈیا', 'فَیِل'],
    'ksh': [
        'Beld', 'Meedije', 'Medie', 'Belld', 'Medium', 'Datei', 'Meedijum',
        'Bild'
    ],
    'ku': ['میدیا', 'پەڕگە', 'Medya', 'Wêne'],
    'kv': ['Медиа', 'Файл', 'Изображение'],
    'kw': ['Restren'],
    'ky': ['Медиа', 'Файл'],
    'la': ['Imago', 'Fasciculus'],
    'lad': ['Dossia', 'Medya', 'Archivo', 'Dosya', 'Imagen', 'Meddia'],
    'lb': ['Fichier', 'Bild'],
    'lbe': ['Медиа', 'Сурат', 'Изображение'],
    'lez': ['Медиа', 'Mediya', 'Файл', 'Şəkil', 'Изображение'],
    'lfn': ['Fix'],
    'li': ['Afbeelding', 'Plaetje', 'Aafbeilding'],
    'lij': ['Immaggine', 'Immagine'],
    'lmo': ['Immagine', 'Imàjine', 'Archivi'],
    'ln': ['Média', 'Fichier'],
    'lo': ['ສື່ອ', 'ສື່', 'ຮູບ'],
    'lrc': ['رسانه', 'تصویر', 'رسانه‌ای', 'جانیا', 'أسگ', 'ڤارئسگأر'],
    'lt': ['Vaizdas', 'Medija'],
    'ltg': ['Medeja', 'Fails'],
    'lv': ['Attēls'],
    'mai': ['मेडिया', 'फाइल'],
    'map-bms': ['Barkas', 'Medhia', 'Gambar', 'Médhia'],
    'mdf': ['Медиа', 'Няйф', 'Изображение'],
    'mg': ['Rakitra', 'Sary', 'Média'],
    'mhr': ['Медиа', 'Файл', 'Изображение'],
    'min': ['Gambar', 'Berkas'],
    'mk': ['Податотека', 'Медија', 'Медиум', 'Слика'],
    'ml': ['പ്രമാണം', 'ചി', 'മീഡിയ', 'പ്ര', 'ചിത്രം'],
    'mn': ['Медиа', 'Файл', 'Зураг'],
    'mr': ['चित्र', 'मिडिया'],
    'mrj': ['Медиа', 'Файл', 'Изображение'],
    'ms': ['Fail', 'Imej'],
    'mt': ['Midja', 'Medja', 'Stampa'],
    'mwl': ['Multimédia', 'Fexeiro', 'Ficheiro', 'Arquivo', 'Imagem'],
    'my': ['ဖိုင်', 'မီဒီယာ'],
    'myv': ['Медия', 'Артовкс', 'Изображение'],
    'mzn': ['رسانه', 'تصویر', 'مه‌دیا', 'مدیا', 'پرونده', 'رسانه‌ای'],
    'nah': ['Mēdiatl', 'Īxiptli', 'Imagen'],
    'nap': ['Fiùra', 'Immagine'],
    'nds': ['Datei', 'Bild'],
    'nds-nl': ['Ofbeelding', 'Afbeelding', 'Bestaand'],
    'ne': ['मीडिया', 'चित्र'],
    'new': ['किपा', 'माध्यम'],
    'nl': ['Bestand', 'Afbeelding'],
    'nn': ['Fil', 'Bilde', 'Filpeikar'],
    'no': ['Fil', 'Medium', 'Bilde'],
    'nov': [],
    'nrm': ['Média', 'Fichier'],
    'nso': ['Seswantšho'],
    'nv': ['Eʼelyaaígíí'],
    'oc': ['Imatge', 'Fichièr', 'Mèdia'],
    'olo': ['Kuva', 'Medii', 'Failu'],
    'or': ['ମାଧ୍ୟମ', 'ଫାଇଲ'],
    'os': ['Ныв', 'Медиа', 'Файл', 'Изображение'],
    'pa': ['ਤਸਵੀਰ', 'ਮੀਡੀਆ'],
    'pcd': ['Média', 'Fichier'],
    'pdc': ['Medium', 'Datei', 'Bild', 'Feil'],
    'pfl': ['Dadai', 'Medium', 'Datei', 'Bild'],
    'pi': ['मीडिया', 'पटिमा'],
    'pl': ['Plik', 'Grafika'],
    'pms': ['Figura', 'Immagine'],
    'pnb': ['میڈیا', 'تصویر', 'فائل'],
    'pnt': ['Εικόνα', 'Αρχείον', 'Εικόναν', 'Μέσον'],
    'ps': ['انځور', 'رسنۍ', 'دوتنه'],
    'pt': ['Multimédia', 'Ficheiro', 'Arquivo', 'Imagem'],
    'qu': ['Midya', 'Imagen', 'Rikcha'],
    'rm': ['Multimedia', 'Datoteca'],
    'rmy': ['Fişier', 'Mediya', 'Chitro', 'Imagine'],
    'ro': ['Fişier', 'Imagine', 'Fișier'],
    'roa-rup': ['Fişier', 'Imagine', 'Fișier'],
    'roa-tara': ['Immagine'],
    'ru': ['Медиа', 'Файл', 'Изображение'],
    'rue': ['Медіа', 'Медиа', 'Файл', 'Изображение', 'Зображення'],
    'rw': ['Dosiye', 'Itangazamakuru'],
    'sa': ['चित्रम्', 'माध्यमम्', 'सञ्चिका', 'माध्यम', 'चित्रं'],
    'sah': ['Миэдьийэ', 'Ойуу', 'Билэ', 'Изображение'],
    'sat': ['ᱨᱮᱫ', 'ᱢᱤᱰᱤᱭᱟ'],
    'sc': ['Immàgini'],
    'scn': ['Immagine', 'Mmàggini', 'Mèdia'],
    'sd': ['عڪس', 'ذريعات', 'فائل'],
    'se': ['Fiila'],
    'sg': ['Média', 'Fichier'],
    'sh': ['Mediji', 'Slika', 'Медија', 'Datoteka', 'Medija', 'Слика'],
    'si': ['රූපය', 'මාධ්‍යය', 'ගොනුව'],
    'sk': ['Súbor', 'Obrázok', 'Médiá'],
    'sl': ['Slika', 'Datoteka'],
    'sq': ['Figura', 'Skeda'],
    'sr': [
        'Датотека',
        'Medij',
        'Slika',
        'Медија',
        'Datoteka',
        'Медиј',
        'Medija',
        'Слика',
    ],
    'srn': ['Afbeelding', 'Gefre'],
    'stq': ['Bielde', 'Bild'],
    'su': ['Média', 'Gambar'],
    'sv': ['Fil', 'Bild'],
    'sw': ['Faili', 'Picha'],
    'szl': ['Plik', 'Grafika'],
    'ta': ['படிமம்', 'ஊடகம்'],
    'tcy': ['ಮಾದ್ಯಮೊ', 'ಫೈಲ್'],
    'te': ['ఫైలు', 'దస్త్రం', 'బొమ్మ', 'మీడియా'],
    'tet': ['Imajen', 'Arquivo', 'Imagem'],
    'tg': ['Акс', 'Медиа'],
    'th': ['ไฟล์', 'สื่อ', 'ภาพ'],
    'ti': ['ፋይል', 'ሜድያ'],
    'tk': ['Faýl'],
    'tl': ['Midya', 'Talaksan'],
    'tpi': ['Fail'],
    'tr': ['Medya', 'Resim', 'Dosya', 'Ortam'],
    'tt': ['Медиа', 'Рәсем', 'Файл', 'Räsem', 'Изображение'],
    'ty': ['Média', 'Fichier'],
    'tyv': ['Медиа', 'Файл', 'Изображение'],
    'udm': ['Медиа', 'Файл', 'Суред', 'Изображение'],
    'ug': ['ۋاسىتە', 'ھۆججەت'],
    'uk': ['Медіа', 'Медиа', 'Файл', 'Изображение', 'Зображення'],
    'ur': ['میڈیا', 'تصویر', 'وسیط', 'زریعہ', 'فائل', 'ملف'],
    'uz': ['Mediya', 'Tasvir', 'Fayl'],
    'vec': ['Immagine', 'Imàjine', 'Mèdia'],
    'vep': ['Pilt', 'Fail'],
    'vi': ['Phương_tiện', 'Tập_tin', 'Hình', 'Tập tin', 'Phương tiện'],
    'vls': ['Afbeelding', 'Ofbeeldienge'],
    'vo': ['Ragiv', 'Magod', 'Nünamakanäd'],
    'wa': ['Imådje'],
    'war': ['Medya', 'Fayl', 'Paypay'],
    'wo': ['Xibaarukaay', 'Dencukaay'],
    'wuu': ['文件', '档案', '图像', '媒体'],
    'xal': ['Аһар', 'Боомг', 'Изображение', 'Зург'],
    'xmf': ['მედია', 'სურათი', 'ფაილი'],
    'yi': ['מעדיע', 'תמונה', 'טעקע', 'בילד'],
    'yo': ['Fáìlì', 'Amóhùnmáwòrán', 'Àwòrán'],
    'za': ['媒体文件', '文件', '档案', '图像', '媒体'],
    'zea': ['Afbeelding', 'Plaetje'],
    'zh': ['媒体文件', 'F', '文件', '媒體', '档案', '图像', '圖像', '媒体', '檔案'],
    'zh-classical': ['文件', '媒體', '圖像', '檔案'],
    'zh-min-nan': ['tóng-àn', '文件', '媒體', 'Mûi-thé', '圖像', '檔案'],
    'zh-yue': [
        '檔',
        '档',
        '文件',
        '图',
        '媒體',
        '圖',
        '档案',
        '图像',
        '圖像',
        '媒体',
        '檔案',
    ],
}

CAT_ALIASES = {
    'ab': ['Категория', 'Акатегориа'],
    'ace': ['Kawan', 'Kategori'],
    'af': ['Kategorie'],
    'ak': ['Nkyekyem'],
    'als': ['Kategorie'],
    'am': ['መደብ'],
    'an': ['Categoría'],
    'ang': ['Flocc'],
    'ar': ['تصنيف'],
    'arc': ['ܣܕܪܐ'],
    'arz': ['تصنيف'],
    'as': ['CAT', 'শ্ৰেণী', 'श्रेणी', 'শ্রেণী'],
    'ast': ['Categoría'],
    'atj': ['Tipanictawin'],
    'av': ['Категория'],
    'ay': ['Categoría'],
    'az': ['Kateqoriya'],
    'azb': ['بؤلمه'],
    'ba': ['Төркөм', 'Категория'],
    'bar': ['Kategorie'],
    'bat-smg': ['Kategorija', 'Kateguorėjė'],
    'bcl': ['Kategorya'],
    'be': ['Катэгорыя'],
    'be-x-old': ['Катэгорыя'],
    'bg': ['Категория'],
    'bh': ['श्रेणी'],
    'bjn': ['Tumbung', 'Kategori'],
    'bm': ['Catégorie'],
    'bn': ['বিষয়শ্রেণী', 'വിഭാഗം'],
    'bpy': ['থাক'],
    'br': ['Rummad'],
    'bs': ['Kategorija'],
    'bug': ['Kategori'],
    'bxr': ['Категори', 'Категория'],
    'ca': ['Categoria'],
    'cbk-zam': ['Categoría'],
    'cdo': ['分類'],
    'ce': ['Категори', 'Тоба', 'Кадегар'],
    'ceb': ['Kategoriya'],
    'ch': ['Katigoria'],
    'ckb': ['پ', 'پۆل'],
    'co': ['Categoria'],
    'crh': ['Категория', 'Kategoriya'],
    'cs': ['Kategorie'],
    'csb': ['Kategòrëjô'],
    'cu': ['Катигорї', 'Категория', 'Катигорїꙗ'],
    'cv': ['Категори'],
    'cy': ['Categori'],
    'da': ['Kategori'],
    'de': ['Kategorie'],
    'din': ['Bekätakthook'],
    'diq': ['Kategoriye', 'Kategori'],
    'dsb': ['Kategorija'],
    'dty': ['श्रेणी'],
    'dv': ['ޤިސްމު'],
    'el': ['Κατηγορία'],
    'eml': ['Categoria'],
    'eo': ['Kategorio'],
    'es': ['CAT', 'Categoría'],
    'et': ['Kategooria'],
    'eu': ['Kategoria'],
    'ext': ['Categoría', 'Categoria'],
    'fa': ['رده'],
    'ff': ['Catégorie'],
    'fi': ['Luokka'],
    'fiu-vro': ['Katõgooria'],
    'fo': ['Bólkur'],
    'fr': ['Catégorie'],
    'frp': ['Catègorie'],
    'frr': ['Kategorie'],
    'fur': ['Categorie'],
    'fy': ['Kategory'],
    'ga': ['Rang', 'Catagóir'],
    'gag': ['Kategori', 'Kategoriya'],
    'gan': ['分類', '分类'],
    'gd': ['Roinn-seòrsa'],
    'gl': ['Categoría'],
    'glk': ['جرگه', 'رده'],
    'gn': ['Ñemohenda'],
    'gom': ['वर्ग', 'श्रेणी'],
    'gor': ['Dalala'],
    'got': ['𐌷𐌰𐌽𐍃𐌰'],
    'gu': ['શ્રેણી', 'CAT', 'શ્રે'],
    'gv': ['Ronney'],
    'hak': ['分類'],
    'haw': ['Māhele'],
    'he': ['קטגוריה', 'קט'],
    'hi': ['श्र', 'श्रेणी'],
    'hif': ['vibhag'],
    'hr': ['CT', 'KT', 'Kategorija'],
    'hsb': ['Kategorija'],
    'ht': ['Kategori'],
    'hu': ['Kategória'],
    'hy': ['Կատեգորիա'],
    'ia': ['Categoria'],
    'id': ['Kategori'],
    'ie': ['Categorie'],
    'ig': ['Ébéonọr', 'Òtù'],
    'ii': ['分类'],
    'ilo': ['Kategoria'],
    'inh': ['ОагӀат'],
    'io': ['Kategorio'],
    'is': ['Flokkur'],
    'it': ['CAT', 'Categoria'],
    'ja': ['カテゴリ'],
    'jbo': ['klesi'],
    'jv': ['Kategori'],
    'ka': ['კატეგორია'],
    'kaa': ['Sanat', 'Kategoriya', 'Санат', 'سانات'],
    'kab': ['Taggayt'],
    'kbd': ['Категория', 'Категориэ'],
    'kbp': ['Catégorie'],
    'kg': ['Kalasi'],
    'kk': ['Sanat', 'Санат', 'سانات'],
    'kl': ['Sumut_atassuseq', 'Kategori', 'Sumut atassuseq'],
    'km': ['ចំនាត់ថ្នាក់ក្រុម', 'ចំណាត់ក្រុម', 'ចំណាត់ថ្នាក់ក្រុម'],
    'kn': ['ವರ್ಗ'],
    'ko': ['분류'],
    'koi': ['Категория'],
    'krc': ['Категория'],
    'ks': ['زٲژ'],
    'ksh': [
        'Saachjropp',
        'Saachjrop',
        'Katejori',
        'Kategorie',
        'Saachjrupp',
        'Kattejori',
        'Sachjrop',
    ],
    'ku': ['Kategorî', 'پۆل'],
    'kv': ['Категория'],
    'kw': ['Class', 'Klass'],
    'ky': ['Категория'],
    'la': ['Categoria'],
    'lad': ['Kateggoría', 'Katēggoría', 'Categoría'],
    'lb': ['Kategorie'],
    'lbe': ['Категория'],
    'lez': ['Категория'],
    'lfn': ['Categoria'],
    'li': ['Categorie', 'Kategorie'],
    'lij': ['Categorîa', 'Categoria'],
    'lmo': ['Categuria', 'Categoria'],
    'ln': ['Catégorie'],
    'lo': ['ໝວດ'],
    'lrc': ['دأسە'],
    'lt': ['Kategorija'],
    'ltg': ['Kategoreja'],
    'lv': ['Kategorija'],
    'mai': ['CA', 'श्रेणी'],
    'map-bms': ['Kategori'],
    'mdf': ['Категорие', 'Категория'],
    'mg': ['Sokajy', 'Catégorie'],
    'mhr': ['Категория', 'Категорий'],
    'min': ['Kategori'],
    'mk': ['Категорија'],
    'ml': ['വിഭാഗം', 'വി', 'വർഗ്ഗം', 'വ'],
    'mn': ['Ангилал'],
    'mr': ['वर्ग'],
    'mrj': ['Категори', 'Категория'],
    'ms': ['Kategori'],
    'mt': ['Kategorija'],
    'mwl': ['Catadorie', 'Categoria'],
    'my': ['ကဏ္ဍ'],
    'myv': ['Категория'],
    'mzn': ['رج', 'رده'],
    'nah': ['Neneuhcāyōtl', 'Categoría'],
    'nap': ['Categurìa', 'Categoria'],
    'nds': ['Kategorie'],
    'nds-nl': ['Categorie', 'Kattegerie', 'Kategorie'],
    'ne': ['श्रेणी'],
    'new': ['पुचः'],
    'nl': ['Categorie'],
    'nn': ['Kategori'],
    'no': ['Kategori'],
    'nrm': ['Catégorie'],
    'nso': ['Setensele'],
    'nv': ['Tʼááłáhági_átʼéego', 'Tʼááłáhági átʼéego'],
    'oc': ['Categoria'],
    'olo': ['Kategourii'],
    'or': ['ବିଭାଗ', 'ଶ୍ରେଣୀ'],
    'os': ['Категори'],
    'pa': ['ਸ਼੍ਰੇਣੀ'],
    'pcd': ['Catégorie'],
    'pdc': ['Abdeeling', 'Kategorie'],
    'pfl': ['Kadegorie', 'Sachgrubb', 'Kategorie'],
    'pi': ['विभाग'],
    'pl': ['Kategoria'],
    'pms': ['Categorìa'],
    'pnb': ['گٹھ'],
    'pnt': ['Κατηγορίαν'],
    'ps': ['وېشنيزه'],
    'pt': ['Categoria'],
    'qu': ['Katiguriya'],
    'rm': ['Categoria'],
    'rmy': ['Shopni'],
    'ro': ['Categorie'],
    'roa-rup': ['Categorie'],
    'roa-tara': ['Categoria'],
    'ru': ['Категория', 'К'],
    'rue': ['Категория', 'Катеґорія'],
    'rw': ['Ikiciro'],
    'sa': ['वर्गः'],
    'sah': ['Категория'],
    'sat': ['ᱛᱷᱚᱠ'],
    'sc': ['Categoria'],
    'scn': ['Catigurìa'],
    'sd': ['زمرو'],
    'se': ['Kategoriija'],
    'sg': ['Catégorie'],
    'sh': ['Kategorija', 'Категорија'],
    'si': ['ප්‍රවර්ගය'],
    'sk': ['Kategória'],
    'sl': ['Kategorija'],
    'sq': ['Kategoria', 'Kategori'],
    'sr': ['Kategorija', 'Категорија'],
    'srn': ['Categorie', 'Guru'],
    'stq': ['Kategorie'],
    'su': ['Kategori'],
    'sv': ['Kategori'],
    'sw': ['Jamii'],
    'szl': ['Kategoryjo', 'Kategoria'],
    'ta': ['பகுப்பு'],
    'tcy': ['ವರ್ಗೊ'],
    'te': ['వర్గం'],
    'tet': ['Kategoría', 'Kategoria'],
    'tg': ['Гурӯҳ'],
    'th': ['หมวดหมู่'],
    'ti': ['መደብ'],
    'tk': ['Kategoriýa'],
    'tl': ['Kategorya', 'Kaurian'],
    'tpi': ['Grup'],
    'tr': ['Kategori', 'KAT'],
    'tt': ['Төркем', 'Törkem', 'Категория'],
    'ty': ['Catégorie'],
    'tyv': ['Аңгылал', 'Категория'],
    'udm': ['Категория'],
    'ug': ['تۈر'],
    'uk': ['Категория', 'Категорія'],
    'ur': ['زمرہ'],
    'uz': ['Turkum', 'Kategoriya'],
    'vec': ['Categoria'],
    'vep': ['Kategorii'],
    'vi': ['Thể_loại', 'Thể loại'],
    'vls': ['Categorie'],
    'vo': ['Klad'],
    'wa': ['Categoreye'],
    'war': ['Kaarangay'],
    'wo': ['Wàll', 'Catégorie'],
    'wuu': ['分类'],
    'xal': ['Янз', 'Әәшл'],
    'xmf': ['კატეგორია'],
    'yi': ['קאטעגאריע', 'קאַטעגאָריע'],
    'yo': ['Ẹ̀ka'],
    'za': ['分类'],
    'zea': ['Categorie'],
    'zh': ['分类', '分類', 'CAT'],
    'zh-classical': ['分類', 'CAT'],
    'zh-min-nan': ['分類', 'Lūi-pia̍t'],
    'zh-yue': ['分类', '分類', '类', '類'],
}


class WikipediaDownloader(DocumentDownloader):

    def __init__(self, download_dir, verbose=False):
        super().__init__()
        self._download_dir = download_dir
        self._verbose = verbose

    def download(self, url):
        urlpath = up.urlparse(url).path[1:]
        output_name = urlpath.replace('/', '-')
        output_file = os.path.join(self._download_dir, output_name)
        if os.path.exists(output_file):
            print(f'bz2 file: {output_file} exists. Not downloading')
        else:
            print(f'Downloading {url} and writing to {output_file}')
            # Download with either wget or s5cmd (aws)
            cmd = ['wget', url, '-O', output_file]
            if self._verbose:
                stdout, stderr = None, None
            else:
                stdout, stderr = subprocess.DEVNULL, subprocess.DEVNULL
            p = subprocess.run(
                cmd,
                stdout=stdout,
                stderr=stderr,
            )
            if p.returncode != 0:
                print(f'Failed to download {url} to {output_file}')

        return output_file


class WikipediaIterator(DocumentIterator):

    def __init__(self, language='en', log_frequency=1000):
        super().__init__()
        self._language = language
        self._log_frequency = log_frequency
        self._counter = 0

    def iterate(self, file_path):
        self._counter = 0
        bname = os.path.split(file_path)[-1]
        input_file = bz2.BZ2File(filename=file_path)
        utf_f = codecs.getreader('utf-8')(input_file)
        context = etree.iterparse(utf_f, events=('end', ))

        for i, (unused_event, elem) in enumerate(context):
            if not elem.tag.endswith('page'):
                continue
            if self._counter > 0 and self._counter % self._log_frequency == 0:
                print(f'Extracted {self._counter} articles from {file_path}')
            self._counter += 1

            namespace = elem.tag[:-4]
            title = elem.find(f'./{namespace}title').text
            ns = elem.find(f'./{namespace}ns').text
            id_ = elem.find(f'./{namespace}id').text
            red_ = elem.find(f'./{namespace}redirect')

            url = f'https://{self._language}.wikipedia.org/wiki/{up.quote(title)}'  # noqa: E501

            # Filter pages that are not in the "main" namespace.
            if ns != '0':
                elem.clear()
                continue

            raw_content = elem.find(
                f'./{namespace}revision/{namespace}text').text
            elem.clear()

            # Filter redirects.
            if raw_content is None or red_ is not None:
                continue

            yield {
                'title': title,
                'id': id_,
                'url': url,
                'language': self._language,
                'source_id': f'{bname}',
            }, raw_content


class WikipediaExtractor(DocumentExtractor):

    def __init__(self, language='en', parser=mwparserfromhell):
        super().__init__()
        self._language = language
        self._parser = parser

    def extract(self, content):
        wikicode = self._parser.parse(content)

        # Filters for magic words / parser instructions -- e.g., __NOTOC__
        re_rm_magic = re.compile('__[A-Z]*__', flags=re.UNICODE)

        # Filters for file/image links.
        media_prefixes = '|'.join(['File', 'Image', 'Media'] +
                                  MEDIA_ALIASES.get(self._language, []))
        re_rm_wikilink = re.compile(f'^(?:{media_prefixes}):',
                                    flags=re.IGNORECASE | re.UNICODE)

        def rm_wikilink(obj):
            return bool(re_rm_wikilink.match(str(obj.title)))

        # Filters for references and tables
        def rm_tag(obj):
            return str(obj.tag) in {'ref', 'table'}

        # Leave category links in-place but remove the category prefixes
        cat_prefixes = '|'.join(['Category'] +
                                CAT_ALIASES.get(self._language, []))
        re_clean_wikilink = re.compile(f'^(?:{cat_prefixes}):',
                                       flags=re.IGNORECASE | re.UNICODE)

        def is_category(obj):
            return bool(re_clean_wikilink.match(str(obj.title)))

        def clean_wikilink(obj):
            text = obj.__strip__()
            text = re.sub(re_clean_wikilink, '', text)
            obj.text = text

        def try_replace_obj(obj):
            try:
                clean_wikilink(obj)
            except ValueError:
                # For unknown reasons, objects are sometimes not found.
                pass

        def try_remove_obj(obj, section):
            try:
                section.remove(obj)
            except ValueError:
                # For unknown reasons, objects are sometimes not found.
                pass

        section_text = []
        # Filter individual sections to clean.
        wiki_code_kwargs = {
            'flat': True,
            'include_lead': True,
            'include_headings': True,
        }
        for section in wikicode.get_sections(**wiki_code_kwargs):
            for obj in section.ifilter_wikilinks(recursive=True):
                if rm_wikilink(obj):
                    try_remove_obj(obj, section)
                elif is_category(obj):
                    try_replace_obj(obj)
                for obj in section.ifilter_tags(matches=rm_tag,
                                                recursive=True):
                    try_remove_obj(obj, section)

            section_text.append(
                re.sub(
                    re_rm_magic,
                    '',
                    section.strip_code().strip(),
                ))
        # Don't return any meta here
        return {}, '\n\n'.join(section_text)


def download_wikipedia(output_path: str,
                       language: str = 'en',
                       dump_date=None,
                       output_type: str = 'jsonl',
                       raw_download_dir=None,
                       keep_raw_download=False,
                       force_download=False,
                       url_limit=None,
                       item_limit=None) -> Dataset:
    """
    Downloads the latest Wikipedia dumps and extracts them using mwparserfromhell

    Args:
      output_path: The path to the root directory of the files
      language: The language of the Wikipedia articles to download
      dump_date: A string formatted as "YYYYMMDD" for the wikipedia dump to use.
        If None, latest dump is used.
      output_type: The file type to save the data as.
      raw_download_dir: Path to store the raw download files for intermediate processing.
        If None, they are stored in a folder named "downloads" under output_path.
      keep_raw_download: If True, keeps the bz2 files that have not been extracted.
      force_download: If False, will skip processing all files in output_paths that already exist and
        directly read from them instead.
      url_limit: The maximum number of raw files to download from the snapshot. If None, all
        files from the range of snapshots are downloaded.
    """  # noqa: E501
    wikipedia_urls = get_wikipedia_urls(language=language, dump_date=dump_date)
    if url_limit:
        wikipedia_urls = wikipedia_urls[:url_limit]
    output_paths = list(
        map(
            lambda url: os.path.join(output_path,
                                     url.split('/')[-1] + f'.{output_type}'),
            wikipedia_urls,
        ))

    if not raw_download_dir:
        raw_download_dir = os.path.join(output_path, 'downloads')
    expand_outdir_and_mkdir(raw_download_dir)

    downloader = WikipediaDownloader(download_dir=raw_download_dir)
    iterator = WikipediaIterator(language=language)
    extractor = WikipediaExtractor(language=language)

    output_format = {
        'text': str,
        'title': str,
        'id': str,
        'url': str,
        'language': str,
        'source_id': str,
        'filename': str,
    }
    dataset = download_and_extract(wikipedia_urls,
                                   output_paths,
                                   downloader,
                                   iterator,
                                   extractor,
                                   output_format,
                                   output_type=output_type,
                                   keep_raw_download=keep_raw_download,
                                   force_download=force_download,
                                   item_limit=item_limit)

    return dataset
