import csv
import json
import logging
from collections import defaultdict
from pathlib import Path
from typing import Dict, List

import pandas as pd
from tqdm import tqdm

"""
Wikidata related
"""


def load_wd_properties(file_path: Path) -> Dict:
    qid2props = {}
    logging.info(f"loading Wikidata item properties from {file_path}")
    with open(file_path, "r") as rf:
        for line in tqdm(rf):
            data = json.loads(line.strip())
            if data["qid"] not in qid2props:
                qid2props[data["qid"]] = {}
            qid2props[data["qid"]][data["pid"]] = data["datavalue"]
    return qid2props


def load_wd_descriptions(file_path: Path) -> Dict:
    qid2desc = {}
    logging.info(f"loading Wikidata item descriptions from {file_path}")
    with open(file_path, "r") as rf:
        for line in tqdm(rf):
            data = json.loads(line.strip())
            if data["qid"] not in qid2desc:
                qid2desc[data["qid"]] = {}
            qid2desc[data["qid"]][data["lang"]] = {
                "label": data["label"],
                "description": data["description"],
            }
    return qid2desc


def load_exclusion_wd_props(file_path: Path) -> Dict[str, List]:
    prop2qids = defaultdict(list)
    logging.info(f"loading list of Wikidata properties to exclude from {file_path}")
    with open(file_path, "r") as rf:
        _ = rf.readline()
        for line in rf:
            prop, _, qid, _ = line.split("\t")
            prop2qids[prop] += [qid]
    return prop2qids


"""
Wikidata-Wikipedia
"""


def load_wd2wikipedia_links(file_path: Path):
    """Links from Wikipedia articles to Wikidata items."""
    logging.info("loading Wikidata to Wikipedia links")
    df = pd.read_csv(file_path, sep="\t", quoting=csv.QUOTE_NONE)
    logging.info(f"# language titles: {len(df)}")
    return df


def write_wd2wikipedia_links(data: Dict, file_path: Path):
    logging.info(f"writing Wikidata to Wikipedia links to {file_path}")
    data.to_csv(file_path, sep="\t", index=False, quoting=csv.QUOTE_NONE)


def get_title2wd(wd2wikipedia_links: pd.DataFrame):
    title2wd = {}
    for _, row in wd2wikipedia_links.iterrows():
        title2wd[(row["Language"], row["Wikipedia Title"])] = row["QID"]
    return title2wd


def load_title2wd(file_path: Path):
    """links from language Wikipedia articles to Wikidata items"""
    logging.info("loading inter-wiki links")
    title2wd = {}
    with open(file_path, "r") as rf:
        header = rf.readline()
        for line in tqdm(rf):
            splits = line.strip().split("\t")
            title2wd[(splits[0], splits[1])] = splits[2]
    return title2wd


"""
Wikidata-Wikinews
"""


def load_wd2wn_links(file_path: Path):
    """Links from Wikinews articles to Wikidata items."""
    logging.info("loading Wikidata to Wikinews links")
    df = pd.read_csv(file_path, sep="\t", quoting=csv.QUOTE_NONE)
    logging.info(f"# language titles: {len(df)}")
    return df


def write_wd2wn_links(data: Dict, file_path: Path):
    logging.info(f"writing Wikidata to Wikinews links to {file_path}")
    data.to_csv(file_path, sep="\t", index=False, quoting=csv.QUOTE_NONE)


def get_wn_title2wd(wd2wn_links: pd.DataFrame):
    title2wd = {}
    for _, row in wd2wn_links.iterrows():
        title2wd[(row["Language"], row["Wikinews Title"])] = row["QID"]
    logging.info(f"Wikinews event titles: {len(title2wd)}")
    return title2wd


def get_wiki_title2wd(label_set: pd.DataFrame):
    # this loads information from final label_desc.tsv from XLEL dataset
    title2wd = {}
    wd2title = {}
    for _, row in label_set.iterrows():
        # used for identifying inlink
        title2wd[(row["Wikipedia Language"], row["Wikipedia Title"])] = row[
            "Wikidata Item"
        ]
        # used for remapping wikinews inlink to the corresponding Wikipedia title
        wd2title[(row["Wikipedia Language"], row["Wikidata Item"])] = row[
            "Wikipedia Title"
        ]
    logging.info(f"Wikipedia event titles: {len(title2wd)}")
    return title2wd, wd2title


"""
Wikinews
"""


def load_wn_dates(dir_path: Path, version: str = "20220301"):

    id2dates = defaultdict(list)

    for lang_dir_path in dir_path.iterdir():
        lg = lang_dir_path.stem
        tsv_path = lang_dir_path / f"{lg}wikinews-{version}-meta.tsv"
        if not tsv_path.is_file():
            continue
        with open(tsv_path, "r") as rf:
            header = rf.readline()
            for line in rf:
                _id, _date = line.strip().split("\t")
                id2dates[(lg, _id)] += [_date]
    logging.info(f"loaded publication dates for {len(id2dates)} articles")

    return id2dates


"""
Wikipedia
"""


def load_wiki_event_descriptions(file_path: Path):
    item2wiki_desc = {}
    df = pd.read_csv(file_path, sep="\t", quoting=csv.QUOTE_NONE)
    logging.info(f"loading Wikipedia event descriptions from {file_path}")
    for _, row in tqdm(df.iterrows()):
        if row["Wikidata Item"] not in item2wiki_desc:
            item2wiki_desc[row["Wikidata Item"]] = {}
        label = row["Wikipedia Title"]
        description = row["Wikipedia Description"]
        item2wiki_desc[row["Wikidata Item"]][row["Wikipedia Language"]] = {
            "label": label if not pd.isnull(label) else "",
            "description": description if not pd.isnull(description) else "",
        }
    return item2wiki_desc


"""
XL-EL
"""


def load_xlel_tsv(
    file_path: Path,
    return_keys: List[str] = ["Wikidata Item"],
    return_iterator: bool = False,
    chunk_size: int = 100000,
):
    logging.info(f"loading xlec data from {file_path}")
    logging.info("storing keys: " + " & ".join(return_keys))
    if return_iterator:
        logging.info(f"using iterator with chunksize: {chunk_size}")
        return pd.read_csv(
            file_path,
            sep="\t",
            quoting=csv.QUOTE_NONE,
            usecols=return_keys,
            iterator=return_iterator,
            chunksize=chunk_size,
        )
    else:
        return pd.read_csv(
            file_path, sep="\t", quoting=csv.QUOTE_NONE, usecols=return_keys
        )


def load_xlel_label_desc(
    file_path: Path,
    return_keys: List[str] = ["Wikidata Item", "Wikipedia Language", "Wikipedia Title"],
):
    logging.info(f"loading xlel label set from {file_path}")
    logging.info("storing keys: " + " & ".join(return_keys))
    return pd.read_csv(file_path, sep="\t", quoting=csv.QUOTE_NONE, usecols=return_keys)
