import os
import sys
import asyncio
import logging
import pathlib
import warnings
import contextlib
import urllib.parse
from functools import partial, lru_cache
from concurrent.futures import ThreadPoolExecutor

import aiohttp
from tqdm import tqdm, tqdm_notebook

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

try:
    import aioftp
except ImportError:  # pragma: nocover
    aioftp = None


try:
    import aiofiles  # pragma: nocover
except ImportError:
    aiofiles = None


SERIAL_MODE = "PARFIVE_SINGLE_DOWNLOAD" in os.environ
DISABLE_RANGE = "PARFIVE_DISABLE_RANGE" in os.environ or SERIAL_MODE

__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
        No longer used, and will be removed in a future release.
    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.
    headers : `dict`
       Request headers to be passed to the server.
       Adds `User-Agent` information about `parfive`, `aiohttp` and `python` if not passed explicitely.
    """

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

        if loop:
            warnings.warn('The loop argument is no longer used, and will be '
                          'removed in a future release.')
        self.max_conn = max_conn if not SERIAL_MODE else 1
        self._init_queues()

        # 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

        self.headers = headers
        if headers is None or 'User-Agent' not in headers:
            self.headers = {
                'User-Agent': f"parfive/{parfive.__version__} aiohttp/{aiohttp.__version__} python/{sys.version[:5]}"}

        self._use_aiofiles = use_aiofiles

    def _init_queues(self):
        # Setup queues
        self.http_queue = _QueueList()
        self.ftp_queue = _QueueList()

    def _generate_tokens(self):
        # Create a Queue with max_conn tokens
        queue = asyncio.Queue(maxsize=self.max_conn)
        for i in range(self.max_conn):
            queue.put_nowait(Token(i + 1))
        return queue

    @property
    @lru_cache()
    def use_aiofiles(self):
        """
        aiofiles will be used if installed and must be explicitly enabled

        PARFIVE_OVERWRITE_ENABLE_AIOFILES takes precedence if present,
        aiofiles will not be used

        finally the Downloader's constructor argument is considered.
        """
        if aiofiles is None:
            return False

        if "PARFIVE_OVERWRITE_ENABLE_AIOFILES" in os.environ:
            return True

        return self._use_aiofiles

    @property
    @lru_cache()
    def default_chunk_size(self):
        """
        aiofiles requires a different default chunk size
        """
        return 1024 if self.use_aiofiles else 100

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

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

    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.Client.context` depending on the protocol.

        Notes
        -----
        Proxy URL is read from the environment variables `HTTP_PROXY` or
        `HTTPS_PROXY`, depending on the protocol of the `url` passed. Proxy
        Authentication `proxy_auth` should be passed as a `aiohttp.BasicAuth`
        object. Proxy Headers `proxy_headers` should be passed as `dict`
        object.
        """
        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.append(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.append(get_file)
        else:
            raise ValueError("URL must start with either 'http' or 'ftp'.")

    @staticmethod
    def _run_in_loop(coro):
        """
        Detect an existing, running loop and run in a separate loop if needed.

        If no loop is running, use asyncio.run to run the coroutine instead.
        """
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop and loop.is_running():
            aio_pool = ThreadPoolExecutor(1)
            new_loop = asyncio.new_event_loop()
            return run_in_thread(aio_pool, new_loop, coro)

        return asyncio.run(coro)

    @staticmethod
    def _configure_debug():  # pragma: no cover

        sh = logging.StreamHandler()
        sh.setLevel(logging.DEBUG)

        formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
        sh.setFormatter(formatter)

        parfive.log.addHandler(sh)
        parfive.log.setLevel(logging.DEBUG)

        aiohttp_logger = logging.getLogger('aiohttp.client')
        aioftp_logger = logging.getLogger('aioftp.client')

        aioftp_logger.addHandler(sh)
        aioftp_logger.setLevel(logging.DEBUG)

        aiohttp_logger.addHandler(sh)
        aiohttp_logger.setLevel(logging.DEBUG)

        parfive.log.debug("Configured parfive to run with debug logging...")

    async def run_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
        -------
        `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``.
        """
        # Setup debug logging before starting a download
        if "PARFIVE_DEBUG" in os.environ:
            self._configure_debug()

        timeouts = timeouts or {"total": os.environ.get("PARFIVE_TOTAL_TIMEOUT", 5 * 60),
                                "sock_read": os.environ.get("PARFIVE_SOCK_READ_TIMEOUT", 90)}

        total_files = self.queued_downloads

        done = set()
        with self._get_main_pb(total_files) as main_pb:
            if len(self.http_queue):
                done.update(await self._run_http_download(main_pb, timeouts))
            if len(self.ftp_queue):
                done.update(await self._run_ftp_download(main_pb, timeouts))

        dl_results = await asyncio.gather(*done, return_exceptions=True)

        results = Results()

        # Iterate through the results and store any failed download errors in
        # the errors list of the results object.
        for res in dl_results:
            if isinstance(res, FailedDownload):
                results.add_error(res.filepath_partial, res.url, res.exception)
                parfive.log.info(f'{res.url} failed to download with exception\n'
                                 f'{res.exception}')
            elif isinstance(res, Exception):
                raise res
            else:
                results.append(res)

        return results

    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
        -------
        `parfive.Results`
            A list of files downloaded.

        Notes
        -----
        This is a synchronous version of `~parfive.Downloader.run_download`, an
        `asyncio` event loop will be created to run the download (in it's own
        thread if a loop is already running).

        The defaults for the `'total'` and `'sock_read'` timeouts can be
        overridden by two environment variables ``PARFIVE_TOTAL_TIMEOUT`` and
        ``PARFIVE_SOCK_READ_TIMEOUT``.
        """
        return self._run_in_loop(self.run_download(timeouts))

    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
        -------
        `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.
        """
        # Reset the queues
        self._init_queues()

        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

    @classmethod
    def simple_download(cls, urls, *, path="./", overwrite=None):
        """
        Download a series of URLs to a single destination.

        Parameters
        ----------
        urls : iterable
            A sequence of URLs to download.

        path : `pathlib.Path`, optional
            The destination directory for the downloaded files.
            Defaults to the current directory.

        overwrite: `bool`, optional
            Overwrite the files at the destination directory. If `False` the
            URL will not be downloaded if a file with the corresponding
            filename already exists.

        Returns
        -------
        `parfive.Results`
            A list of files downloaded.
        """
        dl = cls()
        for url in urls:
            dl.enqueue_file(url, path=path, overwrite=overwrite)
        return dl.download()

    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_http_download(self, main_pb, timeouts):
        async with aiohttp.ClientSession(headers=self.headers) as session:
            self._generate_tokens()
            futures = await self._run_from_queue(
                self.http_queue.generate_queue(),
                self._generate_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.generate_queue(),
            self._generate_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.create_task(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

    async def _get_http(self, session, *, url, filepath_partial, chunksize=None,
                        file_pb=None, token, overwrite, timeouts, max_splits=5, **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.
        max_splits: `int`
            Number of maximum concurrent connections per file.
        kwargs : `dict`
            Extra keyword arguments are passed to `aiohttp.ClientSession.get`.

        Returns
        -------
        `str`
            The name of the file saved.
        """
        if chunksize is None:
            chunksize =  self.default_chunk_size

        timeout = aiohttp.ClientTimeout(**timeouts)
        try:
            scheme = urllib.parse.urlparse(url).scheme
            if 'HTTP_PROXY' in os.environ and scheme == 'http':
                kwargs['proxy'] = os.environ['HTTP_PROXY']
            elif 'HTTPS_PROXY' in os.environ and scheme == 'https':
                kwargs['proxy'] = os.environ['HTTPS_PROXY']

            async with session.get(url, timeout=timeout, **kwargs) as resp:
                parfive.log.debug("%s request made to %s with headers=%s",
                                  resp.request_info.method,
                                  resp.request_info.url,
                                  resp.request_info.headers)
                parfive.log.debug("Response received from %s with headers=%s",
                                  resp.request_info.url,
                                  resp.headers)
                if resp.status != 200:
                    raise FailedDownload(filepath_partial, url, resp)
                else:
                    filepath, skip = get_filepath(filepath_partial(resp, url), overwrite)
                    if skip:
                        parfive.log.debug("File %s already exists and overwrite is False; skipping download.", filepath)
                        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

                    # This queue will contain the downloaded chunks and their offsets
                    # as tuples: (offset, chunk)
                    downloaded_chunk_queue = asyncio.Queue()

                    download_workers = []
                    writer = asyncio.create_task(
                        self._write_worker(downloaded_chunk_queue, file_pb, filepath))

                    if not DISABLE_RANGE and max_splits and resp.headers.get('Accept-Ranges', None) == "bytes":
                        content_length = int(resp.headers['Content-length'])
                        split_length = content_length // max_splits
                        ranges = [
                            [start, start + split_length]
                            for start in range(0, content_length, split_length)
                        ]
                        # let the last part download everything
                        ranges[-1][1] = ''
                        for _range in ranges:
                            download_workers.append(
                                asyncio.create_task(self._http_download_worker(
                                    session, url, chunksize, _range, timeout, downloaded_chunk_queue, **kwargs
                                ))
                            )
                    else:
                        download_workers.append(
                            asyncio.create_task(self._http_download_worker(
                                session, url, chunksize, None, timeout, downloaded_chunk_queue, **kwargs
                            ))
                        )

                    # run all the download workers
                    await asyncio.gather(*download_workers)
                    # join() waits till all the items in the queue have been processed
                    await downloaded_chunk_queue.join()
                    writer.cancel()
                    return str(filepath)

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

    async def _write_worker(self, queue, file_pb, filepath):
        """
        Worker for writing the downloaded chunk to the file.

        The downloaded chunk is put into a asyncio Queue by a download worker.
        This worker gets the chunk from the queue and write it to the file
        using the specified offset of the chunk.

        Parameters
        ----------
        queue: `asyncio.Queue`
             Queue for chunks

        file_pb : `tqdm.tqdm` or `False`
            Should progress bars be displayed for each file downloaded.
        filepath: `pathlib.Path`
            Path to the which the file should be downloaded.
        """
        if self.use_aiofiles:
            await self._async_write_worker(queue, file_pb, filepath)
        else:
            await self._blocking_write_worker(queue, file_pb, filepath)

    async def _async_write_worker(self, queue, file_pb, filepath):
        async with aiofiles.open(filepath, mode="wb") as f:
            while True:
                offset, chunk = await queue.get()

                await f.seek(offset)
                await f.write(chunk)
                await f.flush()

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

                queue.task_done()

    async def _blocking_write_worker(self, queue, file_pb, filepath):
        with open(filepath, "wb") as f:
            while True:
                offset, chunk = await queue.get()

                f.seek(offset)
                f.write(chunk)
                f.flush()

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

                queue.task_done()

    async def _http_download_worker(self, session, url, chunksize, http_range, timeout, queue, **kwargs):
        """
        Worker for downloading chunks from http urls.

        This function downloads the chunk from the specified http range and puts the chunk in the
        asyncio Queue. If no range is specified, then the whole file is downloaded via chunks
        and put in the queue.

        Parameters
        ----------
        session : `aiohttp.ClientSession`
            The `aiohttp.ClientSession` to use to retrieve the files.
        url : `str`
            The url to retrieve.
        chunksize : `int`
            The number of bytes to read into the file at a time.
        http_range: (`int`, `int`) or `None`
            Start and end bytes of the file. In None, then no `Range` header is specified
            in request and the whole file will be downloaded.
        queue: `asyncio.Queue`
             Queue to put the download chunks.
        kwargs : `dict`
            Extra keyword arguments are passed to `aiohttp.ClientSession.get`.
        """
        headers = kwargs.pop('headers', {})
        if http_range:
            headers['Range'] = 'bytes={}-{}'.format(*http_range)
            # init offset to start of range
            offset, _ = http_range
        else:
            offset = 0

        async with session.get(url, timeout=timeout, headers=headers, **kwargs) as resp:
            parfive.log.debug("%s request made to %s with headers=%s",
                              resp.request_info.method,
                              resp.request_info.url,
                              resp.request_info.headers)
            parfive.log.debug("Response received from %s with headers=%s",
                              resp.request_info.url,
                              resp.headers)
            while True:
                chunk = await resp.content.read(chunksize)
                if not chunk:
                    break
                await queue.put((offset, chunk))
                offset += len(chunk)

    async def _get_ftp(self, 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.Client.context`.

        Returns
        -------
        `str`
            The name of the file saved.
        """
        parse = urllib.parse.urlparse(url)
        try:
            async with aioftp.Client.context(parse.hostname, **kwargs) as client:
                parfive.log.debug("Connected to ftp server %s", parse.hostname)
                if parse.username and parse.password:
                    parfive.log.debug("Explicitly Logging in with %s:%s", parse.username, parse.password)
                    await client.login(parse.username, parse.password)

                # This has to be done before we start streaming the file:
                filepath, skip = get_filepath(filepath_partial(None, url), overwrite)
                if skip:
                    parfive.log.debug("File %s already exists and overwrite is False; skipping download.", filepath)
                    return str(filepath)

                if callable(file_pb):
                    total_size = await get_ftp_size(client, parse.path)
                    file_pb = file_pb(position=token.n, unit='B', unit_scale=True,
                                      desc=filepath.name, leave=False, total=total_size)
                else:
                    file_pb = None

                parfive.log.debug("Downloading file %s from %s", parse.path, parse.hostname)
                async with client.download_stream(parse.path) as stream:
                    downloaded_chunks_queue = asyncio.Queue()
                    download_workers = []
                    writer = asyncio.create_task(
                        self._write_worker(downloaded_chunks_queue, file_pb, filepath))

                    download_workers.append(
                        asyncio.create_task(self._ftp_download_worker(
                            stream, downloaded_chunks_queue))
                    )

                    await asyncio.gather(*download_workers)
                    await downloaded_chunks_queue.join()
                    writer.cancel()

                    return str(filepath)

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

    async def _ftp_download_worker(self, stream, queue):
        """
        Similar to `Downloader._http_download_worker`.
        See that function's documentation for more info.

        Parameters
        ----------
        stream: `aioftp.StreamIO`
            Stream of the file to be downloaded.
        queue: `asyncio.Queue`
             Queue to put the download chunks.
        """
        offset = 0
        async for chunk in stream.iter_by_block():
            # Write this chunk to the output file.
            await queue.put((offset, chunk))
            offset += len(chunk)
