"""
Script to download values from the BfG HAD (Bundesanstalt für Gewässerkunde - Hydrologischer Atlas Deutschland).
For given coordinates and layers, the script will download the corresponding values.

Requests are cached to avoid stressing the server too much. It means that values will not be updated
if the server updates the values. The cache is stored in the file 'bafg_cache.sqlite'.

Usage:

    python .\get_values_from_hag.py --help
    usage: get_values_from_hag.py [-h] [--lat LAT] [--lon LON] [--name NAME]
                                [--layers LAYERS [LAYERS ...]] [--tree]

    Get values from the BfG Hydrologischer Atlas Deutschland

    options:
    -h, --help            show this help message and exit
    --lat LAT             Latitude
    --lon LON             Longitude
    --name NAME           Location name
    --layers LAYERS [LAYERS ...]
                            Layers to download
    --tree                Display the tree of layers?

Show the layers tree:

    python .\get_values_from_hag.py --tree
    BfG Hydrologischer Atlas Deutschland
    ├── Teil 1: Grundlagen (id = 0)
    │   ├── 1.1 Orohydrographie (id = 1)
    │   │   └── Höhenschichten in m (id = 2)
    │   ├── 1.2 Gewässernetzdichte (id = 3)
    │   │   └── Gewässernetzdichte in km/km² (id = 4)
    │   ├── 1.3 Bodenübersicht (id = 5)
    │   │   └── Böden (id = 6)
    │   ├── 1.4 Bodenbedeckung (id = 7)
    ...

Download values for a location:

    python get_values_from_hag.py --name Grombühl --lat 49.8023 --lon 9.9449
    Values for Grombühl (49.8023°N, 9.9449°E):
    Layer: Niederschlagshöhe [mm] (id = 35)
        Value: 702 mm
    Layer: mittlere Korrektur Sommerniederschlagshöhen 1961 - 1990 [mm] (id = 39)
        Value: 25 mm
    Layer: mittlere Korrektur Winterniederschlagshöhe 1961 - 1990 [mm] (id = 40)
        Value: 34 mm
    Layer: Abflusshöhe [mm/a] (id = 82)
        Value: 275 mm/a

Specify desired layers by id or by name:

    python get_values_from_hag.py --name Grombühl --lat 49.8023 --lon 9.9449 --layers 129 "Wasserbilanz [mm]"
    Values for Grombühl (49.8023°N, 9.9449°E):
    Layer: Grundwasserneubildung [mm/a] (id = 129)
        Value: 80 mm/a
    Layer: Wasserbilanz [mm] (id = 70)
        Value: 92 mm
"""
import json
import requests
import requests_cache
from pyproj import Transformer
import re
import argparse


# Those constants can also be overwritten by command line arguments.
EXAMPLE_LOCATION = "Lutherkirche Fellbach"
EXAMPLE_LAT, EXAMPLE_LON = 48.80872, 9.27631

# Layers to download. Either the layer id or a substring of the layer name can be used.
EXAMPLES_LAYERS = [
    35,  # 'Niederschlagshöhe [mm]', from 2.5 Mittlere korrigierte jährliche Niederschlagshöhe
    'mittlere Korrektur Sommerniederschlagshöhen',
    'mittlere Korrektur Winterniederschlagshöhe',
    'Abflusshöhe [mm/a]'
]

EPSG_ID = 25832  # EPSG:25832, ETRS89 / UTM zone 32N

URL = "https://geoportal.bafg.de/arcgis1/rest/services/HAD/HAD/MapServer/layers?f=json"
SERVICE = "BfG Hydrologischer Atlas Deutschland"
UNIT = re.compile(r"\[(.*)\]")
TIMEOUT = 30  # seconds


def parse_branch(tree, layer_id=0, level=1):
    layer = tree[layer_id]
    yield (layer['name'].strip(), layer['id'], level)
    sub_layer_infos = layer.get('subLayers', [])
    for sub_layer_info in sub_layer_infos:
        yield from parse_branch(tree, sub_layer_info['id'], level + 1)


def parse_tree(tree):
    for node in tree:
        if not node['parentLayer']:
            yield from parse_branch(tree, node['id'])


def display_tree(tree):
    from rich.console import Console
    from rich.tree import Tree
    from rich.markup import escape
    console = Console()
    viz = Tree(SERVICE)
    parents = {0: viz}
    for name, layer_id, level in parse_tree(tree):
        branch = Tree(escape(f"{name} (id = {layer_id})"))
        parents[level] = branch
        parents[level - 1].add(branch)
    console.print(viz, highlight=True)


def find_layer_by_name(tree, desired_name):
    found = None
    for full_name, layer_id, _level in parse_tree(tree):
        if desired_name in full_name:
            if found:
                raise ValueError(f"'{desired_name}' has been found multiple times!")
            found = full_name, layer_id
    if found:
        return found
    else:
        raise ValueError(f"'{desired_name}' has not been found!")


def find_layer_by_id(tree, desired_id):
    for full_name, layer_id, _level in parse_tree(tree):
        if layer_id == desired_id:
            return full_name, layer_id
    raise ValueError(f"Layer with id {desired_id} has not been found!")


def download_tree(url):
    return get_json(url)['layers']


def get_json(url):
    response = requests.get(url, timeout=TIMEOUT)
    text = response.content
    data = json.loads(text)
    return data


def get_values_from_hag(lat, lon, layer_id):
    transformer = Transformer.from_crs("EPSG:4326",
                                       f"EPSG:{EPSG_ID}",
                                       always_xy=True)
    x, y = transformer.transform(lon, lat)
    x = round(x)
    y = round(y)
    url = f"https://geoportal.bafg.de/arcgis1/rest/services/HAD/HAD/MapServer/{layer_id}/" \
        f"query?f=json&geometry=%7B%22xmin%22%3A{x - 1}%2C%22ymin%22%3A{y - 1}%2C%22" \
        f"xmax%22%3A{x + 1}%2C%22ymax%22%3A{y + 1}%7D&" \
        f"maxAllowableOffset=0&outFields=OBJECTID%2CWert&spatialRel=esriSpatialRelIntersects" \
        f"&where=1%3D1&geometryType=esriGeometryEnvelope&inSR={EPSG_ID}&outSR={EPSG_ID}"
    data = get_json(url)
    try:
        return data['features'][0]['attributes']['Wert']
    except KeyError as exc:
        import pprint
        pprint.pprint(data)
        raise ValueError(f"Could not find a value for layer {layer_id} at {x}, {y}.") from exc


def main(lat, lon, location, layers, should_show_tree):
    requests_cache.install_cache('bafg_cache')
    tree = download_tree(URL)
    if should_show_tree:
        display_tree(tree)
        return

    print(f"Values for {location} ({lat}°N, {lon}°E):")

    for layer_description in layers:
        if isinstance(layer_description, int) or layer_description.isnumeric():
            full_name, layer_id = find_layer_by_id(tree, int(layer_description))
        else:
            full_name, layer_id = find_layer_by_name(tree, layer_description)
        if (match := UNIT.search(full_name)):
            unit = match.group(1)
        else:
            unit = ""
        print(f"  Layer: {full_name} (id = {layer_id})")
        print(f"    Value: {get_values_from_hag(lat, lon, layer_id)} {unit}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=f'Get values from the {SERVICE}')
    parser.add_argument('--lat', type=float, default=EXAMPLE_LAT, help='Latitude')
    parser.add_argument('--lon', type=float, default=EXAMPLE_LON, help='Longitude')
    parser.add_argument('--name', type=str, default=EXAMPLE_LOCATION, help='Location name')
    parser.add_argument('--layers', type=str, nargs='+', default=EXAMPLES_LAYERS, help='Layers to download')
    parser.add_argument('--tree', action='store_true', help='Display the tree of layers?')
    args = parser.parse_args()
    main(args.lat, args.lon, args.name, args.layers, args.tree)