#!/usr/bin/env python3

import os
import sys
import warnings
import logging
import subprocess
import shutil
from glob import glob

import tile_index

LOG = logging.getLogger(__name__)


def group_files(products, input_files):
    """Group input geotiff files by product."""
    groups = {}
    for prod in products:
        prods_files = [f for f in input_files if prod in f]
        if prods_files:
            groups[prod] = prods_files
    if len(groups) != len(products):
        warnings.warn("Not all product geotiffs were provided.")
    return groups


def remap_to_lonlat(itif, otif):
    """Remap a single geotiff by calling gdalwarp."""
    try:
        subprocess.run(['gdalwarp', '-multi', '-wo', 'NUM_THREADS=ALL_CPUS', '-t_srs', 'EPSG:4326', itif, otif], check=True)
    except subprocess.CalledProcessError:
        LOG.error("Could not remap geotiff %s -> %s" % (itif, otif))
        return None
    return otif


def remap_tifs(input_tifs, out_dir, remap_suffix):
    """Remap all input geotiffs to EPSG:4326."""
    for itif in input_tifs:
        ifn = os.path.basename(itif)
        otif = os.path.join(out_dir, ifn.replace('.tif', remap_suffix))
        otif = remap_to_lonlat(itif, otif)
        if otif is not None:
            yield otif


def link_or_copy(input_tifs, out_dir):
    """Hardlink input tifs to output directory."""
    for prod_file in input_tifs:
        out_file = os.path.join(out_dir, os.path.basename(prod_file))
        try:
            os.link(prod_file, out_file)
        except OSError:
            # on different mounts probably?
            shutil.copy2(prod_file, out_file)
        yield out_file


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Take input geotiffs and generate mapserver compatible tiles.")
    parser.add_argument('--remap', action='store_true',
                        help="Remap input geotiffs to EPSG:4326")
    parser.add_argument('--remap-suffix', default='_ll.tif',
                        help="Replace 'tif' with provided suffix when geotiffs are remapped.")
    parser.add_argument('-p', '--products', nargs="*",
                        help="Product names to group together in each "
                             "'layer'. Product name must be in the filename.")
    parser.add_argument('--shape-file', default='{product}.shp',
                        help="Shapefile filename pattern to use and placed "
                             "in the output directory. "
                             "(default: '{product}.shp')")
    parser.add_argument('out_dir',
                        help="Output path to save tile information to (ex. '/data/tiles/{product}')")
    parser.add_argument('input_files', nargs="+",
                        help="Input geotiffs to generate tiles for (separate from product lists with '--')")
    args = parser.parse_args()

    groups = group_files(args.products, args.input_files)
    for prod, prod_files in groups.items():
        out_dir = args.out_dir.format(product=prod)
        os.makedirs(out_dir, exist_ok=True)
        shp_fn = args.shape_file.format(product=prod)
        shp_pathname = os.path.join(out_dir, shp_fn)

        if args.remap:
            # remap if needed
            prod_files = list(remap_tifs(prod_files, out_dir, args.remap_suffix))
        else:
            # hardlink if needed
            prod_files = list(link_or_copy(prod_files, out_dir))

        # get all products in the current directory
        ext = os.path.splitext(prod_files[-1])[-1]
        all_prod_files = sorted(glob(os.path.join(out_dir, '*' + prod + '*' + ext)))

        # create shape file
        LOG.info("Rebuilding shapefile index with:\n\t{}".format(", ".join(all_prod_files)))
        tile_index.index(all_prod_files, shp_pathname)



if __name__ == "__main__":
    sys.exit(main())