'''
Simple Geolocation with FastAPI
'''
import dataclasses
from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
from typing import Annotated, Optional, Union

import geoip2.database
from fastapi import FastAPI, Path
from pydantic import BaseModel

app = FastAPI()

GEOLITE2_ASN_DB = '/usr/local/share/GeoIP/GeoLite2-ASN.mmdb'
GEOLITE2_CITY_DB = '/usr/local/share/GeoIP/GeoLite2-City.mmdb'

@dataclasses.dataclass
class Locality(BaseModel):
    '''
    Locality data
    '''
    city: Optional[str]
    country: Optional[str]
    continent: Optional[str]
    is_eu: bool

@dataclasses.dataclass
class GeoLocation(BaseModel):
    '''
    Geolocation data model
    '''
    ip: Union[IPv6Address,IPv4Address]
    asn: Optional[int]
    asn_org: Optional[str]
    network: Union[IPv6Network,IPv4Network,None]
    locality: Locality

@app.get("/{ipaddress}")
async def root(ipaddress: Annotated[Union[IPv4Address,IPv6Address],
                                    Path(title="The IPAddress to geolocate")]
               ) -> GeoLocation:
    '''
    Look up geolocation for ip in path parameter
    '''
    with (geoip2.database.Reader(GEOLITE2_ASN_DB) as reader_asn,
            geoip2.database.Reader(GEOLITE2_CITY_DB) as reader_city):
        asn_data = reader_asn.asn(ipaddress)
        city_data = reader_city.city(ipaddress)

        return GeoLocation(
            ip=ipaddress,
            asn=asn_data.autonomous_system_number,
            asn_org=asn_data.autonomous_system_organization,
            network=asn_data.network,
            locality=Locality(
                city=city_data.city.name,
                country=city_data.country.iso_code,
                continent=city_data.continent.code,
                is_eu=city_data.country.is_in_european_union
            )
        )