import json
import os
import shutil
import tarfile
from collections import Iterator
from functools import reduce
from itertools import chain
from typing import Any, Callable, Dict, List, Union
from urllib.request import urlopen

import requests
import xmltodict as xml

from .text import _get_abstract, _get_text

searchtext = Union[str, List[str]]
searchtexts = Union[searchtext, List[searchtext]]
textlist = List[Dict[str, Any]]


class Paperset(Iterator):
    def __init__(self, directory: str) -> None:
        """
        The Paperset class:
            __init__ args:
                directory: a string, the directory where the jsons are stored

            description:
                lazy loader for cord-19 text files. Data is not actually loaded
                until indexing, until then it just indexes files. Can be
                indexed with both ints and slices.
        """
        self.directory = directory
        self.dir_dict = {idx: f for idx, f in enumerate(os.listdir(self.directory))}
        self.keys = list(self.dir_dict.keys())

    def _load_file(self, path: str) -> dict:
        with open(f"{self.directory}/{path}") as handle:
            outdict = json.loads(handle.read())
        return outdict

    def __getitem__(self, indices: Union[int, slice]) -> Union[list, dict]:
        slicedkeys = list(self.dir_dict.keys())[indices]
        if not isinstance(slicedkeys, list):
            slicedkeys = [slicedkeys]
        out = [self._load_file(self.dir_dict[key]) for key in slicedkeys]
        if len(out) == 1:
            return out[0]
        else:
            return out

    def __next__(self) -> Dict[Any, Any]:
        key = self.keys.pop()
        return self._load_file(self.dir_dict[key])

    def apply(self, fn: Callable[..., Any]) -> List[Any]:
        return [fn(self._load_file(self.dir_dict[k])) for k in self.dir_dict.keys()]

    def texts(self) -> List[str]:
        return self.apply(_get_text)

    def abstracts(self) -> List[str]:
        return self.apply(_get_abstract)

    def __len__(self) -> int:
        return len(self.dir_dict.keys())


def _search(ps: Union[Paperset, textlist], txt: Any) -> textlist:
    if type(txt[0]) is list:
        raise ValueError("Items of the search cannot be nested lists!")
    if type(txt) is str:
        txt = [txt]
    txt = [x.lower() for x in txt]
    return [
        x
        for x in ps
        if any(c in _get_text(x).lower() for c in txt)
        or any(c in _get_abstract(x).lower() for c in txt)
    ]


def search(
    ps: Union[Paperset, textlist],
    terms: Union[searchtexts, searchtext, List[List[str]]],
) -> textlist:
    if type(terms) is not list:
        raise ValueError("search terms must be a list!!")
    if len(terms) != 1:
        return reduce(lambda x, y: _search(x, y), terms, ps)
    else:
        return _search(ps, terms)


def download(dir: str = ".") -> None:
    site = xml.parse(
        requests.get(
            "https://ai2-semanticscholar-cord-19.s3-us-west-2.amazonaws.com/"
        ).content
    )["ListBucketResult"]["Contents"][::-1][:10]
    key = [x["Key"] for x in site]
    urls = [
        f"https://ai2-semanticscholar-cord-19.s3-us-west-2.amazonaws.com/{k}"
        for k in key
    ]
    keys = [k.split("/")[-1] for k in key]
    data = dict(zip(keys, urls))

    if not os.path.exists(dir):
        os.mkdir(dir)
    for d in data.keys():
        print(f"downloading {data[d]}")
        handle = urlopen(data[d])
        if d.replace(".tar.gz", "") in os.listdir(f"{dir}"):
            shutil.rmtree(f"{dir}/{d.replace('.tar.gz','')}", ignore_errors=True)
        with open(f"{dir}/{d}", "wb") as out:
            while True:
                dat = handle.read(1024)
                if len(dat) == 0:
                    break
                out.write(dat)
    for f in os.listdir(dir):
        if tarfile.is_tarfile(f"{dir}/{f}"):
            print(f"Extracting {dir}/{f}")
            tar = tarfile.open(f"{dir}/{f}", "r:gz")
            tar.extractall(path=dir)
            tar.close()
            os.remove(f"{dir}/{f}")
