import logging
import re
from functools import lru_cache
from typing import Any, Dict, List, Optional, Union

from flashtext import KeywordProcessor

from ref.geonames import GeonamesManager
from ref.geonames import CITIES, STATES, FILE_EXT
from utils.json import load_json


log = logging.getLogger(__name__)


class KeywordsManager:

    def __init__(self, geonames_manager: Optional[GeonamesManager] = None):
        self.gm = geonames_manager or GeonamesManager()
        self.path = self.gm.path
        self.drive = self.path.drive

        self.ix_states_by_name: Dict[str, Any] = {}
        self.kwp: Dict[str, KeywordProcessor] = {}

        self._load_data()

    def _load_data(self):
        filename = f"{self.gm.dir_data}{self.gm.IX_STATES_BY_NAME}.{FILE_EXT}"
        self.ix_states_by_name = load_json(self.drive, filename)

        # {ST: {name, alias}}
        self.kwp[STATES] = self._get_kwp_for_states()
        # {name: flag, {STs: {ST: local-atts}}
        self.kwp[CITIES] = self._get_kwp_for_locations(CITIES)

    def x_get_kwp_for_states(
        self, states: Optional[Union[Dict, List]] = None
    ) -> Optional[KeywordProcessor]:
        """build kwp for US states/territories"""
        if states is None:
            states = self.gm.get(STATES)
        elif isinstance(states, List):
            if isinstance(states[0], str):
                states = self._get_matching_states(states)
            elif not isinstance(states[0], List):
                return None  # invalid element type
        elif not isinstance(states, Dict):
            return None  # invalid type

        if not states:
            return None

        ikws = 0
        kwp = KeywordProcessor()
        for _, data in states.items():
            name = data.get("name", "").strip()
            if name:
                kwp.add_keyword(name)
                ikws += 1

            alias = data.get("alias", "").strip()
            abbreviation = self.gm._extract_abbreviation(alias)
            if abbreviation:
                kwp.add_keyword(abbreviation)
                ikws += 1

        if ikws > 0:
            return kwp

        log.warning("no state keywords added.")
        return None

    def x_get_kwp_for_locations(
        self, locations: Optional[Union[Dict, List, str]] = None
    ) -> Optional[KeywordProcessor]:
        """build kwp for US cities and/or municipalities of interest"""
        if not locations:
            # set of data as default
            locations = self.gm.get(CITIES)
        elif isinstance(locations, str):
            if locations == CITIES:
                # set of data selected by string parameter
                locations = self.gm.get(locations)
            else:
                if locations.find(";"):
                    # string of city names possibly with states delimited by ','
                    locations = locations.split(";")
                else:
                    # string of names without states
                    locations = locations.split(",")
                locations = self._get_matching_locations(locations)
        elif isinstance(locations, List):
            if isinstance(locations[0], str):
                # list of city names
                locations = self._get_matching_locations(locations)
            else:
                # invalid element type
                return None
        elif not isinstance(locations, Dict):
            # invalid type
            return None

        ikws = 0
        kwp = KeywordProcessor()

        for key, data in locations.items():
            # skip if any flag
            if data.get("flag", 0) > 0:
                continue

            # skip if multiple states, flag = 2 (human error)
            states = data.get("states", {})
            state_count = len(states)
            if states and state_count == 1:
                kwp.add_keyword(key)
                ikws += 1
            elif state_count > 1:
                log.info(f"FYI: city {key} has multiple states")

        if ikws > 0:
            log.info(f"{ikws} keywords added.")
            return kwp

        log.warning("no keywords added.")
        return None

    @staticmethod
    @lru_cache(maxsize=128)
    def get_compiled_pattern(
        city_name: str, state: tuple, max_distance: int = 100
    ) -> re.Pattern:
        state_pattern = "|".join(filter(None, state))
        return re.compile(
            rf"{re.escape(city_name)}(?=.{{0,{max_distance}}}(?:{state_pattern}))"
        )

    def find_all(self, text: str) -> Dict[str, str]:
        by_state = {}
        by_city = self.find_all_city_states(text)
        for city in by_city:
            state = by_city[city]
            if city.startswith("["):
                city_name = ""
            else:
                city_name = city
            if state in by_state:
                by_state[state].append(city_name)
            else:
                by_state[state] = [city_name]
        return by_state

    def find_all_city_states(self, text: str) -> Dict[str, str]:
        found_states = self.find_states(text)
        # TODO: verify if outer brackets needed
        found_locations = {self.find_locations(CITIES, text)}

        # filter out states that conflict with found locations
        filtered_states = [
            state
            for state in found_states
            if not any(
                state == self.ix_states_by_name.get(city) or state in city_state
                for city, city_state in found_locations.items()
            )
        ]

        # add remaining states to locations
        found_locations.update({f"[{state}]": state for state in filtered_states})

        return found_locations

    def find_states(self, text: str, kwp: KeywordProcessor = None) -> List[str]:
        kwp = kwp or self.kwp[STATES]
        unique_states = []
        for found in kwp.extract_keywords(text):
            if found in self.ix_states_by_name:
                unique_states.append(self.ix_states_by_name[found])
        return list(set(unique_states))

    def find_locations(
        self, location: str, text: str, kwp: KeywordProcessor = None
    ) -> Dict[str, str]:
        kwp = kwp or self.kwp[location]
        locations = self.gm.get(location)

        unique_locations = {}
        if kwp:
            for unique_name in kwp.extract_keywords(text):
                unique_locations[unique_name] = locations[unique_name]

        # non-unique locations that exist in multiple states and/or flagged
        flagged_locations = self._find_locations_flagged(location, text)
        return {**unique_locations, **flagged_locations}

    def x_get_matching_locations(self, selected: List[str]) -> Dict[str, Any]:
        """single-pass resolution with O(n) performance."""
        sk_locations = set(self.gm.states.keys())
        sk_names = set(self.ix_states_by_name.keys())
        sv_selected = set(selected)

        pk_direct = sv_selected & sk_locations
        pk_name = {self.ix_states_by_name[name] for name in sv_selected & sk_names}
        pk_matches = pk_direct | pk_name

        return {pk: self.gm.states[pk] for pk in pk_matches}

    def x_get_matching_states(self, selected: List[str]) -> Dict[str, Any]:
        """single-pass resolution with O(n) performance."""
        sk_states = set(self.gm.states.keys())
        sk_names = set(self.ix_states_by_name.keys())
        sv_selected = set(selected)

        pk_direct = sv_selected & sk_states
        pk_name = {self.ix_states_by_name[name] for name in sv_selected & sk_names}
        pk_matches = pk_direct | pk_name

        return {pk: self.gm.states[pk] for pk in pk_matches}

    def _find_locations_flagged(self, location: str, text: str) -> Dict[str, str]:
        locations = self.gm.get(location)

        # simple filter bc only one flag value
        found = {}
        for location, locaton_data in locations.items():
            states = locaton_data[STATES]
            if isinstance(states, Dict):
                for state, state_data in states.items():
                    one_row = self._get_subset_states([state])[0]
                    state_names = (one_row[0], one_row[1], one_row[2])
                    state_pattern = KeywordsManager.get_compiled_pattern(
                        location, state_names
                    )
                    if state_pattern.search(text):
                        found[location] = state
                        break

        return found

    def _get_subset_states(self, selected: List[str]) -> List[List[str]]:
        select_set = set(selected)
        return [row for row in self._states_list if row[0] in select_set]
