import logging
from copy import deepcopy
from os import mkdir
from os.path import join
from shutil import rmtree
from typing import Any, Iterator, List, Optional, Tuple, Union

from slugify import slugify

from .base_downloader import BaseDownload, DownloadError
from .downloader import Download
from .loader import load_json, load_text, load_yaml
from .path import get_filename_extension_from_url
from .saver import save_json, save_text, save_yaml
from .typehint import ListDict, ListTuple

logger = logging.getLogger(__name__)


class Retrieve(BaseDownload):
    """Retrieve class which takes in a Download object and can either download, download
    and save or use previously downloaded and saved data. It also allows the use of a
    static fallback when downloading fails.

    Args:
        downloader (Download): Download object
        fallback_dir (str): Directory containing static fallback data
        saved_dir (str): Directory to save or load downloaded data
        temp_dir (str): Temporary directory for when data is not needed after downloading
        save (bool): Whether to save downloaded data. Defaults to False.
        use_saved (bool): Whether to use saved data. Defaults to False.
        prefix (str): Prefix to add to filenames. Defaults to "".
        delete (bool): Whether to delete saved_dir if save is True. Defaults to True.
        log_level (int): Level at which to log messages. Defaults to logging.INFO.
    """

    retrievers = dict()

    def __init__(
        self,
        downloader: Download,
        fallback_dir: str,
        saved_dir: str,
        temp_dir: str,
        save: bool = False,
        use_saved: bool = False,
        prefix: str = "",
        delete: bool = True,
        log_level: int = logging.INFO,
    ):
        self.downloader = downloader
        self.fallback_dir = fallback_dir
        self.saved_dir = saved_dir
        self.temp_dir = temp_dir
        self.save = save
        self.use_saved = use_saved
        self.prefix = prefix
        self.check_flags(saved_dir, save, use_saved, delete)
        self.log_level = log_level

    @staticmethod
    def check_flags(
        saved_dir: str, save: bool, use_saved: bool, delete: bool
    ) -> None:
        """Check flags. Also delete saved_dir if save and delete are True

        Args:
            saved_dir (str): Directory to save or load downloaded data
            save (bool): Whether to save downloaded data
            use_saved (bool): Whether to use saved data
            delete (bool): Whether to delete saved_dir if save is True

        Returns:
            None

        """
        if save:
            if use_saved:
                raise ValueError(
                    "Either the save or use_saved flags can be set to True!"
                )
            if delete:
                rmtree(saved_dir, ignore_errors=True)
                mkdir(saved_dir)

    @staticmethod
    def get_url_logstr(url: str) -> str:
        """Url string that will be logged. It is limited to 100 characters if necessary.

        Args:
            url (str): URL to download

        Returns:
            str: Url string to use in logs

        """
        if len(url) > 100:
            return f"{url[:100]}..."
        return url

    def clone(self, downloader: Download) -> "Retrieve":
        """Clone a given retriever but use the given downloader

        Args:
            downloader (Download): Downloader to use

        Returns:
            Retrieve: Cloned retriever

        """
        return Retrieve(
            downloader,
            fallback_dir=self.fallback_dir,
            saved_dir=self.saved_dir,
            temp_dir=self.temp_dir,
            save=self.save,
            use_saved=self.use_saved,
            prefix=self.prefix,
            delete=False,
        )

    def get_filename(
        self,
        url: str,
        filename: Optional[str] = None,
        possible_extensions: Tuple[str, ...] = tuple(),
        **kwargs: Any,
    ) -> Tuple[str, Any]:
        """Get filename from url and given parameters

        Args:
            url (str): Url from which to get filename
            filename (optional[str]): Filename to use. Defaults to None (infer from url).
            possible_extensions (Tuple[str, ...]): Possible extensions to look for in url
            **kwargs: See below
            format (str): Given extension to look for in url
            file_type (str): Given extension to look for in url

        Returns:
            Tuple[str, Any]: Tuple of (filename, kwargs)

        """
        prefix = kwargs.pop("file_prefix", self.prefix)
        if prefix:
            prefix = f"{prefix}_"
        if filename:
            return f"{prefix}{filename}", kwargs
        filename, extension = get_filename_extension_from_url(
            url, second_last=True, use_query=True
        )
        filename = slugify(filename)
        extensions = list()
        format = kwargs.get("format")
        if format:
            extensions.append(format)
        file_type = kwargs.get("file_type")
        if file_type:
            extensions.append(file_type)
        if possible_extensions:
            extensions.extend(possible_extensions)
        if not extensions:
            return f"{prefix}{filename}{extension}", kwargs
        first_ext = f".{extensions[0].lower()}"
        if not extension:
            return f"{prefix}{filename}{first_ext}", kwargs
        for candidate in extensions:
            if candidate == extension[1:]:
                return f"{prefix}{filename}{extension}", kwargs
        filename = slugify(f"{filename}{extension}")
        return f"{prefix}{filename}{first_ext}", kwargs

    def download_file(
        self,
        url: str,
        filename: Optional[str] = None,
        logstr: Optional[str] = None,
        fallback: bool = False,
        log_level: int = None,
        **kwargs: Any,
    ) -> str:
        """Retrieve file

        Args:
            url (str): URL to download
            filename (Optional[str]): Filename of saved file. Defaults to getting from url.
            logstr (Optional[str]): Text to use in log string to describe download. Defaults to filename.
            fallback (bool): Whether to use static fallback if download fails. Defaults to False.
            log_level (int): Level at which to log messages. Overrides level from constructor.
            **kwargs: Parameters to pass to download_file call

        Returns:
            str: Path to downloaded file

        """
        if log_level is None:
            log_level = self.log_level
        filename, kwargs = self.get_filename(url, filename, **kwargs)
        if not logstr:
            logstr = filename
        if self.save:
            folder = self.saved_dir
        else:
            folder = self.temp_dir
        output_path = join(folder, filename)
        saved_path = join(self.saved_dir, filename)
        if self.use_saved:
            logger.log(log_level, f"Using saved {logstr} in {saved_path}")
            return saved_path
        else:
            try:
                logger.log(
                    log_level,
                    f"Downloading {logstr} from {self.get_url_logstr(url)} into {output_path}",
                )
                return self.downloader.download_file(
                    url, path=output_path, **kwargs
                )
            except DownloadError:
                if not fallback:
                    raise
                fallback_path = join(self.fallback_dir, filename)
                logger.exception(
                    f"{logstr} download failed, using static data {fallback_path}!"
                )
                return fallback_path

    def download_text(
        self,
        url: str,
        filename: Optional[str] = None,
        logstr: Optional[str] = None,
        fallback: bool = False,
        log_level: int = None,
        **kwargs: Any,
    ) -> str:
        """Download text

        Args:
            url (str): URL to download
            filename (Optional[str]): Filename of saved file. Defaults to getting from url.
            logstr (Optional[str]): Text to use in log string to describe download. Defaults to filename.
            fallback (bool): Whether to use static fallback if download fails. Defaults to False.
            log_level (int): Level at which to log messages. Overrides level from constructor.
            **kwargs: Parameters to pass to download_text call

        Returns:
            str: The text from the file

        """
        if log_level is None:
            log_level = self.log_level
        filename, kwargs = self.get_filename(url, filename, **kwargs)
        if not logstr:
            logstr = filename
        saved_path = join(self.saved_dir, filename)
        if self.use_saved:
            logger.log(log_level, f"Using saved {logstr} in {saved_path}")
            text = load_text(saved_path)
        else:
            try:
                logger.log(
                    log_level,
                    f"Downloading {logstr} from {self.get_url_logstr(url)}",
                )
                text = self.downloader.download_text(url, **kwargs)
                if self.save:
                    logger.log(log_level, f"Saving {logstr} in {saved_path}")
                    save_text(text, saved_path)
            except DownloadError:
                if not fallback:
                    raise
                fallback_path = join(self.fallback_dir, filename)
                logger.exception(
                    f"{logstr} download failed, using static data {fallback_path}!"
                )
                text = load_text(fallback_path)
        return text

    def download_yaml(
        self,
        url: str,
        filename: Optional[str] = None,
        logstr: Optional[str] = None,
        fallback: bool = False,
        log_level: int = None,
        **kwargs: Any,
    ) -> Any:
        """Retrieve YAML

        Args:
            url (str): URL to download
            filename (Optional[str]): Filename of saved file. Defaults to getting from url.
            logstr (Optional[str]): Text to use in log string to describe download. Defaults to filename.
            fallback (bool): Whether to use static fallback if download fails. Defaults to False.
            log_level (int): Level at which to log messages. Overrides level from constructor.
            **kwargs: Parameters to pass to download_yaml call

        Returns:
            Any: The data from the YAML file

        """
        if log_level is None:
            log_level = self.log_level
        filename, kwargs = self.get_filename(
            url, filename, ("yaml", "yml"), **kwargs
        )
        if not logstr:
            logstr = filename
        saved_path = join(self.saved_dir, filename)
        if self.use_saved:
            logger.log(log_level, f"Using saved {logstr} in {saved_path}")
            ryaml = load_yaml(saved_path)
        else:
            try:
                logger.log(
                    log_level,
                    f"Downloading {logstr} from {self.get_url_logstr(url)}",
                )
                ryaml = self.downloader.download_yaml(url, **kwargs)
                if self.save:
                    logger.log(log_level, f"Saving {logstr} in {saved_path}")
                    save_yaml(ryaml, saved_path)
            except DownloadError:
                if not fallback:
                    raise
                fallback_path = join(self.fallback_dir, filename)
                logger.exception(
                    f"{logstr} download failed, using static data {fallback_path}!"
                )
                ryaml = load_yaml(fallback_path)
        return ryaml

    def download_json(
        self,
        url: str,
        filename: Optional[str] = None,
        logstr: Optional[str] = None,
        fallback: bool = False,
        log_level: int = None,
        **kwargs: Any,
    ) -> Any:
        """Retrieve JSON

        Args:
            url (str): URL to download
            filename (Optional[str]): Filename of saved file. Defaults to getting from url.
            logstr (Optional[str]): Text to use in log string to describe download. Defaults to filename.
            fallback (bool): Whether to use static fallback if download fails. Defaults to False.
            log_level (int): Level at which to log messages. Overrides level from constructor.
            **kwargs: Parameters to pass to download_json call

        Returns:
            Any: The data from the JSON file

        """
        if log_level is None:
            log_level = self.log_level
        filename, kwargs = self.get_filename(
            url, filename, ("json",), **kwargs
        )
        if not logstr:
            logstr = filename
        saved_path = join(self.saved_dir, filename)
        if self.use_saved:
            logger.log(log_level, f"Using saved {logstr} in {saved_path}")
            rjson = load_json(saved_path)
        else:
            try:
                logger.log(
                    log_level,
                    f"Downloading {logstr} from {self.get_url_logstr(url)}",
                )
                rjson = self.downloader.download_json(url, **kwargs)
                if self.save:
                    logger.log(log_level, f"Saving {logstr} in {saved_path}")
                    save_json(rjson, saved_path)
            except DownloadError:
                if not fallback:
                    raise
                fallback_path = join(self.fallback_dir, filename)
                logger.exception(
                    f"{logstr} download failed, using static data {fallback_path}!"
                )
                rjson = load_json(fallback_path)
        return rjson

    def get_tabular_rows(
        self,
        url: Union[str, ListTuple[str]],
        headers: Union[int, ListTuple[int], ListTuple[str]] = 1,
        dict_form: bool = False,
        has_hxl: bool = False,
        filename: Optional[str] = None,
        logstr: Optional[str] = None,
        fallback: bool = False,
        **kwargs: Any,
    ) -> Tuple[List[str], Iterator[ListDict]]:
        """Returns header of tabular file(s) pointed to by url and an iterator where
        each row is returned as a list or dictionary depending on the dict_rows argument.
        When a list of urls is supplied (in url), then the has_hxl flag indicates if the
        files are HXLated so that the HXL row is only included from the first file.

        The headers argument is either a row number or list of row numbers (in case of
        multi-line headers) to be considered as headers (rows start counting at 1), or
        the actual headers defined as a list of strings. It defaults to 1.
        The dict_form arguments specifies if each row should be returned as a dictionary
        or a list, defaulting to a list.

        Args:
            url (Union[str, ListTuple[str]]): A single or list of URLs or paths to read from
            headers (Union[int, ListTuple[int], ListTuple[str]]): Number of row(s) containing headers or list of headers. Defaults to 1.
            dict_form (bool): Return dict or list for each row. Defaults to False (list)
            has_hxl (bool): Whether files have HXL hashtags. Defaults to False.
            filename (Optional[str]): Filename of saved file. Defaults to getting from url.
            logstr (Optional[str]): Text to use in log string to describe download. Defaults to filename.
            fallback (bool): Whether to use static fallback if download fails. Defaults to False.
            **kwargs: Parameters to pass to download_file call

        Returns:
            Tuple[List[str],Iterator[ListDict]]: Tuple (headers, iterator where each row is a list or dictionary)

        """
        if isinstance(url, list):
            is_list = True
            orig_kwargs = deepcopy(kwargs)
            urls = url
            url = urls[0]
        else:
            is_list = False
        path = self.download_file(url, filename, logstr, fallback, **kwargs)
        if is_list:
            path = [path]
            for url in urls[1:]:
                temp_kwargs = deepcopy(orig_kwargs)
                pth = self.download_file(
                    url, None, logstr, fallback, **temp_kwargs
                )
                path.append(pth)

        kwargs.pop("file_prefix", None)
        return self.downloader.get_tabular_rows(
            path, headers, dict_form, has_hxl=has_hxl, **kwargs
        )

    @classmethod
    def generate_retrievers(
        cls,
        fallback_dir: str,
        saved_dir: str,
        temp_dir: str,
        save: bool = False,
        use_saved: bool = False,
        ignore: ListTuple[str] = tuple(),
        delete: bool = True,
        **kwargs: Any,
    ) -> None:
        """Generate retrievers. Retrievers are generated from downloaders so
        Download.generate_downloaders() needs to have been called first. Each retriever
        can either download, download and save or use previously downloaded and saved
        data. It also allows the use of a static fallback when downloading fails.

        Args:
            fallback_dir (str): Directory containing static fallback data
            saved_dir (str): Directory to save or load downloaded data
            temp_dir (str): Temporary directory for when data is not needed after downloading
            save (bool): Whether to save downloaded data. Defaults to False.
            use_saved (bool): Whether to use saved data. Defaults to False.
            ignore (ListTuple[str]): Don't generate retrievers for these downloaders
            delete (bool): Whether to delete saved_dir if save is True. Defaults to True.
            **kwargs (Any): Any other arguments to pass.

        Returns:
            None
        """
        cls.check_flags(saved_dir, save, use_saved, delete)
        cls.retrievers = dict()
        for name, downloader in Download.downloaders.items():
            if name in ignore:
                continue
            cls.retrievers[name] = cls(
                downloader,
                fallback_dir,
                saved_dir,
                temp_dir,
                save,
                use_saved,
                delete=False,
                **kwargs,
            )

    @classmethod
    def get_retriever(cls, name: Optional[str] = None) -> "Retrieve":
        """Get a generated retriever given a name. If name is not supplied, the default
        one will be returned.

        Args:
            name (Optional[str]): Name of retriever. Defaults to None (get default).

        Returns:
            Retriever: Retriever object
        """
        return cls.retrievers.get(name, cls.retrievers["default"])
