#!/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 anonymize_image, blur_background, clean_image, remove_background, upscale_image

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/editor/', 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):
    img = load_image(src)
    if mode == 'anonymize':
        out = anonymize_image(img)
    elif mode == 'blur':
        out = blur_background(img)
    elif mode == 'clean':
        out = clean_image(img)
    elif mode == 'removebg':
        out = remove_background(img)
    elif mode == 'upscale':
        out = upscale_image(img)
    else:
        out = img
    save_image(out, dest)
    return 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 background")


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


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


@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."""
    editing_files(src, 'clean', desc="Removing objects")


@cli.group('cloud')
def cli_cloud():
    """Manage files on the 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_cloud.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_cloud.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_cloud.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_cloud.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_cloud.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)
        click.echo(json.dumps(meta, indent=2))
    except APIError as error:
        echo_error(error)


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