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

import os
import json
import click
import itertools
from glob import glob

from abraia.utils import load_image, save_image, get_type, temporal_src, list_dir, process_map
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 echo_error(error):
    click.echo('[' + click.style('Error {}'.format(error.code),
                                 fg='red', bold=True) + '] {}'.format(error.message))


def input_files(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


# TODO: Refactor to work in local
def convert_file(src, dest, args):
    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)
        new_ext = 'png' if mode == 'removebg' else 'jpg'
        dirname = src if os.path.isdir(src) else os.path.dirname(src)
        outputs = [os.path.join('output', f"{os.path.splitext(os.path.relpath(src, dirname))[0]}.{new_ext}") for src in inputs]
        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:
            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('files')
def cli_files():
    """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_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)
        click.echo(json.dumps(meta, indent=2))
    except APIError as error:
        echo_error(error)


@cli.group('dataset')
def cli_dataset():
    """Manage dataset commands."""
    pass


@cli_dataset.command()
def list():
    """List available datasets."""
    from abraia.training import dataset
    click.echo(dataset.list_datasets())


@cli_dataset.command()
@click.argument('project')
@click.argument('query', required=False)
def create(project, query=''):
    """Create or update a dataset."""
    from abraia.training import dataset
    files = []
    project = project.strip('/')
    # output_dir = temporal_src(project)
    output_dir = f"{project}/"
    if query:
        files = dataset.search_images(query, limit=100, output_dir=output_dir, verbose=True)
    if os.path.exists(project):
        files = list_dir(project)
        tiffs = [f for f in files if get_type(f) == 'image/tiff']
        if len(tiffs):
            from abraia.multiple import create_visible
            process_map(create_visible, tiffs, desc="Creating images")    
            files = list_dir(project)
    process_map(upload_file, files, itertools.repeat(project + '/'), desc="Uploading images")
    if project not in dataset.list_datasets():
        dataset.save_annotations(project=project, annotations=[])


@cli_dataset.command()
@click.argument('project')
@click.argument('label')
def annotate(project, label):
    """Annotate a dataset using Grounding Dino."""
    from abraia.training import dataset
    images = dataset.list_images(project)
    annotations = dataset.annotate_images(images, [label])
    dataset.save_annotations(project, annotations)


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