import asyncio
import contextlib
import os
import pathlib
import urllib.parse
from concurrent.futures import ThreadPoolExecutor
from functools import partial

import aiohttp
from tqdm import tqdm, tqdm_notebook

from .results import Results
from .utils import (FailedDownload, Token, default_name, get_filepath,
                    get_ftp_size, get_http_size, in_notebook, run_in_thread)

try:
    import aioftp
except ImportError:
    aioftp = None


__all__ = ['Downloader']


class Downloader:
    """
    Download files in parallel.

    Parameters
    ----------

    max_conn : `int`, optional
        The number of parallel download slots.

    progress : `bool`, optional
        If `True` show a main progress bar showing how many of the total files
        have been downloaded. If `False`, no progress bars will be shown at all.

    file_progress : `bool`, optional
        If `True` and ``progress`` is true, show ``max_conn`` progress bars
        detailing the progress of each individual file being downloaded.

    loop : `asyncio.AbstractEventLoop`, optional
        The event loop to use to download the files. If not specified a new
        loop will be created and executed in a new thread so it does not
        interfere with any currently running event loop.

    notebook : `bool`, optional
        If `True` tqdm will be used in notebook mode. If `None` an attempt will
        be made to detect the notebook and guess which progress bar to use.

    overwrite : `bool` or `str`, optional
        Determine how to handle downloading if a file already exists with the
        same name. If `False` the file download will be skipped and the path
        returned to the existing file, if `True` the file will be downloaded
        and the existing file will be overwritten, if `'unique'` the filename
        will be modified to be unique.
    """

    def __init__(self, max_conn=5, progress=True, file_progress=True,
                 loop=None, notebook=None, overwrite=False):

        self.max_conn = max_conn
        self._start_loop(loop)

        # Configure progress bars
        if notebook is None:
            notebook = in_notebook()
        self.progress = progress
        self.file_progress = file_progress if self.progress else False
        self.tqdm = tqdm if not notebook else tqdm_notebook

        self.overwrite = overwrite

    def _start_loop(self, loop):
        # Setup asyncio loops
        if not loop:
            aio_pool = ThreadPoolExecutor(1)
            self.loop = asyncio.new_event_loop()
            self.run_until_complete = partial(run_in_thread, aio_pool, self.loop)
        else:
            self.loop = loop
            self.run_until_complete = self.loop.run_until_complete

        # Setup queues
        self.http_queue = asyncio.Queue(loop=self.loop)
        self.http_tokens = asyncio.Queue(maxsize=self.max_conn, loop=self.loop)
        self.ftp_queue = asyncio.Queue(loop=self.loop)
        self.ftp_tokens = asyncio.Queue(maxsize=self.max_conn, loop=self.loop)
        for i in range(self.max_conn):
            self.http_tokens.put_nowait(Token(i + 1))
            self.ftp_tokens.put_nowait(Token(i + 1))

    @property
    def queued_downloads(self):
        """
        The total number of files already queued for download.
        """

        return self.http_queue.qsize() + self.ftp_queue.qsize()

    def enqueue_file(self, url, path=None, filename=None, overwrite=None, **kwargs):
        """
        Add a file to the download queue.

        Parameters
        ----------

        url : `str`
            The URL to retrieve.

        path : `str`, optional
            The directory to retrieve the file into, if `None` defaults to the
            current directory.

        filename : `str` or `callable`, optional
            The filename to save the file as. Can also be a callable which
            takes two arguments the url and the response object from opening
            that URL, and returns the filename. (Note, for FTP downloads the
            response will be ``None``.) If `None` the HTTP headers will be read
            for the filename, or the last segment of the URL will be used.

        overwrite : `bool` or `str`, optional
            Determine how to handle downloading if a file already exists with the
            same name. If `False` the file download will be skipped and the path
            returned to the existing file, if `True` the file will be downloaded
            and the existing file will be overwritten, if `'unique'` the filename
            will be modified to be unique. If `None` the value set when
            constructing the `~parfive.Downloader` object will be used.

        kwargs : `dict`
            Extra keyword arguments are passed to `aiohttp.ClientSession.get`
            or `aioftp.ClientSession` depending on the protocol.
        """
        overwrite = overwrite or self.overwrite

        if path is None and filename is None:
            raise ValueError("Either path or filename must be specified.")
        elif path is None:
            path = './'

        path = pathlib.Path(path)
        if not filename:
            filepath = partial(default_name, path)
        elif callable(filename):
            filepath = filename
        else:
            # Define a function because get_file expects a callback
            def filepath(*args):
                return path / filename

        scheme = urllib.parse.urlparse(url).scheme

        if scheme in ('http', 'https'):
            get_file = partial(self._get_http, url=url, filepath_partial=filepath,
                               overwrite=overwrite, **kwargs)
            self.http_queue.put_nowait(get_file)
        elif scheme == 'ftp':
            if aioftp is None:
                raise ValueError("The aioftp package must be installed to download over FTP.")
            get_file = partial(self._get_ftp, url=url, filepath_partial=filepath,
                               overwrite=overwrite, **kwargs)
            self.ftp_queue.put_nowait(get_file)
        else:
            raise ValueError("URL must start with either 'http' or 'ftp'.")

    def download(self, timeouts=None):
        """
        Download all files in the queue.

        Parameters
        ----------
        timeouts : `dict`, optional
            Overrides for the default timeouts for http downloads. Supported
            keys are any accepted by the `aiohttp.ClientTimeout` class. Defaults
            to 5 minutes for total session timeout and 90 seconds for socket
            read timeout.

        Returns
        -------
        filenames : `parfive.Results`
            A list of files downloaded.

        Notes
        -----
        The defaults for the `'total'` and `'sock_read'` timeouts can be
        overridden by two environment variables ``PARFIVE_TOTAL_TIMEOUT`` and
        ``PARFIVE_SOCK_READ_TIMEOUT``.

        """
        timeouts = timeouts or {"total": os.environ.get("PARFIVE_TOTAL_TIMEOUT", 5 * 60),
                                "sock_read": os.environ.get("PARFIVE_SOCK_READ_TIMEOUT", 90)}
        try:
            future = self.run_until_complete(self._run_download(timeouts))
        finally:
            self.loop.stop()
        dlresults = future.result()

        results = Results()

        # Iterate through the results and store any failed download errors in
        # the errors list of the results object.
        for res in dlresults:
            if isinstance(res, FailedDownload):
                results.add_error(res.filepath_partial, res.url, res.exception)
            elif isinstance(res, Exception):
                raise res
            else:
                results.append(res)

        return results

    def retry(self, results):
        """
        Retry any failed downloads in a results object.

        .. note::
            This will start a new event loop.

        Parameters
        ----------

        results : `parfive.Results`
            A previous results object, the ``.errors`` property will be read
            and the downloads retried.

        Returns
        -------

        results : `parfive.Results`
            A modified version of the input ``results`` with all the errors from
            this download attempt and any new files appended to the list of
            file paths.

        """
        # Restart the loop.
        self._start_loop(None)

        for err in results.errors:
            self.enqueue_file(err.url, filename=err.filepath_partial)

        new_res = self.download()

        results += new_res
        results._errors = new_res._errors

        return results

    def _get_main_pb(self, total):
        """
        Return the tqdm instance if we want it, else return a contextmanager
        that just returns None.
        """
        if self.progress:
            return self.tqdm(total=total, unit='file',
                             desc="Files Downloaded",
                             position=0)
        else:
            return contextlib.contextmanager(lambda: iter([None]))()

    async def _run_download(self, timeouts):
        """
        Download all files in the queue.

        Returns
        -------

        results : `parfive.Results`
            A list of filenames which successfully downloaded. This list also
            has an attribute ``errors`` which lists any failed urls and their
            error.
        """
        total_files = self.http_queue.qsize() + self.ftp_queue.qsize()
        done = set()
        with self._get_main_pb(total_files) as main_pb:
            if not self.http_queue.empty():
                done.update(await self._run_http_download(main_pb, timeouts))
            if not self.ftp_queue.empty():
                done.update(await self._run_ftp_download(main_pb, timeouts))

        # Return one future to represent all the results.
        return asyncio.gather(*done, return_exceptions=True)

    async def _run_http_download(self, main_pb, timeouts):
        async with aiohttp.ClientSession(loop=self.loop) as session:
            futures = await self._run_from_queue(self.http_queue, self.http_tokens,
                                                 main_pb, session=session, timeouts=timeouts)

            # Wait for all the coroutines to finish
            done, _ = await asyncio.wait(futures)

            return done

    async def _run_ftp_download(self, main_pb, timeouts):
        futures = await self._run_from_queue(self.ftp_queue, self.ftp_tokens,
                                             main_pb, timeouts=timeouts)
        # Wait for all the coroutines to finish
        done, _ = await asyncio.wait(futures)

        return done

    async def _run_from_queue(self, queue, tokens, main_pb, *, session=None, timeouts):
        futures = []
        while not queue.empty():
            get_file = await queue.get()
            token = await tokens.get()
            file_pb = self.tqdm if self.file_progress else False
            future = asyncio.ensure_future(get_file(session, token=token,
                                                    file_pb=file_pb, timeouts=timeouts))

            def callback(token, future, main_pb):
                tokens.put_nowait(token)
                # Update the main progressbar
                if main_pb and not future.exception():
                    main_pb.update(1)

            future.add_done_callback(partial(callback, token, main_pb=main_pb))
            futures.append(future)

        return futures

    @staticmethod
    async def _get_http(session, *, url, filepath_partial, chunksize=100,
                        file_pb=None, token, overwrite, timeouts, **kwargs):
        """
        Read the file from the given url into the filename given by ``filepath_partial``.

        Parameters
        ----------

        session : `aiohttp.ClientSession`
            The `aiohttp.ClientSession` to use to retrieve the files.

        url : `str`
            The url to retrieve.

        filepath_partial : `callable`
            A function to call which returns the filepath to save the url to.
            Takes two arguments ``resp, url``.

        chunksize : `int`
            The number of bytes to read into the file at a time.

        file_pb : `tqdm.tqdm` or `False`
            Should progress bars be displayed for each file downloaded.

        token : `parfive.downloader.Token`
            A token for this download slot.

        kwargs : `dict`
            Extra keyword arguments are passed to `aiohttp.ClientSession.get`.

        Returns
        -------

        filepath : `str`
            The name of the file saved.

        """
        timeout = aiohttp.ClientTimeout(**timeouts)
        try:
            async with session.get(url, timeout=timeout, **kwargs) as resp:
                if resp.status != 200:
                    raise FailedDownload(filepath_partial, url, resp)
                else:
                    filepath, skip = get_filepath(filepath_partial(resp, url), overwrite)
                    if skip:
                        return str(filepath)
                    if callable(file_pb):
                        file_pb = file_pb(position=token.n, unit='B', unit_scale=True,
                                          desc=filepath.name, leave=False,
                                          total=get_http_size(resp))
                    else:
                        file_pb = None
                    with open(str(filepath), 'wb') as fd:
                        while True:
                            chunk = await resp.content.read(chunksize)
                            if not chunk:
                                # Close the file progressbar
                                if file_pb is not None:
                                    file_pb.close()

                                return str(filepath)

                            # Write this chunk to the output file.
                            fd.write(chunk)

                            # Update the progressbar for file
                            if file_pb is not None:
                                file_pb.update(chunksize)

        except Exception as e:
            raise FailedDownload(filepath_partial, url, e)

    @staticmethod
    async def _get_ftp(session=None, *, url, filepath_partial,
                       file_pb=None, token, overwrite, timeouts, **kwargs):
        """
        Read the file from the given url into the filename given by ``filepath_partial``.

        Parameters
        ----------

        session : `None`
            A placeholder for API compatibility with ``_get_http``

        url : `str`
            The url to retrieve.

        filepath_partial : `callable`
            A function to call which returns the filepath to save the url to.
            Takes two arguments ``resp, url``.

        file_pb : `tqdm.tqdm` or `False`
            Should progress bars be displayed for each file downloaded.

        token : `parfive.downloader.Token`
            A token for this download slot.

        kwargs : `dict`
            Extra keyword arguments are passed to `~aioftp.ClientSession`.

        Returns
        -------

        filepath : `str`
            The name of the file saved.

        """
        parse = urllib.parse.urlparse(url)
        try:
            async with aioftp.ClientSession(parse.hostname, **kwargs) as client:
                if parse.username and parse.password:
                    client.login(parse.username, parse.password)

                # This has to be done before we start streaming the file:
                total_size = await get_ftp_size(client, parse.path)
                async with client.download_stream(parse.path) as stream:
                    filepath, skip = get_filepath(filepath_partial(None, url), overwrite)
                    if skip:
                        return str(filepath)
                    if callable(file_pb):
                        file_pb = file_pb(position=token.n, unit='B', unit_scale=True,
                                          desc=filepath.name, leave=False, total=total_size)
                    else:
                        file_pb = None

                    with open(str(filepath), 'wb') as fd:
                        async for chunk in stream.iter_by_block():
                            # Write this chunk to the output file.
                            fd.write(chunk)

                            # Update the progressbar for file
                            if file_pb is not None:
                                file_pb.update(len(chunk))

                        # Close the file progressbar
                        if file_pb is not None:
                            file_pb.close()

                        return str(filepath)

        except Exception as e:
            raise FailedDownload(filepath_partial, url, e)
