import asyncio
import logging
import re
from typing import Optional

import aiohttp
import numpy as np
import pandas as pd
import pycountry

from utils.base_attributes_generation import get_fan_attribute_values
from .enrichment import Enrichment

component_logger = logging.getLogger().getChild("enrichment.enrich_geo_pelias")


class EnrichGeographicsPelias(Enrichment):
    """
    Enrich address from our self-hosted Pelias Geocoding API
    """

    source_attribute_ids = [
        12,  # 'userCountry'
        13,  # 'userCity'
        15,  # 'userAddress'
        17,  # 'userState',
        20,  # 'userLocationLat'
        21,  # 'userLocationLon'
    ]

    result_attribute_ids = [
        54,  # 'enrCountry'
        55,  # 'enrLocality' - matches [95] enrCity
        56,  # 'enrAdminAreaLvl1' - matches [99] enrState
        57,  # 'enrAdminAreaLvl2' (county)
        58,  # 'enrLongitude'
        59,  # 'enrLatitude'
        60,  # 'enrCountryISO2', - matches [96] enrCountryIso2
    ]
    data_gathering_rule = "strict"

    def __init__(
        self,
        collection_id: int,
        schema: str,
        user_id: str,
        default_value: str = "unknown",
        dummy_enrichment=False,
        **kwargs,
    ):
        super().__init__(collection_id, schema, user_id, default_value)
        self.dummy_enrichment = dummy_enrichment
        self.semaphore = None

    async def async_request(self, url, params, session, retry_once=True):
        async with self.semaphore:
            try:
                async with session.get(url, params=params) as response:
                    if response.status == 200:
                        return await response.json()
                    else:
                        component_logger.warning(
                            f"Request to Pelias API {response.request_info.real_url} failed with code {response.status}{' , retrying...' if retry_once else ''}"
                        )
                        if retry_once:
                            return await self.async_request(url, params, session, retry_once=False)
                        else:
                            return None
            except aiohttp.ClientConnectorError as e:
                component_logger.error(f"Couldn't connect to Pelias API!", exc_info=e)
                return None
            except aiohttp.ServerDisconnectedError:
                if retry_once:
                    component_logger.warning("Pelias API rejected the request, retrying...")
                    return await self.async_request(url, params, session, retry_once=False)
                else:
                    component_logger.error("Pelias API rejected the request twice...")
                    return None
            except Exception as e:
                component_logger.error(f"There was a problem querying Pelias API!", exc_info=e)
                return None

    async def _forward_geocode(
        self, session, search_string: str, size=1, boundary_country=None, layers="locality,localadmin,borough"
    ):
        """ Perform forward geocoding (address -> formatted address + lat & lon)
            - params: https://github.com/pelias/documentation/blob/master/search.md
        """
        if search_string is None or search_string == "":
            return None
        params = dict(text=search_string, size=size, lang="en", layers=layers)
        if boundary_country is not None and boundary_country != self.default_value:
            params["boundary.country"] = boundary_country

        return await self.async_request("http://pelias.fansifter.cloud/v1/search", params, session)

    async def _reverse_geocode(self, session, lat: float, lon: float, size=1, layers="locality,localadmin,borough"):
        """ Perform reverse geocoding (lat & lon -> formatted address)
            - params: https://github.com/pelias/documentation/blob/master/reverse.md
        """
        params = {
            "point.lat": lat,
            "point.lon": lon,
            "size": size,
            "layers": layers,
            "lang": "en",
        }

        return await self.async_request("http://pelias.fansifter.cloud/v1/reverse", params, session)

    @staticmethod
    def _normalize_coordinate_value(value) -> Optional[float]:
        if value is None:
            return None
        # removing all non-digit characters except . and -
        coord_value = re.sub(r"[^-\.\d]", "", str(value))
        try:
            normalized_value = float(coord_value)
        except ValueError:
            return None
        return normalized_value

    @staticmethod
    def _normalize_country(row, col) -> Optional[str]:
        """ Uses pycountry module to map ISO2 and ISO3 to full country name. Also fixes capitalization.
         - Has special handling for UK (since UK is not ISO2, but people often confuse it as such).
         - Switches "None" string to actual None (so it would be bucketed together with "unknown")
         """
        x = row[col]

        if x is None:
            return x
        else:
            x = x.strip()
        if x.lower() == "uk":
            x = "gb"
        if x.lower() == "none":
            return None  # we can't parse this further anyway

        long_name = None
        try:
            if len(x) < 2:
                country = None
            elif len(x) == 2:
                country = pycountry.countries.get(alpha_2=x).name
            elif len(x) == 3:
                country = pycountry.countries.get(alpha_3=x).name
            else:
                country = pycountry.countries.get(official_name=x) or pycountry.countries.get(name=x)
            if country is not None:
                long_name = country.common_name if hasattr(country, "common_name") else country.name
        except Exception:
            pass

        return long_name

    @staticmethod
    def _get_country_iso2(row, col) -> Optional[str]:
        """ uses pycountry module to map full country name to ISO2
        (required for fb, google audience files generation)"""
        x = row[col]
        if x is None:
            return x
        else:
            x = x.strip()
        try:
            iso2_name = pycountry.countries.get(name=x).alpha_2
        except Exception as e:
            iso2_name = None
        return iso2_name

    @staticmethod
    def _normalize_region(row, col, col2) -> Optional[str]:
        """ If 2-3 characters, return the longer state name. Otherwise, return None
        """
        x = row[col]
        country_iso2 = row[col2]
        if x is None:
            return x
        else:
            x = x.strip()

        if x.lower() in ["none", "undefined"]:
            return None  # we can't parse this further anyway

        """ See if we perhaps get country code. We can have pycountry subdivisions for other than US as well.
        Otherwise, we have to assume it's US :shrug: """
        if not country_iso2:
            country_iso2 = "US"

        long_name = None
        try:
            # In some countries, state code can be only one letter (e.g. Ireland; Spain) or three letters (e.g. UK)
            long_name = pycountry.subdivisions.get(code=f"{country_iso2}-{x.upper()}").name
        except Exception:
            pass

        return long_name

    def is_na(self, value):
        return pd.isna(value) or value == self.default_value

    async def _enrich_with_geocoding_data(self, df):
        self.semaphore = asyncio.Semaphore(100)  # allowing 100 concurrent async requests to Pelias API

        async with aiohttp.ClientSession(
            # 3 hours timeout, should be enough to enrich all rows even in the hugest files
            timeout=aiohttp.ClientTimeout(total=3 * 60 * 60),
        ) as session:
            enriched_indices = []
            tasks = []

            for i, row in df.iterrows():
                # if no source data is present in this row
                if np.all([pd.isna(row.get(str(col))) for col in self.source_attribute_ids]):
                    continue

                # if latitude & longitude are present - perform reverse geocoding (lat & lon -> formatted address)
                if pd.notna(row.get("20")) and pd.notna(row.get("21")):
                    lat = self._normalize_coordinate_value(row.get("20"))
                    lon = self._normalize_coordinate_value(row.get("21"))
                    tasks.append(asyncio.ensure_future(self._reverse_geocode(session, lat=lat, lon=lon)))
                    enriched_indices.append(i)
                # otherwise, perform forward geocoding (address -> formatted address + lat & lon)
                else:
                    search_string = ", ".join(
                        [
                            str(el)
                            #                Address,          City,         State,       Country
                            for el in [row.get("15"), row.get("13"), row.get("17"), row.get("12")]
                            if pd.notna(el)
                        ]
                    )
                    if pd.notna(row.get("13")) or pd.notna(row.get("15")):  # City or Address is not None
                        layers = "borough,localadmin,locality,county,macrocounty,region,macroregion,country"
                    elif pd.notna(row.get("17")):  # State is not None
                        layers = "region,macroregion,country"
                    else:  # only Country is not None
                        layers = "dependency,country"
                    tasks.append(
                        asyncio.ensure_future(
                            self._forward_geocode(
                                session, search_string, boundary_country=row.get("enrCountryISO2"), layers=layers
                            )
                        )
                    )
                    enriched_indices.append(i)

            geocodes = await asyncio.gather(*tasks)

            for i, geocode in zip(enriched_indices, geocodes):
                if geocode is None or len(geocode["features"]) == 0:
                    continue
                place = geocode["features"][0]

                properties = place["properties"]

                country_iso = properties.get("country_a") or properties.get("dependency_a") or self.default_value
                try:
                    if len(country_iso) == 3:
                        country_iso = pycountry.countries.get(alpha_3=country_iso).alpha_2
                except Exception:
                    pass

                # It may happen that Pelias guessed a completely wrong place in a different country
                # compared to the country obtained from pycountry library, which we trust more
                # So we perform a sanity check to make sure we're not overwriting good data with wrong data
                if (
                    pd.isna(df.at[i, "enrCountryISO2"])
                    or df.at[i, "enrCountryISO2"] == self.default_value
                    or df.at[i, "enrCountryISO2"] == country_iso
                ):
                    df.at[i, "enrCountryISO2"] = country_iso
                    df.at[i, "enrCountry"] = (
                        properties.get("country") or properties.get("dependency") or self.default_value
                    )
                    df.at[i, "enrAdminAreaLvl1"] = (
                        properties.get("region") or properties.get("macroregion") or self.default_value
                    )
                    df.at[i, "enrAdminAreaLvl2"] = (
                        properties.get("county") or properties.get("macrocounty") or self.default_value
                    )
                    df.at[i, "enrLocality"] = (
                        properties.get("locality")
                        or properties.get("localadmin")
                        or properties.get("borough")
                        or self.default_value
                    )
                    longitude, latitude = place["geometry"]["coordinates"]
                    df.at[i, "enrLongitude"] = longitude
                    df.at[i, "enrLatitude"] = latitude
                    # potentially we can also save 'street', 'postalcode', 'label' (formatted single-line address)

    async def enrich_with_geo_coordinates_where_missing(self, df):
        """
        We could have used only pycountry for enriching enrCountry and enrAdminAreaLvl1 without Pelias data,
        but then we need to get lat & lon of these countries separately using enriched data
        """
        self.semaphore = asyncio.Semaphore(100)  # allowing 100 concurrent async requests to Pelias API

        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=3 * 60 * 60),) as session:
            enriched_indices = []
            tasks = []

            for i, row in df.iterrows():
                if self.is_na(row.get("enrCountry")) or (
                    not self.is_na(pd.notna(row.get("enrLongitude"))) and not self.is_na(row.get("enrLatitude"))
                ):
                    continue
                search_string = ", ".join(
                    [str(el) for el in [row.get("enrAdminAreaLvl1"), row.get("enrCountry")] if not self.is_na(el)]
                )
                layers = (
                    "region,macroregion,locality,borough"
                    if not self.is_na(row.get("enrAdminAreaLvl1"))
                    else "dependency,country"
                )
                tasks.append(asyncio.ensure_future(self._forward_geocode(session, search_string, layers=layers)))
                enriched_indices.append(i)

            geocodes = await asyncio.gather(*tasks)

            for i, geocode in zip(enriched_indices, geocodes):
                if geocode is None or len(geocode["features"]) == 0:
                    continue
                longitude, latitude = geocode["features"][0]["geometry"]["coordinates"]
                df.at[i, "enrLongitude"] = longitude
                df.at[i, "enrLatitude"] = latitude

    def _run(self):
        source_df = self.available_data["source"]

        apply = (
            source_df.parallel_apply
            if len(source_df) > 1000 and hasattr(source_df, "parallel_apply")
            else source_df.apply
        )

        source_df["enrCountry"] = apply(self._normalize_country, axis=1, args=["12"])
        source_df["enrCountryISO2"] = apply(self._get_country_iso2, axis=1, args=["enrCountry"])
        source_df["enrAdminAreaLvl1"] = apply(self._normalize_region, axis=1, args=["17", "enrCountryISO2"])

        loop = asyncio.get_event_loop()
        loop.run_until_complete(self._enrich_with_geocoding_data(source_df))
        loop.run_until_complete(self.enrich_with_geo_coordinates_where_missing(source_df))

        source_df.set_index("row_id", inplace=True)

        source_df = source_df[["fan_id", *[col for col in source_df.columns if col.startswith("enr")]]]

        for column in source_df.columns[1:]:
            source_df[column].fillna(value=self.default_value, inplace=True)

        self.result_df = source_df

    def _generate_fan_attributes(self):
        """ Getting a list of attributes that are present in this enrichment"""
        condition = self.system_fields["a_id"].isin(self.result_attribute_ids)
        attributes_df = self.system_fields[condition][["a_system_name", "a_id"]]
        attributes_df.rename(columns={"a_system_name": "name", "a_id": "attribute_id"}, inplace=True)

        fan_attributes_df = get_fan_attribute_values(
            source_df=self.result_df, attributes_df=attributes_df, collection_id=self.enrichment_collection_id
        )
        self.fan_attribute = fan_attributes_df