#!/usr/bin/env python3
"""get tag from http://demo.illustration2vec.net/."""
# note:
# - error 'ERROR: Request Entity Too Large' for file 1.1 mb
# <span style="color:red;">ERROR: Request Entity Too Large</span>
from collections import OrderedDict
from pathlib import Path
from pprint import pformat
import imghdr
import logging
import os
import shutil
import time
import urllib
import hashlib

import click
import requests
import structlog
import peewee
from PIL import Image

from i2vec_cli import models
from i2vec_cli.requests_session import Session, convert_raw_to_hydrus
from i2vec_cli.sha256 import sha256_checksum
from i2vec_cli.utils import user_data_dir, thumb_folder


def is_url(path):
    """Return True if path is url, False otherwise."""
    scheme = urllib.parse.urlparse(path).scheme
    if scheme in ('http', 'https'):
        return True
    return False


def is_ext_equal(file_ext, imghdr_ext):
    """compare file extension with result from imghdr_ext."""
    if not imghdr_ext:
        return False
    if file_ext.lower() == '.{}'.format(imghdr_ext):
        return True
    if file_ext.lower() in ('.jpg', '.jpeg') and imghdr_ext == 'jpeg':
        return True
    return False


def download(url, no_clobber):
    """download url.

    Args:
        url: URL to be downloaded.
        no_clobber: Skip download if file already exist.

    Returns:
        Downloaded filename or existing file if `no_clobber` is `True`
    """
    log = structlog.getLogger()

    basename = os.path.basename(url)
    if os.path.isfile(basename) and no_clobber:
        return basename

    response = requests.get(url, stream=True)
    with open(basename, 'wb') as out_file:
        shutil.copyfileobj(response.raw, out_file)
    name, ext = os.path.splitext(basename)
    imghdr_ext = imghdr.what(basename)
    ext_equal = is_ext_equal(file_ext=ext, imghdr_ext=imghdr_ext)

    if not imghdr_ext:
        log.debug("imghdr can't recognize file", file=basename)
        return basename
    else:
        new_basename = '{}.{}'.format(name, imghdr_ext)
        new_basename_exist = os.path.isfile(new_basename)

    if ext_equal:
        log.debug('Extension is equal', file_ext=ext, imghdr_ext=imghdr_ext)
        return basename
    elif not ext_equal:
        if new_basename_exist and not no_clobber:
            log.debug('Replace existing file', old=basename, new=new_basename)
            shutil.move(basename, new_basename)
        elif not new_basename_exist:
            log.debug('Rename file ext', file=basename, new_ext=imghdr_ext)
            shutil.move(basename, new_basename)
        else:
            log.debug('Not replace/rename file', no_clobber=no_clobber, new_basename=new_basename)
        return new_basename
    else:
        log.debug(
            'Unknown condition',
            file=basename,
            ext_equal=ext_equal,
            new_basename_exist=new_basename_exist,
            imghdr_ext=imghdr_ext
        )
    # just return base name if any error happen
    return basename


def validate_close_delay(ctx, param, value):
    """validate close delay."""
    try:
        value = int(value)
    except Exception as e:
        raise click.BadParameter(
            'Error when validate close delay: value={}, error={}'.format(value, e))
    if value >= -1:
        return value
    else:
        raise click.BadParameter('Close delay have to be bigger or equal than -1')


def delay_close(close_delay):
    """delay when closing the program."""
    log = structlog.getLogger()
    if close_delay == -1:
        click.pause()
    elif close_delay == 0:
        log.debug('No close delay')
    elif close_delay > 0:
        time.sleep(close_delay)
    else:
        log.error('Invalid close delay', v=close_delay)


def md5_checksum(fname):
    hash_md5 = hashlib.md5()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def create_thumbnail(path, thumb_path):
    """create thumbnail."""
    size = 320, 320
    try:
        im = Image.open(path)
        im.thumbnail(size)
        im.save(thumb_path, "JPEG")
    except IOError:
        raise IOError("cannot create thumbnail for", path)


def get_print_result(path, db_path, format, session):
    """get print result."""
    # compatibility
    p = path

    sha256 = sha256_checksum(p)
    md5 = md5_checksum(p)
    thumb_path = os.path.join(user_data_dir, 'thumb', '{}.jpg'.format(sha256))
    try:
        load_res = models.load_result(db=db_path, sha256=sha256, md5=md5)
    except models.Image.DoesNotExist:
        load_res = None
    if load_res:
        tags = {'prediction': load_res}
    else:
        tags = session.get_tags(path=p)
        try:
            models.save_result(
                db=db_path, sha256=sha256, md5=md5, prediction=tags['prediction'])
        except peewee.IntegrityError as e:
            log.debug(str(e))
        except keyError as e:
            log.debug(str(tags))
    if not os.path.isfile(thumb_path):
        create_thumbnail(p, thumb_path)

    if format == 'dict':
        return tags
    if format == 'hydrus':
        return convert_raw_to_hydrus(tags)
    else:
        return pformat(tags['prediction'])


@click.command()
@click.option('--format', type=click.Choice(['raw', 'hydrus']), default='raw')
@click.option('-d', '--debug', is_flag=True, help="Enable debug.")
@click.option('-nc', '--no-clobber', is_flag=True, help="Skip download url when file exist.")
@click.option(
    '--close-delay', default=0, help="Close delay of the program.", callback=validate_close_delay)
@click.option(
    '--driver', default=None, help="Driver for browser (deprecated).",
    type=click.Choice(['firefox', 'phantomjs', 'chrome', 'zope.testbrowser', 'django']))
@click.option('--dump-html', is_flag=True, help="Dump html table for debugging (deprecated).")
@click.argument('path', nargs=-1)
def main(format, path, debug, no_clobber, close_delay, driver=None, dump_html=False):
    """get tag from illustration2vec."""
    if debug:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
    structlog.configure_once(logger_factory=structlog.stdlib.LoggerFactory())
    log = structlog.getLogger()

    if not path:
        raise ValueError('PATH required.')

    # init folder
    os.makedirs(user_data_dir, exist_ok=True)
    os.makedirs(thumb_folder, exist_ok=True)

    # database
    db_path = os.path.join(user_data_dir, 'main.db')
    if not os.path.isfile(db_path):
        Path(db_path).touch()
    models.database.init(db_path)
    try:
        models.init_all_tables()
    except peewee.OperationalError:
        log.debug('Table already created')

    session = Session(driver=driver)
    try:
        for p in path:
            if os.path.isfile(p):
                print('path:{}'.format(os.path.basename(p)))
            elif is_url(p):
                print('url:{}'.format(p))
                p = download(p, no_clobber=no_clobber)
            else:
                log.error('Unknown path format or path is not exist', path=p)
                continue
            result = get_print_result(
                path=p, db_path=db_path, format=format, session=session)
            print(result)
    finally:
        delay_close(close_delay)
        if hasattr(session, 'browser'):
            session.browser.quit()


if __name__ == '__main__':
    main()

[((182, 1, 182, 16), 'click.command', 'click.command', ({}, {}), '()', False, 'import click\n'), ((184, 1, 184, 66), 'click.option', 'click.option', (), '', False, 'import click\n'), ((185, 1, 185, 93), 'click.option', 'click.option', (), '', False, 'import click\n'), ((186, 1, 187, 98), 'click.option', 'click.option', (), '', False, 'import click\n'), ((191, 1, 191, 94), 'click.option', 'click.option', (), '', False, 'import click\n'), ((192, 1, 192, 33), 'click.argument', 'click.argument', (), '', False, 'import click\n'), ((58, 10, 58, 31), 'structlog.getLogger', 'structlog.getLogger', ({}, {}), '()', False, 'import structlog\n'), ((60, 15, 60, 36), 'os.path.basename', 'os.path.basename', ({(60, 32, 60, 35): 'url'}, {}), '(url)', False, 'import os\n'), ((64, 15, 64, 45), 'requests.get', 'requests.get', (), '', False, 'import requests\n'), ((67, 16, 67, 42), 'os.path.splitext', 'os.path.splitext', ({(67, 33, 67, 41): 'basename'}, {}), '(basename)', False, 'import os\n'), ((68, 17, 68, 38), 'imghdr.what', 'imghdr.what', ({(68, 29, 68, 37): 'basename'}, {}), '(basename)', False, 'import imghdr\n'), ((118, 10, 118, 31), 'structlog.getLogger', 'structlog.getLogger', ({}, {}), '()', False, 'import structlog\n'), ((130, 15, 130, 28), 'hashlib.md5', 'hashlib.md5', ({}, {}), '()', False, 'import hashlib\n'), ((153, 13, 153, 31), 'i2vec_cli.sha256.sha256_checksum', 'sha256_checksum', ({(153, 29, 153, 30): 'p'}, {}), '(p)', False, 'from i2vec_cli.sha256 import sha256_checksum\n'), ((200, 10, 200, 31), 'structlog.getLogger', 'structlog.getLogger', ({}, {}), '()', False, 'import structlog\n'), ((206, 4, 206, 45), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((207, 4, 207, 44), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((210, 14, 210, 52), 'os.path.join', 'os.path.join', ({(210, 27, 210, 40): 'user_data_dir', (210, 42, 210, 51): '"""main.db"""'}, {}), "(user_data_dir, 'main.db')", False, 'import os\n'), ((213, 4, 213, 33), 'i2vec_cli.models.database.init', 'models.database.init', ({(213, 25, 213, 32): 'db_path'}, {}), '(db_path)', False, 'from i2vec_cli import models\n'), ((219, 14, 219, 36), 'i2vec_cli.requests_session.Session', 'Session', (), '', False, 'from i2vec_cli.requests_session import Session, convert_raw_to_hydrus\n'), ((31, 13, 31, 40), 'urllib.parse.urlparse', 'urllib.parse.urlparse', ({(31, 35, 31, 39): 'path'}, {}), '(path)', False, 'import urllib\n'), ((61, 7, 61, 31), 'os.path.isfile', 'os.path.isfile', ({(61, 22, 61, 30): 'basename'}, {}), '(basename)', False, 'import os\n'), ((66, 8, 66, 50), 'shutil.copyfileobj', 'shutil.copyfileobj', ({(66, 27, 66, 39): 'response.raw', (66, 41, 66, 49): 'out_file'}, {}), '(response.raw, out_file)', False, 'import shutil\n'), ((76, 29, 76, 57), 'os.path.isfile', 'os.path.isfile', ({(76, 44, 76, 56): 'new_basename'}, {}), '(new_basename)', False, 'import os\n'), ((113, 14, 113, 82), 'click.BadParameter', 'click.BadParameter', ({(113, 33, 113, 81): '"""Close delay have to be bigger or equal than -1"""'}, {}), "('Close delay have to be bigger or equal than -1')", False, 'import click\n'), ((120, 8, 120, 21), 'click.pause', 'click.pause', ({}, {}), '()', False, 'import click\n'), ((141, 13, 141, 29), 'PIL.Image.open', 'Image.open', ({(141, 24, 141, 28): 'path'}, {}), '(path)', False, 'from PIL import Image\n'), ((157, 19, 157, 73), 'i2vec_cli.models.load_result', 'models.load_result', (), '', False, 'from i2vec_cli import models\n'), ((171, 11, 171, 37), 'os.path.isfile', 'os.path.isfile', ({(171, 26, 171, 36): 'thumb_path'}, {}), '(thumb_path)', False, 'import os\n'), ((177, 15, 177, 42), 'i2vec_cli.requests_session.convert_raw_to_hydrus', 'convert_raw_to_hydrus', ({(177, 37, 177, 41): 'tags'}, {}), '(tags)', False, 'from i2vec_cli.requests_session import Session, convert_raw_to_hydrus\n'), ((179, 15, 179, 42), 'pprint.pformat', 'pformat', ({(179, 23, 179, 41): "tags['prediction']"}, {}), "(tags['prediction'])", False, 'from pprint import pformat\n'), ((196, 8, 196, 48), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((198, 8, 198, 47), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((211, 11, 211, 34), 'os.path.isfile', 'os.path.isfile', ({(211, 26, 211, 33): 'db_path'}, {}), '(db_path)', False, 'import os\n'), ((215, 8, 215, 32), 'i2vec_cli.models.init_all_tables', 'models.init_all_tables', ({}, {}), '()', False, 'from i2vec_cli import models\n'), ((183, 31, 183, 62), 'click.Choice', 'click.Choice', ({(183, 44, 183, 61): "['raw', 'hydrus']"}, {}), "(['raw', 'hydrus'])", False, 'import click\n'), ((190, 9, 190, 87), 'click.Choice', 'click.Choice', ({(190, 22, 190, 86): "['firefox', 'phantomjs', 'chrome', 'zope.testbrowser', 'django']"}, {}), "(['firefox', 'phantomjs', 'chrome', 'zope.testbrowser', 'django'])", False, 'import click\n'), ((165, 12, 166, 82), 'i2vec_cli.models.save_result', 'models.save_result', (), '', False, 'from i2vec_cli import models\n'), ((199, 44, 199, 76), 'structlog.stdlib.LoggerFactory', 'structlog.stdlib.LoggerFactory', ({}, {}), '()', False, 'import structlog\n'), ((222, 15, 222, 32), 'os.path.isfile', 'os.path.isfile', ({(222, 30, 222, 31): 'p'}, {}), '(p)', False, 'import os\n'), ((84, 12, 84, 47), 'shutil.move', 'shutil.move', ({(84, 24, 84, 32): 'basename', (84, 34, 84, 46): 'new_basename'}, {}), '(basename, new_basename)', False, 'import shutil\n'), ((124, 8, 124, 31), 'time.sleep', 'time.sleep', ({(124, 19, 124, 30): 'close_delay'}, {}), '(close_delay)', False, 'import time\n'), ((212, 8, 212, 21), 'pathlib.Path', 'Path', ({(212, 13, 212, 20): 'db_path'}, {}), '(db_path)', False, 'from pathlib import Path\n'), ((87, 12, 87, 47), 'shutil.move', 'shutil.move', ({(87, 24, 87, 32): 'basename', (87, 34, 87, 46): 'new_basename'}, {}), '(basename, new_basename)', False, 'import shutil\n'), ((223, 39, 223, 58), 'os.path.basename', 'os.path.basename', ({(223, 56, 223, 57): 'p'}, {}), '(p)', False, 'import os\n')]