"""
Resources:
 * https://www.riserid.eu/data/user_upload/downloads/info-pdf.s/Diverses/Liste-Amtlicher-Gemeindeschluessel-AGS-2015.pdf
 * https://de.wikipedia.org/wiki/Liste_der_Landkreise_in_Deutschland
 * Test the query: https://query.wikidata.org/#SELECT%20DISTINCT%20%3Fcity%20%3FcityLabel%20%3Fags%20WHERE%20%7B%0A%20%20%7B%20%3Fcity%20rdfs%3Alabel%20%22Ahrweiler%22.%20%7D%0A%20%20UNION%0A%20%20%7B%20%3Fcity%20rdfs%3Alabel%20%22Ahrweiler%22%40de.%20%7D%0A%20%20UNION%0A%20%20%7B%20%3Fcity%20rdfs%3Alabel%20%22Ahrweiler%22%40en.%20%7D%0A%20%20%7B%20%3Fcity%20wdt%3AP440%20%3Fags.%20%7D%0A%20%20UNION%0A%20%20%7B%20%3Fcity%20wdt%3AP439%20%3Fags.%20%7D%0A%20%20SERVICE%20wikibase%3Alabel%20%7B%20bd%3AserviceParam%20wikibase%3Alanguage%20%22de%2Cen%22.%20%7D%0A%7D%0A
"""

from typing import Tuple, List, Optional

import requests

query_template = """
SELECT DISTINCT ?city ?cityLabel ?ags WHERE {{
  {{ ?city rdfs:label "{0}". }}
  UNION
  {{ ?city rdfs:label "{0}"@de. }}
  UNION
  {{ ?city rdfs:label "{0}"@en. }}
  {{ ?city wdt:P440 ?ags. }}
  UNION
  {{ ?city wdt:P439 ?ags. }}
  SERVICE wikibase:label {{ bd:serviceParam wikibase:language "de,en". }}
}}
"""
wikidata_sparql = "https://query.wikidata.org/sparql"


def city_to_ags_all(city_name: str) -> List[Tuple[str, str]]:
    query = query_template.format(city_name)
    response = requests.get(wikidata_sparql, {"format": "json", "query": query})
    response.raise_for_status()
    data = response.json()
    values = data["results"]["bindings"]
    pairs = set()
    for i in values:
        city = i["cityLabel"]["value"]
        value = i["ags"]["value"]
        if len(value) != 5 and len(value) != 8:
            raise RuntimeError(
                "Invalid Amtlicher Gemeindeschlüssel: '{}'".format(value)
            )
        pairs.add((city, value))

    return list(pairs)


def city_to_ags(city_name: str, district: bool) -> Optional[str]:
    """ Returns the Amtliche Gemeindeschlüssel"""
    ags_list = city_to_ags_all(city_name)
    if len(ags_list) == 1:
        return ags_list[0][1]

    # Disambiguate between district and city
    if len(ags_list) == 2:
        if len(ags_list[0][1]) == 5 and len(ags_list[1][1]) == 8:
            if district:
                return ags_list[0][1]
            else:
                return ags_list[1][1]
        elif len(ags_list[0][1]) == 8 and len(ags_list[1][1]) == 5:
            if district:
                return ags_list[1][1]
            else:
                return ags_list[0][1]

    return None

@@ -2027,16 +2027,65 @@
 list) ==
+ 0:%0A        return None%0A    elif len(ags_list) ==
  1:%0A    
@@ -2114,16 +2114,36 @@
 %5B1%5D%0A
-%0A
+    else:%0A   
     
+ 
 # 
-D
+Try to d
 isam
@@ -2184,159 +2184,164 @@
     
-if len(ags_list) == 2:%0A        if len(ags_list%5B0%5D%5B1%5D) == 5 and len(ags_list%5B1%5D%5B1%5D) == 8:%0A            if district:%0A                return ags_list%5B0%5D%5B1%5D
+    districts = %5B%5D%0A        cities = %5B%5D%0A        for single_ags in ags_list:%0A            if len(single_ags%5B1%5D) == 8:%0A                cities.append(single_ags)
 %0A   
@@ -2371,37 +2371,44 @@
         
-return ags_list%5B1%5D%5B1%5D
+districts.append(single_ags)
 %0A       
@@ -2412,66 +2412,43 @@
     
-el
 if 
-len(ags_list%5B0%5D%5B1%5D) == 8 and len(ags_list%5B1%5D%5B1%5D
+district and len(districts
 ) == 
-5
+1
 :%0A  
@@ -2461,81 +2461,77 @@
     
-if district:%0A                return ags_list%5B1%5D%5B1%5D%0A            else:%0A    
+return districts%5B0%5D%5B1%5D%0A        if not district and len(cities) == 1:%0A
     
@@ -2541,32 +2541,30 @@
      return 
-ags_list
+cities
 %5B0%5D%5B1%5D%0A%0A    

Make city to ags work with multiple districts