import argparse
import os
import requests
import sys
import urllib

from voussoirkit import bytestring
from voussoirkit import dotdict
from voussoirkit import pathclass
from voussoirkit import pipeable
from voussoirkit import ratelimiter
from voussoirkit import vlogging

log = vlogging.getLogger(__name__, 'downloady')

USERAGENT = '''
'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko)
Chrome/42.0.2311.152 Safari/537.36'
'''.strip().replace('\n', ' ')

HEADERS = {
    'User-Agent': USERAGENT,
}

FILENAME_BADCHARS = '*?"<>|\r\n'

CHUNKSIZE = 4 * bytestring.KIBIBYTE
TIMEOUT = 60
TEMP_EXTENSION = '.downloadytemp'

if os.name == 'nt':
    SPECIAL_FILENAMES = pathclass.WINDOWS_RESERVED_NAMES
else:
    SPECIAL_FILENAMES = [os.devnull]
SPECIAL_FILENAMES = [os.path.normcase(x) for x in SPECIAL_FILENAMES]

class DownloadyException(Exception):
    pass

class NotEnoughBytes(DownloadyException):
    pass

class ServerNoRange(DownloadyException):
    pass

class SpecialPath:
    def __init__(self, path):
        self.absolute_path = path

    def open(self, *args, **kwargs):
        return open(self.absolute_path, *args, **kwargs)

def download_file(
        url,
        localname=None,
        auth=None,
        bytespersecond=None,
        callback_progress=None,
        do_head=True,
        headers=None,
        overwrite=False,
        raise_for_undersized=True,
        timeout=None,
        verbose=False,
        verify_ssl=True,
    ):
    plan = prepare_plan(
        url,
        localname,
        auth=auth,
        bytespersecond=bytespersecond,
        callback_progress=callback_progress,
        do_head=do_head,
        headers=headers,
        overwrite=overwrite,
        raise_for_undersized=raise_for_undersized,
        timeout=timeout,
        verify_ssl=verify_ssl,
    )

    if plan is None:
        return

    return download_plan(plan)

def download_plan(plan):
    if not isinstance(plan.download_into, SpecialPath):
        plan.download_into.parent.makedirs(exist_ok=True)
        plan.download_into.touch()

    if plan.plan_type in ['resume', 'partial']:
        file_handle = plan.download_into.open('r+b')
        file_handle.seek(plan.seek_to)
        bytes_downloaded = plan.seek_to

    elif plan.plan_type == 'fulldownload':
        file_handle = plan.download_into.open('wb')
        bytes_downloaded = 0

    if plan.header_range_min is not None:
        plan.headers['range'] = 'bytes={min}-{max}'.format(
            min=plan.header_range_min,
            max=plan.header_range_max,
        )

    log.info('Downloading %s into "%s"', plan.url, plan.real_localname.absolute_path)

    download_stream = request(
        'get',
        plan.url,
        stream=True,
        auth=plan.auth,
        headers=plan.headers,
        timeout=plan.timeout,
        verify_ssl=plan.verify_ssl,
    )

    if plan.remote_total_bytes is None:
        # Since we didn't do a head, let's fill this in now.
        plan.remote_total_bytes = int(download_stream.headers.get('Content-Length', 0))

    callback_progress = plan.callback_progress
    if callback_progress is not None:
        callback_progress = callback_progress(plan.remote_total_bytes)

    for chunk in download_stream.iter_content(chunk_size=CHUNKSIZE):
        bytes_downloaded += len(chunk)
        file_handle.write(chunk)
        if callback_progress is not None:
            callback_progress.step(bytes_downloaded)

        if plan.limiter is not None and bytes_downloaded < plan.remote_total_bytes:
            plan.limiter.limit(len(chunk))

    file_handle.close()

    # Don't try to rename /dev/null or other special names
    if isinstance(plan.real_localname, SpecialPath):
        return plan.real_localname

    temp_localsize = plan.download_into.size
    undersized = plan.plan_type != 'partial' and temp_localsize < plan.remote_total_bytes
    if undersized and plan.raise_for_undersized:
        message = 'File does not contain expected number of bytes. Received {size} / {total}'
        message = message.format(size=temp_localsize, total=plan.remote_total_bytes)
        raise NotEnoughBytes(message)

    if plan.download_into != plan.real_localname:
        os.rename(plan.download_into.absolute_path, plan.real_localname.absolute_path)

    return plan.real_localname

def prepare_plan(
        url,
        localname,
        auth=None,
        bytespersecond=None,
        callback_progress=None,
        do_head=True,
        headers=None,
        overwrite=False,
        raise_for_undersized=True,
        timeout=TIMEOUT,
        verify_ssl=True,
    ):
    # Chapter 1: File existence
    headers = headers or {}
    user_provided_range = 'range' in headers

    url = sanitize_url(url)
    if localname in [None, '']:
        localname = basename_from_url(url)

    if is_special_file(localname):
        real_localname = SpecialPath(localname)
        temp_localname = SpecialPath(localname)
        real_exists = False
        temp_exists = False
    else:
        localname = pathclass.Path(localname)
        if localname.is_dir:
            localname = localname.with_child(basename_from_url(url))
        localname = sanitize_filename(localname.absolute_path)
        real_localname = pathclass.Path(localname)
        temp_localname = real_localname.add_extension(TEMP_EXTENSION)
        real_exists = real_localname.exists
        temp_exists = temp_localname.exists

    if real_exists and overwrite is False and not user_provided_range:
        log.debug('File exists and overwrite is off. Nothing to do.')
        return None

    if isinstance(real_localname, SpecialPath):
        temp_localsize = 0
    else:
        temp_localsize = int(temp_exists and temp_localname.size)

    # Chapter 2: Ratelimiting
    if bytespersecond is None:
        limiter = None
    elif isinstance(bytespersecond, ratelimiter.Ratelimiter):
        limiter = bytespersecond
    else:
        limiter = ratelimiter.Ratelimiter(allowance=bytespersecond)

    # Chapter 3: Extracting range
    if user_provided_range:
        user_range_min = int(headers['range'].split('bytes=')[1].split('-')[0])
        user_range_max = headers['range'].split('-')[1]
        if user_range_max != '':
            user_range_max = int(user_range_max)
    else:
        user_range_min = None
        user_range_max = None

    # Chapter 4: Server range support
    # Always include a range on the first request to figure out whether the
    # server supports it. Use 0- to get correct remote_total_bytes
    if user_provided_range and not do_head:
        raise DownloadyException('Cannot determine range support without the head request')

    temp_headers = headers.copy()
    temp_headers.update({'range': 'bytes=0-'})

    if do_head:
        # I'm using a GET instead of an actual HEAD here because some servers respond
        # differently, even though they're not supposed to.
        head = request('get', url, stream=True, headers=temp_headers, auth=auth)
        remote_total_bytes = int(head.headers.get('content-length', 0))
        server_respects_range = (head.status_code == 206 and 'content-range' in head.headers)
        head.connection.close()
    else:
        remote_total_bytes = None
        server_respects_range = False

    if user_provided_range and not server_respects_range:
        raise ServerNoRange('Server did not respect your range header')

    # Chapter 5: Plan definitions
    plan_base = {
        'url': url,
        'auth': auth,
        'callback_progress': callback_progress,
        'limiter': limiter,
        'headers': headers,
        'real_localname': real_localname,
        'raise_for_undersized': raise_for_undersized,
        'remote_total_bytes': remote_total_bytes,
        'timeout': timeout,
        'verify_ssl': verify_ssl,
    }
    plan_fulldownload = dotdict.DotDict(
        plan_base,
        download_into=temp_localname,
        header_range_min=None,
        header_range_max=None,
        plan_type='fulldownload',
        seek_to=0,
    )
    plan_resume = dotdict.DotDict(
        plan_base,
        download_into=temp_localname,
        header_range_min=temp_localsize,
        header_range_max='',
        plan_type='resume',
        seek_to=temp_localsize,
    )
    plan_partial = dotdict.DotDict(
        plan_base,
        download_into=real_localname,
        header_range_min=user_range_min,
        header_range_max=user_range_max,
        plan_type='partial',
        seek_to=user_range_min,
    )

    # Chapter 6: Redeem your meal vouchers here
    if real_exists:
        if overwrite:
            os.remove(real_localname.absolute_path)

        if user_provided_range:
            return plan_partial

        return plan_fulldownload

    elif temp_exists and temp_localsize > 0:
        if overwrite:
            return plan_fulldownload

        if user_provided_range:
            return plan_partial

        if server_respects_range:
            log.info('Resume from byte %d' % plan_resume.seek_to)
            return plan_resume

    else:
        if user_provided_range:
            return plan_partial

        return plan_fulldownload

    raise DownloadyException('No plan was chosen?')

class Progress1:
    def __init__(self, total_bytes):
        self.limiter = ratelimiter.Ratelimiter(allowance=8, mode='reject')
        self.limiter.balance = 1
        self.total_bytes = max(1, total_bytes)
        self.divisor = bytestring.get_appropriate_divisor(total_bytes)
        self.total_format = bytestring.bytestring(total_bytes, force_unit=self.divisor)
        self.downloaded_format = '{:>%d}' % len(self.total_format)
        self.blank_char = ' '
        self.solid_char = '█'

    def step(self, bytes_downloaded):
        percent = bytes_downloaded / self.total_bytes
        percent = min(1.00, percent)
        if self.limiter.limit(1) is False and percent < 1.00:
            return

        downloaded_string = bytestring.bytestring(bytes_downloaded, force_unit=self.divisor)
        downloaded_string = self.downloaded_format.format(downloaded_string)
        block_count = 50
        solid_blocks = self.solid_char * int(block_count * percent)
        statusbar = solid_blocks.ljust(block_count, self.blank_char)
        statusbar = self.solid_char + statusbar + self.solid_char

        end = '\n' if percent == 1 else ''
        message = '\r{bytes_downloaded} {statusbar} {total_bytes}'
        message = message.format(
            bytes_downloaded=downloaded_string,
            total_bytes=self.total_format,
            statusbar=statusbar,
        )
        pipeable.stderr(message, end=end)

class Progress2:
    def __init__(self, total_bytes):
        self.total_bytes = max(1, total_bytes)
        self.limiter = ratelimiter.Ratelimiter(allowance=8, mode='reject')
        self.limiter.balance = 1
        self.total_bytes_string = '{:,}'.format(self.total_bytes)
        self.bytes_downloaded_string = '{:%d,}' % len(self.total_bytes_string)

    def step(self, bytes_downloaded):
        percent = bytes_downloaded / self.total_bytes
        percent = min(1.00, percent)
        if self.limiter.limit(1) is False and percent < 1.00:
            return

        percent *= 100
        percent_string = f'{percent:08.4f}'
        bytes_downloaded_string = self.bytes_downloaded_string.format(bytes_downloaded)

        end = '\n' if percent == 100 else ''
        message = '\r{bytes_downloaded} / {total_bytes} / {percent}%'
        message = message.format(
            bytes_downloaded=bytes_downloaded_string,
            total_bytes=self.total_bytes_string,
            percent=percent_string,
        )
        pipeable.stderr(message, end=end)

def basename_from_url(url):
    '''
    Determine the local filename appropriate for a URL.
    '''
    localname = urllib.parse.unquote(url)
    localname = localname.rstrip('/')
    localname = localname.split('?')[0]
    localname = localname.rsplit('/', 1)[-1]
    return localname

def is_special_file(file):
    if isinstance(file, pathclass.Path):
        return False
    file = pathclass.normalize_sep(file)
    file = file.rsplit(os.sep)[-1]
    file = os.path.normcase(file)
    return file in SPECIAL_FILENAMES

def request(method, url, headers=None, timeout=TIMEOUT, verify_ssl=True, **kwargs):
    if headers is None:
        headers = {}
    else:
        headers = headers.copy()

    for (key, value) in HEADERS.items():
        headers.setdefault(key, value)

    session = requests.Session()
    a = requests.adapters.HTTPAdapter(max_retries=30)
    b = requests.adapters.HTTPAdapter(max_retries=30)
    session.mount('http://', a)
    session.mount('https://', b)
    session.max_redirects = 40

    method = {
        'get': session.get,
        'head': session.head,
        'post': session.post,
    }[method]

    req = method(url, headers=headers, timeout=timeout, verify=verify_ssl, **kwargs)
    req.raise_for_status()
    return req

def sanitize_filename(text, exclusions=''):
    to_remove = FILENAME_BADCHARS
    for exclude in exclusions:
        to_remove = to_remove.replace(exclude, '')

    for char in to_remove:
        text = text.replace(char, '')

    (drive, path) = os.path.splitdrive(text)
    path = path.replace(':', '')
    text = drive + path

    return text

def sanitize_url(url):
    url = url.replace('%3A//', '://')
    return url

def download_argparse(args):
    url = args.url

    url = pipeable.input(url, split_lines=False)
    callback = {
        None: Progress1,
        '1': Progress1,
        '2': Progress2,
    }.get(args.callback, args.callback)

    bytespersecond = args.bytespersecond
    if bytespersecond is not None:
        bytespersecond = bytestring.parsebytes(bytespersecond)

    headers = {}
    if args.range is not None:
        headers['range'] = 'bytes=%s' % args.range

    retry = args.retry
    if not retry:
        retry = 1

    while retry != 0:
        # Negative numbers permit infinite retries.
        try:
            download_file(
                url=url,
                localname=args.localname,
                bytespersecond=bytespersecond,
                callback_progress=callback,
                do_head=args.no_head is False,
                headers=headers,
                overwrite=args.overwrite,
                timeout=args.timeout,
                verbose=True,
                verify_ssl=args.no_ssl is False,
            )
        except (NotEnoughBytes, requests.exceptions.ConnectionError, requests.exceptions.ChunkedEncodingError):
            retry -= 1
            if retry == 0:
                raise
        else:
            break

@vlogging.main_decorator
def main(argv):
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('url')
    parser.add_argument('localname', nargs='?', default=None)
    parser.add_argument('-c', '--callback', dest='callback', default=Progress1)
    parser.add_argument('-bps', '--bytespersecond', dest='bytespersecond', default=None)
    parser.add_argument('-ow', '--overwrite', dest='overwrite', action='store_true')
    parser.add_argument('-r', '--range', dest='range', default=None)
    parser.add_argument('--timeout', dest='timeout', type=int, default=TIMEOUT)
    parser.add_argument('--retry', dest='retry', nargs='?', type=int, default=1)
    parser.add_argument('--no-head', dest='no_head', action='store_true')
    parser.add_argument('--no-ssl', dest='no_ssl', action='store_true')
    parser.set_defaults(func=download_argparse)

    args = parser.parse_args(argv)
    return args.func(args)

if __name__ == '__main__':
    raise SystemExit(main(sys.argv[1:]))
