#!/usr/bin/env python3
import signal
signal.signal(signal.SIGINT, signal.SIG_DFL)

import os
import sys
import getopt
import logging
import json
import functools
import time
import traceback
from concurrent.futures import ProcessPoolExecutor, as_completed

import ccbrowse
from ccbrowse import utils
from ccbrowse.exceptions import StorageNotAvailable
from ccbrowse.products import PRODUCTS, Calipso, CloudSat, NaturalEarth


def usage():
    sys.stderr.write('''Usage: {program_name} [option]... TYPE FILE...
       {program_name} --help
Try `{program_name} --help' for more information.
'''.format(program_name=program_name))


def print_help():
    sys.stdout.write('''Usage: {program_name} [OPTIONS] TYPE FILE...
       {program_name} --help

Import data from FILE into profile specified in configuration file.

Positional arguments:
  TYPE             product type
  FILE             product file

Optional arguments:
  -c FILE          configuration file (default: config.json)
  -j N             number of parallel jobs (default: number of CPU cores)
  -l LAYER         import only specified profile layer
  --overwrite      overwrite existing tiles
  -s               print statistics
  --skip           skip tiles that exist
  --hard           hard import
  --help           print this help information and exit
  -z ZOOM          import only specified zoom level

Supported product types:
  calipso
  cloudsat
  naturalearth

Report bugs to <ccplot-general@lists.sourceforge.net>
'''.format(program_name=program_name))


def save_decorator(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        try: return f(*args, **kwargs)
        except KeyboardInterrupt: print()
    return wrapper


@save_decorator
def save(product, profile, layer=None, zoom=None,
         statistics=False, soft=False, overwrite=False, skip=False):
    if statistics: stat = dict(n=0, read=0, save=0)

    XX = []

    if soft:
        product_name = list(PRODUCTS.keys())[list(PRODUCTS.values()).index(type(product))]
        filename = os.path.abspath(product.filename)
        root = profile.get_root()
        if filename.startswith(root): filename = filename[len(root):]

        ref = {
            'product': product_name,
            'filename': filename,
            'offset': product.offset(),
            'bounds': product.bounds(),
        }
        obj = {'ref': [ref]}
        try: profile.save(obj)
        except StorageNotAvailable: pass
        else:
            for level in sorted(profile['zoom'].keys()):
                for l in product.layers():
                    X = product.xrange(l, level)
                    XX += [[l, level, X]]
            return XX

    for level in sorted(profile['zoom'].keys()):
        if zoom != None and level != zoom: continue
        for l in product.layers():
            if layer != None and l != layer: continue
            line = None
            X = product.xrange(l, level)
            Z = product.zrange(l, level)
            size = len(X)*len(Z)
            i = 0
            for x in X:
                for z in Z:
                    line = '%s level %s tiles %d--%d [%d/%d] %.f%%' % \
                           (l, level, X[0], X[-1], i, size, 100.0*i/size)
                    if sys.stdout.isatty():
                        if line: sys.stdout.write('\r\033[K' + line)
                        sys.stdout.flush()

                    i = i + 1

                    if statistics: t1 = time.clock()

                    if skip:
                        obj = profile.load({
                            'layer': l,
                            'zoom': level,
                            'x': x,
                            'z': z,
                        }, exclude=['data'], dereference=False)
                        if obj is not None: continue

                    if not soft:
                        tile = product.tile(l, level, x, z)
                    else:
                        tile = {
                            'layer': l,
                            'zoom': level,
                            'x': x,
                            'z': z,
                            'ref': [ref],
                        }

                    if statistics:
                        t2 = time.clock()
                        stat['read'] += t2 - t1

                    profile.save(tile, append=(not overwrite))

                    if statistics:
                        stat['save'] += time.clock() - t2
                        stat['n'] += 1

            XX += [[l, level, X]]

            if line: print('\r\033[K%s level %s tiles %d--%d' % \
                           (l, level, X[0], X[-1]))
    if statistics and stat['n'] > 0:
        print('Statistics: read and processing %.2f ms/tile, save %.2f ms/tile' % \
              (1000*stat['read']/stat['n'], 1000*stat['save']/stat['n']))

    return XX


def worker_init(config):
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    global profile
    profile = ccbrowse.Profile(config, cache_size=0, write_availability=False)


def worker(file_type, filename, layer, zoom, options):
    print(filename)
    try:
        if file_type == 'calipso':
            product = Calipso(filename, profile)
            return save(product, profile, layer=layer, zoom=zoom, **options)
        elif file_type == 'cloudsat':
            product = CloudSat(filename, profile)
            return save(product, profile, layer=layer, zoom=zoom, **options)
        elif file_type == 'naturalearth':
            naturalearth = NaturalEarth(filename, profile)
            naturalearth.save(layer=layer)
    except Exception:
        logging.error('%s: %s' % (filename, traceback.format_exc()))

if __name__ == '__main__':
    program_name = 'ccbrowse import'
    logging.basicConfig(format=program_name+': %(message)s', level=logging.INFO)
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'c:j:l:sz:',
            ['help', 'overwrite', 'skip', 'hard'])
    except getopt.GetoptError as e:
        logging.error(e)
        usage()
        sys.exit(1)

    config_filename = 'config.json'
    njobs = os.cpu_count()
    file_type = None
    layer = None
    zoom = None

    options = {
        'statistics': False,
        'overwrite': False,
        'skip': False,
        'soft': True,
    }

    for opt, value in opts:
        if opt == '--help':
            print_help()
            sys.exit(0)
        elif opt == '-c':
            config_filename = value
        elif opt == '-j':
            try:
                njobs = int(value)
                assert(njobs > 0)
            except (ValueError, AssertionError):
                logging.error('Invalid number of jobs "%s"' % value)
                sys.exit(1)
        elif opt == '-l':
            layer = value
        elif opt == '-s':
            options['statistics'] = True
        elif opt == '-z':
            zoom = value
        elif opt == '--overwrite':
            options['overwrite'] = True
        elif opt == '--skip':
            options['skip'] = True
        elif opt == '--hard':
            options['soft'] = False

    if len(args) < 2:
        usage()
        sys.exit(1)

    file_type = args[0]
    filenames = args[1:]

    if not file_type in ['calipso', 'cloudsat', 'naturalearth']:
        logging.error('Invalid file type "%s"' % file_type)
        sys.exit(1)

    try:
        with open(config_filename) as fp:
            config = json.load(fp)
    except IOError as e:
        logging.error('%s: %s' % (e.filename, e.strerror))
        sys.exit(1)

    try:
        tasks = [
            (file_type, filename, layer, zoom, options)
            for filename in filenames
        ]
        with ccbrowse.Profile(config, cache_size=0) as profile:
            with ProcessPoolExecutor(
                max_workers=njobs,
                initializer=worker_init,
                initargs=(config,),
            ) as ex:
                fs = [ex.submit(worker, *task) for task in tasks]
                for f in as_completed(fs):
                    res = f.result()
                    if res is None: continue
                    for layer, level, X in res:
                        intervals = utils.intervals(X)
                        profile.update_availability(layer, level, intervals)
    except IOError as e:
        if e.filename is not None and e.strerror is not None:
            logging.error('%s: %s' % (e.filename, e.strerror))
        else:
            logging.error(e)
