#!/usr/bin/env python
from builtins import input

import os
import json
import click
import itertools
from glob import glob
from tqdm import tqdm
from urllib.parse import urlparse
from concurrent.futures import ProcessPoolExecutor

from abraia.utils import load_image, save_image
from abraia.editing import clean_image, blur_background

from abraia import config
from abraia import Abraia
from abraia import APIError
from abraia import __version__


abraia = Abraia()


def process_map(task, *values, desc='', max_workers=3):
    with ProcessPoolExecutor(max_workers) as exe:
        with tqdm(total=len(values[0]), desc=desc) as pbar:
            for result in exe.map(task, *values):
                pbar.set_postfix_str(result)
                pbar.update(1)


def is_url(url):
    res = urlparse(url)
    if res.scheme and res.netloc:
        return True
    return False


def echo_error(error):
    click.echo('[' + click.style('Error {}'.format(error.code),
                                 fg='red', bold=True) + '] {}'.format(error.message))


def input_files(src):
    if isinstance(src, str) and src.startswith('http'):
        return [src]
    src = os.path.join(src, '**/*') if os.path.isdir(src) else src
    return glob(src, recursive=True)


@click.group('abraia')
@click.version_option(__version__)
def cli():
    """Abraia CLI tool"""
    pass


@cli.command()
def configure():
    """Configure the abraia api key."""
    # click.echo('Go to [' + click.style('https://abraia.me/console/', fg='green') + '] to see your account information\n')
    click.echo('Go to [' + click.style('https://abraia.me/console/', fg='green') + '] to get your user id and key\n')
    try:
        abraia_id, abraia_key = config.load()
        abraia_id = click.prompt('Abraia Id', default=abraia_id)
        abraia_key = click.prompt('Abraia Key', default=abraia_key)
        config.save(abraia_id, abraia_key)
    except:
        pass


def convert_file(src, dest, args):
    if src.startswith('http'):
        path, args['url'] = src, src
    else:
        path = abraia.upload_file(src, 'batch/')
    abraia.transform_image(path, dest, args)
    return dest


def editing_file(src, dest, mode):
    if mode == 'blur':
        img = load_image(src)
        out = blur_background(img)
        save_image(out, dest)
        return dest
    output = f"export/{os.path.basename(dest)}"
    path = abraia.upload_file(src, 'batch/')
    if mode == 'anonymize':
        abraia.anonymize_image(path, output)
    elif mode == 'upscale':
        abraia.upscale_image(path, output)
    else:
        abraia.remove_background(path, output)
    return abraia.download_file(output, dest)


def editing_files(src, mode, desc="Editing"):
    try:
        inputs = input_files(src)
        ext = 'png' if mode == 'removebg' else 'jpg'
        dirname = src if os.path.isdir(src) else os.path.dirname(src)
        paths = [os.path.splitext(os.path.relpath(src, dirname))[0] for src in inputs]
        outputs = [os.path.join('output', f"{path}.{ext}") for path in paths]
        process_map(editing_file, inputs, outputs, itertools.repeat(mode), desc=desc)
    except APIError as error:
        echo_error(error)


@cli.group('editing')
def cli_editing():
    """Convert and edit images in bulk."""
    pass


@cli_editing.command()
@click.option('--width', help='Resize to specified width', type=int)
@click.option('--height', help='Resize to specified height', type=int)
@click.option('--mode', help='Select the resize mode', type=click.Choice(['pad', 'crop', 'thumb']))
@click.option('--format', help='Convert to specified image format', type=click.Choice(['jpeg', 'png', 'webp']))
@click.argument('src')
def convert(src, width, height, mode, format):
    """Convert, resize, and optimize images in bulk."""
    try:
        args = {'width': width, 'height': height, 'mode': mode, 'format': format, 'quality': 'auto'}
        filenames = input_files(src)
        dests = []
        dirname = src if os.path.isdir(src) else os.path.dirname(src)
        for filename in filenames:
            if filename.startswith('http'):
                filename = os.path.basename(urlparse(filename).path) or 'screenshot.png'
            path, ext = os.path.splitext(os.path.relpath(filename, dirname))
            oext = format if format is not None else (ext and ext[1:])
            dests.append(os.path.join('output', f"{path}.{oext}"))
        process_map(convert_file, filenames, dests, itertools.repeat(args), desc="Converting")
    except APIError as error:
        echo_error(error)


@cli_editing.command()
@click.argument('src')
def removebg(src):
    """Remove the images background to make them transparent."""
    editing_files(src, 'removebg', desc="Removing")


@cli_editing.command()
@click.argument('src')
def upscale(src):
    """Upscale and enhance images increasing the resolution."""
    editing_files(src, 'upscale', desc="Upscaling")


@cli_editing.command()
@click.argument('src')
def anonymize(src):
    """Anonymize image blurring faces and car license plates."""
    editing_files(src, 'anonymize', desc="Anonymizing")


@cli_editing.command()
@click.argument('src')
def blur(src):
    """Blur the image background to focus attention on the main object."""
    editing_files(src, 'blur', desc="Blur background")


@cli_editing.command()
@click.argument('src')
def clean(src):
    """Clean images removing unwanted objects with inpainting."""
    img = load_image(src)
    clean_image(img)


@cli.group('files')
def cli_files():
    """Commands related to cloud storage."""
    pass


def upload_file(file, folder):
    return abraia.upload_file(file, folder)


def download_file(path, folder):
    dest = os.path.join(folder, os.path.basename(path))
    return abraia.download_file(path, dest)


def remove_file(path):
    return abraia.remove_file(path)


def format_output(files, folders=[]):
    output = '\n'.join(['{:>28}  {}/'.format('', click.style(f['name'], fg='blue', bold=True)) for f in folders]) + '\n'
    output += '\n'.join(['{}  {:>7}  {}'.format(f['date'], f['size'], f['name']) for f in files])
    output += '\ntotal {}'.format(len(files))
    return output


@cli_files.command()
@click.argument('folder', required=False, default='')
def list(folder):
    """List files in abraia."""
    try:
        files, folders = abraia.list_files(folder)
        click.echo(format_output(files, folders))
    except APIError as error:
        echo_error(error)


@cli_files.command()
@click.argument('src', type=click.Path())
@click.argument('folder', required=False, default='')
def upload(src, folder):
    """Upload files to abraia."""
    try:
        files = input_files(src)
        process_map(upload_file, files, itertools.repeat(folder), desc="Uploading")
    except APIError as error:
        echo_error(error)


@cli_files.command()
@click.argument('path')
@click.argument('folder', required=False, default='')
def download(path, folder):
    """Download files from abraia."""
    try:
        files = abraia.list_files(path)[0]
        paths = [file['path'] for file in files]
        process_map(download_file, paths, itertools.repeat(folder), desc="Downloading")
    except APIError as error:
        echo_error(error)


@cli_files.command()
@click.argument('path')
def remove(path):
    """Remove files from abraia."""
    try:
        files = abraia.list_files(path)[0]
        click.echo(format_output(files))
        if files and click.confirm('Are you sure you want to remove the files?'):
            paths = [file['path'] for file in files]
            process_map(remove_file, paths, desc="Removing")
    except APIError as error:
        echo_error(error)


@cli_files.command()
@click.option('--remove', help='Remove file metadata', is_flag=True)
@click.argument('path')
def metadata(path, remove):
    """Load or remove file metadata."""
    try:
        if remove:
            abraia.remove_metadata(path)
        meta = abraia.load_metadata(path)
        # metadata = abraia.load_metadata('usain.jpg')
        # abraia.save_json('usain.json', metadata)
        click.echo(json.dumps(meta, indent=2))
    except APIError as error:
        echo_error(error)


if __name__ == '__main__':
    if not abraia.userid:
        configure()
    else:
        cli()
