# steps to upload the geoip2 database to the stage location
# 1. download the geoip2 to local folder: pip install --target ./geoip2 geoip2
# 2. go into the geoip2 folder and zip the folder: zip -r ../geoip2.zip .
# 3. upload the zip file to the stage location
def create_get_country_city_from_ip(session: Session):
    def get_country_city_from_ip(ip: str) -> str:
        path = os.path.join(sys._xoptions["snowflake_import_directory"], "GeoIP2-City.mmdb")
        try:
            with geoip2.database.Reader(path) as reader:
                response = reader.city(ip)
                country = response.country.name
                city = response.city.name
        except Exception:
            country = None
            city = None
        return json.dumps({"country": country, "city": city})

    session.udf.register(
        func=get_country_city_from_ip,
        return_type=StringType(),
        input_types=[StringType()],
        name="get_country_city_from_ip",
        is_permanent=True,
        stage_location="@SUNO_MAXMIND_DB",
        imports=["@SUNO_MAXMIND_DB/geoip2.zip", "@SUNO_MAXMIND_DB/GeoIP2-City.mmdb"],
        replace=True,
    )