# -*- coding: utf-8 -*-
'''
Tools for Web Flayer
'''
# Python
import os
import re
import sys
import time
import random
import pprint
import urllib

# 3rd party
import requests
from termcolor import colored
import psycopg2
from psycopg2.extras import Json
from bs4 import BeautifulSoup

# Internal
import flayer.event


class Output(object):
    '''
    Used for outputting data
    '''
    def __init__(self, opts):
        '''
        Initialize
        '''
        self.opts = opts

    def action(self, msg, force=False):
        '''
        Something is currently happening
        '''
        if not self.opts['daemon'] or force is True:
            print(colored(msg, self.opts.get('action_color', 'green')))

    def info(self, msg, force=False):
        '''
        Informational only
        '''
        if not self.opts['daemon'] or force is True:
            print(colored(msg, self.opts.get('info_color', 'cyan')))

    def warn(self, msg, force=False):
        '''
        Something is possibly wrong, but not enough to stop running
        '''
        if not self.opts['daemon'] or force is True:
            print(colored(msg, self.opts.get('warn_color', 'yellow')))

    def error(self, msg, force=False):
        '''
        Something is wrong enough to halt execution
        '''
        if not self.opts['daemon'] or force is True:
            print(colored(msg, self.opts.get('error_color', 'red'), attrs=['bold']))


def process_url(url_uuid, url, content, parsers):
    '''
    Process a URL
    '''
    fun = None
    for mod in parsers:
        if fun is not None:
            break
        if not mod.endswith('.func_map'):
            continue
        fun = parsers[mod](url)
    fun(url_uuid, url, content)


def get_url(
        url,
        parent=None,
        referer=None,
        dbclient=None,
        client=requests,
        opts=None,
        context=None,
    ):
    '''
    Download a URL (if necessary) and store it
    '''
    out = Output(opts)

    headers = opts['headers'].copy()
    data = opts.get('data', None)

    if referer:
        headers['referer'] = referer

    if flayer.db.check_domain_wait(dbclient, url) is False:
        # We need to put this URL back into the queue
        queue_urls([url], dbclient, opts)
    flayer.db.pattern_wait(dbclient, url)
    flayer.db.set_domain_wait(dbclient, opts, url)

    wait = 0
    if opts.get('no_db_cache') is True:
        # Skip all the DB stuff and just download the URL
        req = client.request(
            opts['method'],
            url,
            headers=headers,
            data=data,
            verify=bool(opts.get('verify', True)),
        )
        req.raise_for_status()
        if opts.get('include_headers') is True:
            out.info(pprint.pformat(dict(req.headers)))
        content = req.text
        if opts['random_wait'] is True:
            wait = int(opts.get('wait', 10))
            time.sleep(random.randrange(1, wait))
        if url not in opts['warned']:
            opts['warned'].append(url)
        return 0, content

    cur = dbclient.cursor()
    exists = False

    # Check for URL in DB
    cur.execute('''
        SELECT uuid, url, last_retrieved
        FROM urls
        WHERE url = %s
    ''', [url])

    if cur.rowcount < 1:
        # URL has never been retrieved
        cur.execute('''
            INSERT INTO urls
            (url) VALUES (%s)
            RETURNING uuid
        ''', [url])
        dbclient.commit()
        url_uuid = cur.fetchone()[0]
        out.action('{} has not been retrieved before, new UUID is {}'.format(url, url_uuid))
    else:
        # URL has been retrieved, get its UUID
        url_uuid = cur.fetchone()[0]
        out.warn('{} exists, UUID is {}'.format(url, url_uuid))
        exists = True
        if url not in opts['warned']:
            opts['warned'].append(url)

    # Save referer relationships
    if parent:
        try:
            cur.execute('''
                INSERT INTO referers
                (url_uuid, referer_uuid)
                VALUES
                (%s, %s)
            ''', [url_uuid, parent])
            dbclient.commit()
        except psycopg2.IntegrityError:
            # This relationship already exists
            dbclient.rollback()

    if opts['force_directories'] and not opts['save_path']:
        opts['save_path'] = '.'

    # Check for content
    cur.execute('''
        SELECT data, uuid
        FROM content
        WHERE url_uuid = %s
        ORDER BY retrieved
        LIMIT 1
    ''', [url_uuid])
    if cur.rowcount < 1:
        try:
            if opts['save_path']:
                req = client.request(
                    opts['method'],
                    url,
                    headers=headers,
                    data=data,
                    verify=bool(opts.get('verify', True)),
                    stream=True,
                )
                content, req_headers = _save_path(url, url_uuid, req, wait, opts, context, dbclient)
            else:
                req = client.request(
                    opts['method'],
                    url,
                    headers=headers,
                    data=data,
                    verify=bool(opts.get('verify', True)),
                )
                content = req.text
                req_headers = req.headers
        except requests.exceptions.ConnectionError as exc:
            out.error('Error downloading {}:'.format(url))
            out.error(exc)
            return 0, ''
        except requests.exceptions.InvalidSchema as exc:
            out.error('Error downloading {}:'.format(url))
            out.error(exc)
            return 0, ''
        if url not in opts['warned']:
            opts['warned'].append(url)
        if opts.get('include_headers') is True:
            out.info(pprint.pformat(dict(req_headers)))
        if content:
            cur.execute('''
                    INSERT INTO content
                    (url_uuid, data) VALUES (%s, %s)
                ''',
                [
                    url_uuid,
                    Json({
                        'content': content.replace('\x00', ''),
                        'status': req.status_code,
                    })
                ]
            )
            dbclient.commit()
    else:
        if opts['force'] is True:
            row_id = cur.fetchone()[1]
            if opts['save_path']:
                req = client.request(
                    opts['method'],
                    url,
                    headers=headers,
                    data=data,
                    verify=bool(opts.get('verify', True)),
                    stream=True,
                )
                content, req_headers = _save_path(url, url_uuid, req, wait, opts, context, dbclient)
            else:
                req = client.request(
                    opts['method'],
                    url,
                    headers=headers,
                    data=data,
                    verify=bool(opts.get('verify', True)),
                )
                content = req.text
                req_headers = req.headers
            if url not in opts['warned']:
                opts['warned'].append(url)
            if opts.get('include_headers') is True:
                out.info(pprint.pformat(dict(req_headers)))
            if content:
                cur.execute('''
                        UPDATE content
                        SET url_uuid = %s, data = %s
                        WHERE uuid = %s
                    ''',
                    [
                        url_uuid,
                        Json({'content': content}),
                        row_id
                    ]
                )
                dbclient.commit()
        else:
            content = cur.fetchone()[0]['content']

    flayer.db.pattern_wait(dbclient, url)
    flayer.db.set_domain_wait(dbclient, opts, url)

    if exists is False:
        if opts['random_wait'] is True:
            wait = int(opts.get('wait', 10))
            time.sleep(random.randrange(1, wait))
    return url_uuid, content


def _save_path(url, url_uuid, req, wait, opts, context, dbclient):
    '''
    Save the URL to a path
    '''
    urlcomps = urllib.parse.urlparse(url)
    if opts['force_directories']:
        newpath = urlcomps[2].lstrip('/')
        file_name = os.path.join(opts['save_path'], urlcomps[1], newpath)
    else:
        file_name = os.path.join(opts['save_path'], urlcomps[2].split('/')[-1])
    return status(req, url, url_uuid, file_name, wait, opts, context, dbclient)


def status(
        req,
        media_url,
        url_uuid,
        file_name,
        wait=0,
        opts=None,
        context=None,
        dbclient=None,
    ):
    '''
    Show status of the download
    '''
    out = Output(opts)

    if opts is None:
        opts = {}

    if context is None:
        context = {}

    file_name = _rename(media_url, file_name, opts)

    cache_dir = '/'.join(file_name.split('/')[:-1])
    try:
        os.makedirs(cache_dir, mode=0o0755, exist_ok=True)
    except PermissionError as exc:
        out.error('Cannot create directory {}: {}'.format(cache_dir, exc))

    is_text = False
    req_headers = req.headers
    for header in list(req_headers):
        if header.lower().startswith('content-type'):
            if req_headers[header].startswith('text'):
                is_text = True
    content = ''

    cur = dbclient.cursor()
    agent_id = opts.get('id', 'unknown')
    cur.execute(
        'INSERT INTO active_dl (url_uuid, started_by) VALUES (%s, %s)',
        [url_uuid, agent_id]
    )

    cur.execute('SELECT url FROM urls WHERE uuid = %s', [url_uuid])
    root_url = cur.fetchone()[0]

    flayer.db.pattern_wait(dbclient, media_url)
    flayer.db.set_domain_wait(dbclient, opts, media_url)

    out.action('Downloading: {}'.format(media_url))
    if os.path.exists(file_name):
        if opts['overwrite']:
            out.warn('... {} exists, overwriting'.format(file_name))
        else:
            out.warn('... {} exists, skipping'.format(file_name))
            return None, {}
    if not opts['daemon']:
        sys.stdout.write(colored('...Saving to: ', 'green'))
    out.info(file_name)
    buffer_size = 4096
    total = int(req.headers.get('Content-Length', 0))
    count = 0
    try:
        point = int(total / 100)
        #increment = int(total / buffer_size)
    except ZeroDivisionError:
        out.error('Divide by zero error, status not available')
        point = 0
        #increment = 0
    start_time = time.time()
    last_time = time.time()
    delay_blocks = 0
    delay_count = 0

    context['dl_data'] = {
        'url': root_url,
        'media_url': media_url,
        'url_uuid': url_uuid,
        'bytes_total': '',
        'bytes_elapsed': '',
        'time_total': '',
        'time_left': '',
        'kbsec': 0,
    }
    flayer.event.fire('flayer/{}/download'.format(opts['id']), {root_url: 'started'}, opts)
    try:
        with open(file_name, 'wb') as fhp:
            #old_time = time.time()
            try:
                for block in req.iter_content(buffer_size):
                    if opts.get('hard_stop'):
                        queue_urls([media_url], dbclient, opts)
                        break
                    if opts.get('abort'):
                        break
                    if is_text is True:
                        content += str(block)
                    fhp.write(block)
                    count += buffer_size
                    delay_blocks += buffer_size
                    delay_count += 1
                    #old_time = time.time()
                    time_delay = time.time() - last_time
                    if time_delay >= float(1):
                        last_time = time.time()
                        try:
                            blocks_left = int((total - count) / buffer_size)
                        except ZeroDivisionError:
                            blocks_left = 0
                        kbsec = (buffer_size / 1024) * delay_count
                        try:
                            seconds_left = ((blocks_left * buffer_size) / 1024) / kbsec
                        except ZeroDivisionError:
                            seconds_left = 0
                        minutes_left = int(seconds_left / 60)
                        minsecs_left = seconds_left % 60
                        time_left = '%d:%02d' % (minutes_left, minsecs_left)
                        seconds_elapsed = time.time() - start_time
                        seconds_total = seconds_elapsed + seconds_left
                        minutes_total = int(seconds_total / 60)
                        minsecs_total = int(seconds_total % 60)
                        time_total = '%d:%02d' % (minutes_total, minsecs_total)
                        try:
                            percent = int(count / point)
                        except ZeroDivisionError:
                            percent = 0
                        context['dl_data']['bytes_total']   = total  # pylint: disable=bad-whitespace
                        context['dl_data']['bytes_elapsed'] = count  # pylint: disable=bad-whitespace
                        context['dl_data']['time_total']    = time_total  # pylint: disable=bad-whitespace
                        context['dl_data']['time_left']     = time_left  # pylint: disable=bad-whitespace
                        context['dl_data']['kbsec']         = kbsec  # pylint: disable=bad-whitespace
                        if not opts['daemon']:
                            sys.stdout.write('\x1b[2K\r')
                            sys.stdout.write(
                                colored('Total size is {} '.format(sizeof_fmt(total)), 'green'))
                            sys.stdout.write(colored('({} bytes), '.format(total), 'green'))
                            sys.stdout.write(colored('{}%, '.format(str(percent)), 'cyan'))
                            sys.stdout.write(colored(kbsec, 'cyan'))
                            sys.stdout.write(colored(' KiB/s, ', 'cyan'))
                            sys.stdout.write(colored('{}/{} left'.format(time_left, time_total), 'cyan'))
                            sys.stdout.flush()
                        delay_blocks = 0
                        delay_count = 0
            except OSError as exc:
                out.error('OS Error: {}'.format(exc))
                out.error('Media URL: {}'.format(media_url))
            except ProtocolError as exc:
                out.error('Protocol Error: {}'.format(exc))
                out.error('Media URL: {}'.format(media_url))
            except Exception as exc:
                out.error('Exception: {}'.format(exc))
                out.error('Media URL: {}'.format(media_url))
    except OSError as exc:
        out.error('There was an error opening {}: {}'.format(file_name, exc))

    del context['dl_data']

    if opts.get('hard_stop') or opts.get('abort'):
        os.remove(file_name)

    if is_text is True and opts.get('save_html', True) is False:
        os.remove(file_name)

    if not content:
        content = None

    cur.execute('DELETE FROM active_dl WHERE url_uuid = %s', [url_uuid])
    flayer.event.fire('flayer/{}/download'.format(opts['id']), {root_url: 'complete'}, opts)

    flayer.db.pattern_wait(dbclient, media_url)
    flayer.db.set_domain_wait(dbclient, opts, media_url)

    if not opts['daemon']:
        print()
    time.sleep(wait)

    return content, req_headers


def sizeof_fmt(num, suffix='B'):
    '''
    Show human-readable sizes
    '''
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s " % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s " % (num, 'Yi', suffix)


def dbsave_media(cur, media_url, url_uuid, file_name, dbclient):
    '''
    Save a media item into the database, once it's been downloaded

    cur: Database cursor

    media_url: The URL of the image/video that was downloaded

    url_uuid: The UUID of the parent of the media_url

    file_name: The place where the media_url was downloaded to
    '''
    try:
        cur.execute('''
            INSERT INTO urls (url) values (%s) RETURNING uuid
        ''', [media_url])
        dbclient.commit()
        new_id = cur.fetchone()[0]
    except psycopg2.IntegrityError:
        # This relationship already exists
        dbclient.rollback()
        cur.execute('''
            SELECT uuid FROM urls WHERE url = %s
        ''', [media_url])
        new_id = cur.fetchone()[0]

    try:
        cur.execute('''
            INSERT INTO referers (url_uuid, referer_uuid) values (%s, %s)
        ''', [new_id, url_uuid])
        dbclient.commit()
    except psycopg2.IntegrityError:
        # This relationship already exists
        dbclient.rollback()

    cur.execute('''
        SELECT COUNT(*) FROM content WHERE url_uuid = %s
    ''', [new_id])
    if cur.fetchone()[0] < 1:
        cur.execute('''
            INSERT INTO content
            (url_uuid, cache_path)
            VALUES
            (%s, %s)
        ''', [new_id, file_name])
        dbclient.commit()


def queue_urls(links, dbclient, opts):
    '''
    Check the database for any queued URLS, and add to the list
    '''
    out = Output(opts)

    cur = dbclient.cursor()
    if isinstance(links, str):
        links = [links]
    for url in links:
        if opts.get('force') is not True and not opts.get('queue_id'):
            # Check for URL in DB
            cur.execute('''
                SELECT uuid
                FROM urls
                WHERE url = %s
            ''', [url])
            if cur.rowcount > 0:
                if url not in opts['warned']:
                    out.info('URL has already been downloaded; use --force if necessary')
                else:
                    if url not in opts['warned']:
                        opts['warned'].append(url)
                continue

        fields = ['url']
        args = [url]
        if opts.get('queue_id') is not None:
            fields.append('uuid')
            args.append(opts['queue_id'])

        if 'refresh_interval' in opts:
            fields.append('refresh_interval')
            args.append(opts['refresh_interval'])

        if 'overwrite' not in opts:
            opts['overwrite'] = False

        fields.append('overwrite')
        args.append(opts['overwrite'])

        query = 'INSERT INTO dl_queue ({}) VALUES ({})'.format(
            ', '.join(fields),
            ', '.join(['%s' for arg in range(len(args))])
        )

        try:
            cur.execute(query, args)
            dbclient.commit()
        except psycopg2.IntegrityError:
            # This URL is already queued
            dbclient.rollback()
    cur.execute('SELECT count(*) FROM dl_queue')
    return cur.fetchone()[0]


def reprocess_urls(urls, patterns, dbclient=None):
    '''
    Reprocess the cached URLs which matches the pattern(s)
    '''
    if not urls:
        urls = []

    if isinstance(patterns, str):
        patterns = [patterns]

    cur = dbclient.cursor()
    wheres = ['url~%s'] * len(patterns)
    query = 'SELECT url FROM urls WHERE {}'.format(' OR '.join(wheres))
    cur.execute(query, patterns)
    for row in cur.fetchall():
        urls.append(row[0])

    return urls


def queue_regexp(urls, pattern, dbclient, opts):
    '''
    Add the URLs matching the pattern to the download queue
    '''
    expr = re.compile(pattern)
    links = []
    for url in urls:
        if expr.search(url):
            links.append(url)
    queue_urls(links, dbclient, opts)


def _rename(media_url, file_name, opts):
    '''
    When files are downloaded using status, rename as per a template
    '''
    out = Output(opts)

    template = opts.get('rename_template', '')
    if not template:
        return file_name

    urlcomps = urllib.parse.urlparse(media_url)
    replacements = {
        'host': urlcomps[1].split(':')[0],
        'path': '/'.join(urlcomps[2].split('/')[:-2])
    }

    # File extensions
    if '.' in urlcomps[2].split('/')[-1]:
        replacements['ext'] = urlcomps[2].split('/')[-1].split('.')[-1]
    else:
        replacements['ext'] = ''

    if not opts.get('rename_count'):
        opts['rename_count'] = opts.get('rename_count_start', 0)

    if opts.get('rename_count_padding'):
        try:
            opts['rename_count_padding'] = int(opts['rename_count_padding'])
        except ValueError:
            out.warn('--rename-count-padding must be an integer, using 0')
            opts['rename_count_padding'] = 0
        template = template.replace('{count}', '{count:0>{rename_count_padding}}')
        replacements['rename_count_padding'] = opts['rename_count_padding']
    replacements['count'] = str(opts['rename_count'])
    opts['rename_count'] += 1

    file_name = os.path.join(opts['save_path'], template.format(**replacements))

    return file_name


def parse_links(url, content, level, opts):
    '''
    Return the links from an HTML page
    '''
    out = Output(opts)

    hrefs = []
    try:
        # Get ready to do some html parsing
        soup = BeautifulSoup(content, 'html.parser')
        # Generate absolute URLs for every link on the page
        url_comps = urllib.parse.urlparse(url)
        tags = soup.find_all('a')
        if opts['search_src'] is True:
            tags = tags + soup.find_all(src=True)
        for link in tags:
            if level > int(opts['level']):
                continue
            href = urllib.parse.urljoin(url, link.get('href'))
            if opts['search_src'] is True and not link.get('href'):
                href = urllib.parse.urljoin(url, link.get('src'))
            link_comps = urllib.parse.urlparse(href)
            if link.text.startswith('javascript'):
                continue
            if int(opts.get('level', 0)) > 0 and int(opts.get('level', 0)) < 2:
                continue
            if opts['span_hosts'] is not True:
                if not link_comps[1].startswith(url_comps[1].split(':')[0]):
                    continue
            hrefs.append(href.split('#')[0])
        # Render the page, and print it along with the links
        if opts.get('render', False) is True:
            out.info(soup.get_text())
        return hrefs
    except TypeError:
        # This URL probably isn't HTML
        return []
