import importlib
import json
import logging
import os
import pprint

import geometamaker
import natcap.invest
import pint

from natcap.invest import utils
from . import gettext
from .unit_registry import u


LOGGER = logging.getLogger(__name__)

# Specs for common arg types ##################################################
WORKSPACE = {
    "name": gettext("workspace"),
    "about": gettext(
        "The folder where all the model's output files will be written. If "
        "this folder does not exist, it will be created. If data already "
        "exists in the folder, it will be overwritten."),
    "type": "directory",
    "contents": {},
    "must_exist": False,
    "permissions": "rwx",
}

SUFFIX = {
    "name": gettext("file suffix"),
    "about": gettext(
        "Suffix that will be appended to all output file names. Useful to "
        "differentiate between model runs."),
    "type": "freestyle_string",
    "required": False,
    "regexp": "[a-zA-Z0-9_-]*"
}

N_WORKERS = {
    "name": gettext("taskgraph n_workers parameter"),
    "about": gettext(
        "The n_workers parameter to provide to taskgraph. "
        "-1 will cause all jobs to run synchronously. "
        "0 will run all jobs in the same process, but scheduling will take "
        "place asynchronously. Any other positive integer will cause that "
        "many processes to be spawned to execute tasks."),
    "type": "number",
    "units": u.none,
    "required": False,
    "expression": "value >= -1"
}

METER_RASTER = {
    "type": "raster",
    "bands": {
        1: {
            "type": "number",
            "units": u.meter
        }
    }
}
AOI = {
    "type": "vector",
    "fields": {},
    "geometries": {"POLYGON", "MULTIPOLYGON"},
    "name": gettext("area of interest"),
    "about": gettext(
        "A map of areas over which to aggregate and "
        "summarize the final results."),
}
LULC = {
    "type": "raster",
    "bands": {1: {"type": "integer"}},
    "about": gettext(
        "Map of land use/land cover codes. Each land use/land cover type "
        "must be assigned a unique integer code."),
    "name": gettext("land use/land cover")
}
DEM = {
    "type": "raster",
    "bands": {
        1: {
            "type": "number",
            "units": u.meter
        }
    },
    "about": gettext("Map of elevation above sea level."),
    "name": gettext("digital elevation model")
}
PRECIP = {
    "type": "raster",
    "bands": {
        1: {
            "type": "number",
            "units": u.millimeter/u.year
        }
    },
    "about": gettext("Map of average annual precipitation."),
    "name": gettext("precipitation")
}
ET0 = {
    "name": gettext("reference evapotranspiration"),
    "type": "raster",
    "bands": {
        1: {
            "type": "number",
            "units": u.millimeter
        }
    },
    "about": gettext("Map of reference evapotranspiration values.")
}
SOIL_GROUP = {
    "type": "raster",
    "bands": {1: {"type": "integer"}},
    "about": gettext(
        "Map of soil hydrologic groups. Pixels may have values 1, 2, 3, or 4, "
        "corresponding to soil hydrologic groups A, B, C, or D, respectively."),
    "name": gettext("soil hydrologic group")
}
THRESHOLD_FLOW_ACCUMULATION = {
    "expression": "value >= 0",
    "type": "number",
    "units": u.pixel,
    "about": gettext(
        "The number of upslope pixels that must flow into a pixel "
        "before it is classified as a stream."),
    "name": gettext("threshold flow accumulation")
}
LULC_TABLE_COLUMN = {
    "type": "integer",
    "about": gettext(
        "LULC codes from the LULC raster. Each code must be a unique "
        "integer.")
}

# Specs for common outputs ####################################################
TASKGRAPH_DIR = {
    "type": "directory",
    "about": (
        "Cache that stores data between model runs. This directory contains no "
        "human-readable data and you may ignore it."),
    "contents": {
        "taskgraph.db": {}
    }
}
FILLED_DEM = {
    "about": gettext("Map of elevation after any pits are filled"),
    "bands": {1: {
        "type": "number",
        "units": u.meter
    }}
}
FLOW_ACCUMULATION = {
    "about": gettext("Map of flow accumulation"),
    "bands": {1: {
        "type": "number",
        "units": u.none
    }}
}
FLOW_DIRECTION = {
    "about": gettext(
        "MFD flow direction. Note: the pixel values should not "
        "be interpreted directly. Each 32-bit number consists "
        "of 8 4-bit numbers. Each 4-bit number represents the "
        "proportion of flow into one of the eight neighboring "
        "pixels."),
    "bands": {1: {"type": "integer"}}
}
FLOW_DIRECTION_D8 = {
    "about": gettext(
        "D8 flow direction."),
    "bands": {1: {"type": "integer"}}
}
SLOPE = {
    "about": gettext(
        "Percent slope, calculated from the pit-filled "
        "DEM. 100 is equivalent to a 45 degree slope."),
    "bands": {1: {"type": "percent"}}
}
STREAM = {
    "about": "Stream network, created using flow direction and flow accumulation derived from the DEM and Threshold Flow Accumulation. Values of 1 represent streams, values of 0 are non-stream pixels.",
    "bands": {1: {"type": "integer"}}
}

FLOW_DIR_ALGORITHM = {
    "flow_dir_algorithm": {
        "type": "option_string",
        "options": {
            "D8": {
                "display_name": gettext("D8"),
                "description": "D8 flow direction"
            },
            "MFD": {
                "display_name": gettext("MFD"),
                "description": "Multiple flow direction"
            }
        },
        "about": gettext("Flow direction algorithm to use."),
        "name": gettext("flow direction algorithm")
    }
}

# geometry types ##############################################################
# the full list of ogr geometry types is in an enum in
# https://github.com/OSGeo/gdal/blob/master/gdal/ogr/ogr_core.h

POINT = {'POINT'}
LINESTRING = {'LINESTRING'}
POLYGON = {'POLYGON'}
MULTIPOINT = {'MULTIPOINT'}
MULTILINESTRING = {'MULTILINESTRING'}
MULTIPOLYGON = {'MULTIPOLYGON'}

LINES = LINESTRING | MULTILINESTRING
POLYGONS = POLYGON | MULTIPOLYGON
POINTS = POINT | MULTIPOINT
ALL_GEOMS = LINES | POLYGONS | POINTS


def format_unit(unit):
    """Represent a pint Unit as user-friendly unicode text.

    This attempts to follow the style guidelines from the NIST
    Guide to the SI (https://www.nist.gov/pml/special-publication-811):
    - Use standard symbols rather than spelling out
    - Use '/' to represent division
    - Use the center dot ' · ' to represent multiplication
    - Combine denominators into one, surrounded by parentheses

    Args:
        unit (pint.Unit): the unit to format

    Raises:
        TypeError if unit is not an instance of pint.Unit.

    Returns:
        String describing the unit.
    """
    if not isinstance(unit, pint.Unit):
        raise TypeError(
            f'{unit} is of type {type(unit)}. '
            f'It should be an instance of pint.Unit')

    # Optionally use a pre-set format for a particular unit
    custom_formats = {
        u.pixel: gettext('number of pixels'),
        u.year_AD: '',  # don't need to mention units for a year input
        u.other: '',    # for inputs that can have any or multiple units
        # For soil erodibility (t*h*ha/(ha*MJ*mm)), by convention the ha's
        # are left on top and bottom and don't cancel out
        # pint always cancels units where it can, so add them back in here
        # this isn't a perfect solution
        # see https://github.com/hgrecco/pint/issues/1364
        u.t * u.hr / (u.MJ * u.mm): 't · h · ha / (ha · MJ · mm)',
        u.none: gettext('unitless')
    }
    if unit in custom_formats:
        return custom_formats[unit]

    # look up the abbreviated symbol for each unit
    # `formatter` expects an iterable of (unit, exponent) pairs, which lives in
    # the pint.Unit's `_units` attribute.
    unit_items = [(u.get_symbol(key), val) for key, val in unit._units.items()]
    formatted_unit = pint.formatting.formatter(
        unit_items,
        as_ratio=True,
        single_denominator=True,
        product_fmt=" · ",
        division_fmt='/',
        power_fmt="{}{}",
        parentheses_fmt="({})",
        exp_call=pint.formatting._pretty_fmt_exponent)

    if 'currency' in formatted_unit:
        formatted_unit = formatted_unit.replace('currency', gettext('currency units'))
    return formatted_unit


def serialize_args_spec(spec):
    """Serialize an MODEL_SPEC dict to a JSON string.

    Args:
        spec (dict): An invest model's MODEL_SPEC.

    Raises:
        TypeError if any object type within the spec is not handled by
        json.dumps or by the fallback serializer.

    Returns:
        JSON String
    """

    def fallback_serializer(obj):
        """Serialize objects that are otherwise not JSON serializeable."""
        if isinstance(obj, pint.Unit):
            return format_unit(obj)
        # Sets are present in 'geometries' attributes of some args
        # We don't need to worry about deserializing back to a set/array
        # so casting to string is okay.
        elif isinstance(obj, set):
            return str(obj)
        raise TypeError(f'fallback serializer is missing for {type(obj)}')

    return json.dumps(spec, default=fallback_serializer)


# accepted geometries for a vector will be displayed in this order
GEOMETRY_ORDER = [
    'POINT',
    'MULTIPOINT',
    'LINESTRING',
    'MULTILINESTRING',
    'POLYGON',
    'MULTIPOLYGON']

INPUT_TYPES_HTML_FILE = 'input_types.html'


def format_required_string(required):
    """Represent an arg's required status as a user-friendly string.

    Args:
        required (bool | str | None): required property of an arg. May be
            `True`, `False`, `None`, or a conditional string.

    Returns:
        string
    """
    if required is None or required is True:
        return gettext('required')
    elif required is False:
        return gettext('optional')
    else:
        # assume that the about text will describe the conditional
        return gettext('conditionally required')


def format_geometries_string(geometries):
    """Represent a set of allowed vector geometries as user-friendly text.

    Args:
        geometries (set(str)): set of geometry names

    Returns:
        string
    """
    # sort the geometries so they always display in a consistent order
    sorted_geoms = sorted(
        geometries,
        key=lambda g: GEOMETRY_ORDER.index(g))
    return '/'.join(gettext(geom).lower() for geom in sorted_geoms)


def format_permissions_string(permissions):
    """Represent a rwx-style permissions string as user-friendly text.

    Args:
        permissions (str): rwx-style permissions string

    Returns:
        string
    """
    permissions_strings = []
    if 'r' in permissions:
        permissions_strings.append(gettext('read'))
    if 'w' in permissions:
        permissions_strings.append(gettext('write'))
    if 'x' in permissions:
        permissions_strings.append(gettext('execute'))
    return ', '.join(permissions_strings)


def format_options_string_from_dict(options):
    """Represent a dictionary of option: description pairs as a bulleted list.

    Args:
        options (dict): the dictionary of options to document, where keys are
            options and values are dictionaries describing the options.
            They may have either or both 'display_name' and 'description' keys,
            for example:
            {'option1': {'display_name': 'Option 1', 'description': 'the first option'}}

    Returns:
        list of RST-formatted strings, where each is a line in a bullet list
    """
    lines = []
    for key, info in options.items():
        display_name = info['display_name'] if 'display_name' in info else key
        if 'description' in info:
            lines.append(f'- {display_name}: {info["description"]}')
        else:
            lines.append(f'- {display_name}')
    # sort the options alphabetically
    # casefold() is a more aggressive version of lower() that may work better
    # for some languages to remove all case distinctions
    return sorted(lines, key=lambda line: line.casefold())


def format_options_string_from_list(options):
    """Represent options as a comma-separated list.

    Args:
        options (list[str]): the set of options to document

    Returns:
        string of comma-separated options
    """
    return ', '.join(options)


def capitalize(title):
    """Capitalize a string into title case.

    Args:
        title (str): string to capitalize

    Returns:
        capitalized string (each word capitalized except linking words)
    """

    def capitalize_word(word):
        """Capitalize a word, if appropriate."""
        if word in {'of', 'the'}:
            return word
        else:
            return word[0].upper() + word[1:]

    title = ' '.join([capitalize_word(word) for word in title.split(' ')])
    title = '/'.join([capitalize_word(word) for word in title.split('/')])
    return title


def format_type_string(arg_type):
    """Represent an arg type as a user-friendly string.

    Args:
        arg_type (str|set(str)): the type to format. May be a single type or a
            set of types.

    Returns:
        formatted string that links to a description of the input type(s)
    """
    # some types need a more user-friendly name
    # all types are listed here so that they can be marked up for translation
    type_names = {
        'boolean': gettext('true/false'),
        'csv': gettext('CSV'),
        'directory': gettext('directory'),
        'file': gettext('file'),
        'freestyle_string': gettext('text'),
        'integer': gettext('integer'),
        'number': gettext('number'),
        'option_string': gettext('option'),
        'percent': gettext('percent'),
        'raster': gettext('raster'),
        'ratio': gettext('ratio'),
        'vector': gettext('vector')
    }

    def format_single_type(arg_type):
        """Represent a type as a link to the corresponding Input Types section.

        Args:
            arg_type (str): the type to format.

        Returns:
            formatted string that links to a description of the input type
        """
        # Represent the type as a string. Some need a more user-friendly name.
        # we can only use standard docutils features here, so no :ref:
        # this syntax works to link to a section in a different page, but it
        # isn't universally supported and depends on knowing the built page name.
        if arg_type == 'freestyle_string':
            section_name = 'text'
        elif arg_type == 'option_string':
            section_name = 'option'
        elif arg_type == 'boolean':
            section_name = 'truefalse'
        elif arg_type == 'csv':
            section_name = 'csv'
        else:
            section_name = arg_type

        return f'`{type_names[arg_type]} <{INPUT_TYPES_HTML_FILE}#{section_name}>`__'

    if isinstance(arg_type, set):
        return ' or '.join(format_single_type(t) for t in sorted(arg_type))
    else:
        return format_single_type(arg_type)


def describe_arg_from_spec(name, spec):
    """Generate RST documentation for an arg, given an arg spec.

    This is used for documenting:
        - a single top-level arg
        - a row or column in a CSV
        - a field in a vector
        - an item in a directory

    Args:
        name (str): Name to give the section. For top-level args this is
            arg['name']. For nested args it's typically their key in the
            dictionary one level up.
        spec (dict): A arg spec dictionary that conforms to the InVEST args
            spec specification. It must at least have the key `'type'`, and
            whatever other keys are expected for that type.
    Returns:
        list of strings, where each string is a line of RST-formatted text.
        The first line has the arg name, type, required state, description,
        and units if applicable. Depending on the type, there may be additional
        lines that are indented, that describe details of the arg such as
        vector fields and geometries, option_string options, etc.
    """
    type_string = format_type_string(spec['type'])
    in_parentheses = [type_string]

    # For numbers and rasters that have units, display the units
    units = None
    if spec['type'] == 'number':
        units = spec['units']
    elif spec['type'] == 'raster' and spec['bands'][1]['type'] == 'number':
        units = spec['bands'][1]['units']
    if units:
        units_string = format_unit(units)
        if units_string:
            # pybabel can't find the message if it's in the f-string
            translated_units = gettext("units")
            in_parentheses.append(f'{translated_units}: **{units_string}**')

    if spec['type'] == 'vector':
        in_parentheses.append(format_geometries_string(spec["geometries"]))

    # Represent the required state as a string, defaulting to required
    # It doesn't make sense to include this for boolean checkboxes
    if spec['type'] != 'boolean':
        # get() returns None if the key doesn't exist in the dictionary
        required_string = format_required_string(spec.get('required'))
        in_parentheses.append(f'*{required_string}*')

    # Nested args may not have an about section
    if 'about' in spec:
        sanitized_about_string = spec["about"].replace("_", "\\_")
        about_string = f': {sanitized_about_string}'
    else:
        about_string = ''

    first_line = f"**{name}** ({', '.join(in_parentheses)}){about_string}"

    # Add details for the types that have them
    indented_block = []
    if spec['type'] == 'option_string':
        # may be either a dict or set. if it's empty, the options are
        # dynamically generated. don't try to document them.
        if spec['options']:
            if isinstance(spec['options'], dict):
                indented_block.append(gettext('Options:'))
                indented_block += format_options_string_from_dict(spec['options'])
            else:
                formatted_options = format_options_string_from_list(spec['options'])
                indented_block.append(gettext('Options:') + f' {formatted_options}')

    elif spec['type'] == 'csv':
        if 'columns' not in spec and 'rows' not in spec:
            first_line += gettext(
                ' Please see the sample data table for details on the format.')

    # prepend the indent to each line in the indented block
    return [first_line] + ['\t' + line for line in indented_block]


def describe_arg_from_name(module_name, *arg_keys):
    """Generate RST documentation for an arg, given its model and name.

    Args:
        module_name (str): invest model module containing the arg.
        *arg_keys: one or more strings that are nested arg keys.

    Returns:
        String describing the arg in RST format. Contains an anchor named
        <arg_keys[0]>-<arg_keys[1]>...-<arg_keys[n]>
        where underscores in arg keys are replaced with hyphens.
    """
    # import the specified module (that should have an MODEL_SPEC attribute)
    module = importlib.import_module(module_name)
    # start with the spec for all args
    # narrow down to the nested spec indicated by the sequence of arg keys
    spec = module.MODEL_SPEC['args']
    for i, key in enumerate(arg_keys):
        # convert raster band numbers to ints
        if arg_keys[i - 1] == 'bands':
            key = int(key)
        try:
            spec = spec[key]
        except KeyError:
            keys_so_far = '.'.join(arg_keys[:i + 1])
            raise ValueError(
                f"Could not find the key '{keys_so_far}' in the "
                f"{module_name} model's MODEL_SPEC")

    # format spec into an RST formatted description string
    if 'name' in spec:
        arg_name = capitalize(spec['name'])
    else:
        arg_name = arg_keys[-1]

    # anchor names cannot contain underscores. sphinx will replace them
    # automatically, but lets explicitly replace them here
    anchor_name = '-'.join(arg_keys).replace('_', '-')
    rst_description = '\n\n'.join(describe_arg_from_spec(arg_name, spec))
    return f'.. _{anchor_name}:\n\n{rst_description}'


def write_metadata_file(datasource_path, spec, lineage_statement, keywords_list):
    """Write a metadata sidecar file for an invest output dataset.

    Args:
        datasource_path (str) - filepath to the invest output
        spec (dict) -  the invest specification for ``datasource_path``
        lineage_statement (str) - string to describe origin of the dataset.
        keywords_list (list) - sequence of strings

    Returns:
        None

    """
    resource = geometamaker.describe(datasource_path)
    resource.set_lineage(lineage_statement)
    # a pre-existing metadata doc could have keywords
    words = resource.get_keywords()
    resource.set_keywords(set(words + keywords_list))

    if 'about' in spec:
        resource.set_description(spec['about'])
    attr_spec = None
    if 'columns' in spec:
        attr_spec = spec['columns']
    if 'fields' in spec:
        attr_spec = spec['fields']
    if attr_spec:
        for key, value in attr_spec.items():
            about = value['about'] if 'about' in value else ''
            units = format_unit(value['units']) if 'units' in value else ''
            try:
                resource.set_field_description(
                    key, description=about, units=units)
            except KeyError as error:
                # fields that are in the spec but missing
                # from model results because they are conditional.
                LOGGER.debug(error)
    if 'bands' in spec:
        for idx, value in spec['bands'].items():
            try:
                units = format_unit(spec['bands'][idx]['units'])
            except KeyError:
                units = ''
            resource.set_band_description(idx, units=units)

    resource.write()


def generate_metadata(model_module, args_dict):
    """Create metadata for all items in an invest model output workspace.

    Args:
        model_module (object) - the natcap.invest module containing
            the MODEL_SPEC attribute
        args_dict (dict) - the arguments dictionary passed to the
            model's ``execute`` function.

    Returns:
        None

    """
    file_suffix = utils.make_suffix_string(args_dict, 'results_suffix')
    formatted_args = pprint.pformat(args_dict)
    lineage_statement = (
        f'Created by {model_module.__name__}.execute(\n{formatted_args})\n'
        f'Version {natcap.invest.__version__}')
    keywords = [model_module.MODEL_SPEC['model_id'], 'InVEST']

    def _walk_spec(output_spec, workspace):
        for filename, spec_data in output_spec.items():
            if 'type' in spec_data and spec_data['type'] == 'directory':
                if 'taskgraph.db' in spec_data['contents']:
                    continue
                _walk_spec(
                    spec_data['contents'],
                    os.path.join(workspace, filename))
            else:
                pre, post = os.path.splitext(filename)
                full_path = os.path.join(workspace, f'{pre}{file_suffix}{post}')
                if os.path.exists(full_path):
                    try:
                        write_metadata_file(
                            full_path, spec_data, lineage_statement, keywords)
                    except ValueError as error:
                        # Some unsupported file formats, e.g. html
                        LOGGER.debug(error)

    _walk_spec(model_module.MODEL_SPEC['outputs'], args_dict['workspace_dir'])
