import csv
import io
import json
import logging
import sys
from typing import Any, Dict, List, Optional, Union
import subprocess

from flashtext import KeywordProcessor

from arkeo.path import PathManager
from utils.file import (
    extract_extension_from_filename,
    get_delimiter_from_file_extension,
)
from utils.json import load_json
from utils.schema import SchemaTransformer
from utils.text import to_type_int


log = logging.getLogger(__name__)


STATES = "states"
CITIES = "cities"
DATA_NAMES = (STATES, CITIES)

CITY_STATE = "city-state"
DIM_CITY_STATE_ADJACENCY = 25

IX_STATES_BY_NAME = "states_by_name"
IX_CITY_STATE_1X1 = "city_state_1x1"
IX_CITY_STATE_1XN = "city_state_1xn"
INDEX_NAMES = [IX_STATES_BY_NAME, IX_CITY_STATE_1X1, IX_CITY_STATE_1XN]

FILE_EXT = "json"

KEY_MODE = "key_mode"  # for cities
MODE_NORMAL = 0
MODE_COMMON_WORD = 1  # city name is common word
MODE_MULTIPLE_STATES = 2  # city name found in more than one state
MODE_SUPPLEMENTAL = 3  # city name not curated
MODE_DATA_ISSUE = 9  # city-state relationship error
KEY_MODES = [MODE_COMMON_WORD, MODE_MULTIPLE_STATES]

KEY_CITY = "city"
KEY_STATE = "state"
KEY_POPULATION = "population"
KEY_SAVED = "saved"

KEY_FEDERAL_CODE = "feature_code"


class GeonamesManager:
    """read-only geonames data from JSON files"""

    def __init__(self, path_manager: Optional[PathManager] = None):
        self.path: PathManager = path_manager or PathManager()
        self.drive = self.path.drive
        self.data: Dict[str, Any] = {}
        self._idx: Dict[str, Any] = {}
        self.kwp: Dict[str, KeywordProcessor] = {}

        self._load_data()

    @property
    def dir_data(self) -> str:
        return self.path.reference + "geonames/"

    @property
    def dir_arc(self) -> str:
        return self.dir_data + "bak/"

    @property
    def idx(self) -> Dict[str, Any]:
        if not self._idx:
            self._load_indexes()
        return self._idx

    @property
    def path_geonames(self) -> str:
        return self.dir_arc + "US_with_header.txt"

    def _load_indexes(self, index: Optional[str] = None) -> None:
        """load index from JSON files."""
        index_names = [index] if index and index in INDEX_NAMES else INDEX_NAMES
        for index_name in index_names:
            index = {}
            if index_name == IX_STATES_BY_NAME:
                for key, values in self.data[STATES].items():
                    name = values.get("name", "")
                    if name:
                        index[name] = key
                    alias = values.get("alias", "")
                    if alias:
                        index[alias] = key

            elif index_name == IX_CITY_STATE_1X1:
                self._idx[index_name] = {
                    city: data
                    for city, data in self.data[CITIES].items()
                    if data[KEY_MODE] == MODE_NORMAL
                }
            elif index_name == IX_CITY_STATE_1XN:
                self._idx[index_name] = {
                    city: data
                    for city, data in self.data[CITIES].items()
                    if data[KEY_MODE] > MODE_NORMAL
                }

    def _load_data(self, data: Optional[str] = None) -> None:
        """load data from JSON files."""
        data_names = [data] if data and data in DATA_NAMES else DATA_NAMES
        for data_name in data_names:
            file_path = f"{self.dir_data}{data_name}.{FILE_EXT}"
            self.data[data_name] = load_json(self.drive, file_path)

            if data_name == CITIES:
                self._correct_city_mode()

        if not data:
            self.data[CITY_STATE] = {}

    def get(self, geoname: Optional[str] = "") -> Dict[str, Any]:
        if geoname in DATA_NAMES:
            return self.data[geoname]
        if geoname in INDEX_NAMES:
            return self.idx[geoname]

        log.warning(f"{geoname} does not exist.")
        return {}

    def find(
        self, text: str, city_states: Optional[Union[Dict, List]] = None
    ) -> Dict[str, Any]:
        state_cities = self.find_cities(text, city_states)
        states = self.find_states(text)

        new_states = set(states) - set(state_cities.keys())

        for state in new_states:
            state_cities[state] = []
        return state_cities

    def find_cities(
        self, text: str, cities: Optional[Union[Dict, List]] = None
    ) -> Dict[str, Any]:
        state_cities = {}

        # city 1:1 state
        if not (self.kwp and self.kwp.get(CITIES) and len(self.kwp[CITIES]) > 0):
            kwp_1x1 = KeywordProcessor()
            kwp_1x1.add_keywords_from_list(self.get_keywords_cities(cities))
            self.kwp[CITIES] = kwp_1x1

        matches_1x1 = self.kwp[CITIES].extract_keywords(text)
        for city_name in matches_1x1:
            city_data = self.data[CITIES][city_name]
            state_code = next(iter(city_data[STATES]))
            if state_code in state_cities:
                if not city_name in state_cities[state_code]:
                    state_cities[state_code].update({KEY_CITY: city_name})
            else:
                state_cities[state_code] = {KEY_CITY: city_name}

        # city 1:n states
        if self.data[CITY_STATE] and len(self.data[CITY_STATE]) > 0:
            if not (
                self.kwp and self.kwp.get(CITY_STATE) and len(self.kwp[CITY_STATE]) > 0
            ):
                kwp_1xn = KeywordProcessor(case_sensitive=True)
                kwp_1xn.add_keywords_from_list(list(self.data[CITY_STATE].keys()))
                self.kwp[CITY_STATE] = kwp_1xn

            matches_1xn = self.kwp[CITY_STATE].extract_keywords(text, span_info=True)
            for city in matches_1xn:
                city_name, ifrom, ithru = city
                ifrom = ithru + 1
                ithru = ifrom + DIM_CITY_STATE_ADJACENCY

                city_data = self.data[CITY_STATE][city_name]
                city_states = city_data.get(STATES, {})
                for state_code in city_states.keys():
                    state_data = self.data[STATES][state_code]
                    state_name = state_data.get("name", "")
                    state_alias = state_data.get("alias", "")
                    if any(
                        identifier in text[ifrom:ithru]
                        for identifier in [state_code, state_name, state_alias]
                        if identifier
                    ):
                        if state_code in state_cities:
                            if not city_name in state_cities[state_code]:
                                state_cities[state_code].update({KEY_CITY: city_name})
                        else:
                            state_cities[state_code] = {KEY_CITY: city_name}
                        break

        return state_cities

    def find_states(
        self, text: str, states: Optional[Union[Dict, List]] = None
    ) -> List[str]:
        if not (self.kwp and self.kwp.get(STATES) and len(self.kwp[STATES]) > 0):
            kwp = KeywordProcessor()
            kwp.add_keywords_from_list(self.get_keywords_states(states))
            self.kwp[STATES] = kwp
        keyword_matches = self.kwp[STATES].extract_keywords(text)

        matched_states = []
        for match in keyword_matches:
            for code, data in self.data[STATES].items():
                if match in (data["name"], data["alias"]):
                    matched_states.append(code)
                    break
        return matched_states

    def get_keywords_cities(
        self, cities: Optional[Union[Dict, List]] = None
    ) -> List[str]:
        """build keywords list for US cities and/or municipalities of interest"""
        # dict if search other than city name # like state
        cities_dict = self._normalize_cities(cities)
        if not cities_dict:
            return []

        return list(cities_dict.keys())

    def get_keywords_states(
        self, states: Optional[Union[Dict, List]] = None
    ) -> List[str]:
        """build keywords list for US states and/or territories of interest"""

        states_dict = self._normalize_states(states)
        if not states_dict:
            return []

        names = (data.get("name", "").strip() for data in states_dict.values())

        abbrs = (
            self._extract_abbreviation(data.get("alias", "").strip())
            for data in states_dict.values()
        )

        keywords = [keyword for keyword in names if keyword]
        keywords.extend(abbr for abbr in abbrs if abbr)

        return keywords

    def overwrite_all(self) -> bool:
        """overwrite all using defaults."""
        st = SchemaTransformer(self.drive)
        for data_name in DATA_NAMES:
            if not self.overwrite_from_dsv(data_name, transformer=st):
                log.warning(f"failed to overwrite all; break at {data_name}.")
                return False
        return True

    def overwrite_from_dsv(
        self,
        geoname: str,
        file_path: str = "",
        transformer: Optional[SchemaTransformer] = None,
    ) -> bool:
        """overwrite existing data from file."""
        if geoname not in (DATA_NAMES):
            return False

        transformer = transformer or SchemaTransformer(self.drive)
        file_path = file_path or f"{self.path.dir_data}{geoname}.psv"

        schema_name = KEY_STATE if geoname == STATES else KEY_CITY
        schema_path = f"{self.path.dir_data}geo{schema_name}-schema.json"

        if transformer.process_file(file_path, schema_path):
            self._load_data(geoname)
            return True
        return False

    def _extract_dsv_from_geonames(
        self, reference_file: str, check_id: int = 0
    ) -> bool:
        """stage state and city psv files from geonames."""

        file_ext = extract_extension_from_filename(reference_file)
        delimiter = get_delimiter_from_file_extension(file_ext)
        st = SchemaTransformer(self.path.drive)

        # references ~ geonames initial city selection
        content = self.path.drive.read_or_none(self.dir_arc + reference_file)
        buffer = io.StringIO(content)
        references = csv.DictReader(buffer, delimiter=delimiter)

        refs = {}
        for irow, row in enumerate(references, start=1):
            refs[(row[KEY_CITY], row[KEY_STATE])] = {
                KEY_MODE: to_type_int(row[KEY_MODE]),
                KEY_POPULATION: to_type_int(row[KEY_POPULATION]),
                KEY_SAVED: 0,
            }

        try:
            path_cities = self.dir_data + "cities.psv"
            path_states = self.dir_data + "states.psv"
            path_exceptions = self.dir_data + "exceptions.psv"
            with open(self.path_geonames, "r", encoding="utf-8") as infile, open(
                path_cities, "w", encoding="utf-8"
            ) as cities_file, open(
                path_states, "w", encoding="utf-8"
            ) as states_file, open(
                path_exceptions, "w", encoding="utf-8"
            ) as exceptions_file:

                geonames = csv.DictReader(infile, delimiter="\t")
                gcs_wr = csv.writer(cities_file, delimiter=delimiter)
                gss_wr = csv.writer(states_file, delimiter=delimiter)

                cities_found = 0
                states_found = 0

                for irow, row in enumerate(geonames, start=1):
                    geo_name = row["name"].strip()
                    state = row["admin1_code"].strip()
                    pop = to_type_int(row[KEY_POPULATION])

                    fedcd = row[KEY_FEDERAL_CODE]
                    if fedcd == "ADM1":
                        # state
                        subset = []
                        subset.append(state)
                        subset.append(geo_name)

                        airport_code, alias = self._extract_alternate_names(
                            geo_name, row["alternatenames"]
                        )
                        log.debug(
                            f'unused: airport code="{airport_code}"; alias="{alias}".'
                        )
                        subset.append("")
                        subset.append(row["geonameid"])
                        subset.append(pop)
                        subset.append(row["latitude"])
                        subset.append(row["longitude"])
                        subset.append(row["modification_date"])

                        if states_found == 0:
                            schema = load_json(
                                self.path.drive, self.dir_data + "geostate-schema.json"
                            )
                            gss_wr.writerow(st.extract_field_names(schema))

                        gss_wr.writerow(subset)
                        states_found += 1

                    else:
                        # check if city/state/population in references
                        city_state = (geo_name, state)
                        if city_state in refs and pop > 0:
                            ref_data = refs[city_state]
                            mode = ref_data.get(KEY_MODE, 0)
                            ref_pop = ref_data.get(KEY_POPULATION, 0)
                            geonameid = to_type_int(row["geonameid"])
                            if geonameid == check_id:
                                log.warning(
                                    f"line [{irow}]: city {geo_name} with {pop} people vs ref {ref_pop}"
                                )
                            save = abs(ref_pop - pop) / ref_pop < 0.3
                            if save:
                                # row.append(mode)
                                # writer.writerow(row)
                                subset = []
                                subset.append(geo_name)
                                subset.append(mode)
                                subset.append(state)
                                subset.append(geonameid)  # row['geonameid'])
                                subset.append(pop)
                                subset.append(row["latitude"])
                                subset.append(row["longitude"])

                                airport_code, alias = self._extract_alternate_names(
                                    geo_name, row["alternatenames"]
                                )
                                subset.append(airport_code or "")
                                subset.append("")  # LOL, NO.
                                log.debug(f'unused: alias="{alias}".')

                                subset.append(row["modification_date"])

                                if cities_found == 0:
                                    schema = load_json(
                                        self.path.drive,
                                        self.dir_data + "geocity-schema.json",
                                    )
                                    gcs_wr.writerow(st.extract_field_names(schema))

                                gcs_wr.writerow(subset)
                                ref_data[KEY_SAVED] = 1
                                cities_found += 1

                                log.debug(
                                    f"found: [{fedcd}] {geo_name}, {state} (pop: {pop} ~ {ref_pop} ~> {KEY_SAVED if save else ''})"
                                )
                            else:
                                log.warning(
                                    f"unsaved: [{fedcd}] {geo_name}, {state} (pop: {pop} ~ {ref_pop}"
                                )

                gxs_wr = csv.writer(exceptions_file, delimiter=delimiter)

                buffer.seek(0)  # recycle!
                rows = csv.reader(buffer, delimiter=delimiter)
                # as reader vs DictReader to retain line format
                for irow, row in enumerate(rows):
                    if irow == 0:  # headers
                        gxs_wr.writerow(row)
                    else:  # data
                        ref_key = (row[0], row[2])  # city, state
                        if ref_key in refs and refs[ref_key][KEY_SAVED] == 0:
                            gxs_wr.writerow(row)

        except Exception as e:
            print(e)
            sys.exit(1)

        print(f"\nextracted {cities_found} referenced records.")
        return True

    def _extract_abbreviation(self, alias) -> Optional[str]:
        if not alias:
            return None

        for candidate in alias.split(","):
            if candidate.endswith("."):
                return candidate

        return None

    def _extract_alternate_names(self, original: str, alternates: str) -> List[str]:
        if not alternates:
            return "", ""

        alternates = alternates.split(",")
        airport_code = ""
        alias = []
        for alternate in alternates:
            if alternate and alternate.isascii():
                if not airport_code:
                    if (
                        len(alternate) == 3
                        and alternate.isupper()
                        and alternate.isalpha()
                    ):
                        airport_code = alternate
                        continue
                if alternate != original:
                    alias.append(alternate)

        return airport_code or "", ", ".join(alias) if alias else ""

    def _lookup_cities(self, identifiers: List[str]) -> Dict[str, Any]:
        """resolve cities: curated 1x1, stateful ~> O(n) performance 2x."""
        ids = set(identifiers)
        cities_1x1 = self.idx.get(IX_CITY_STATE_1X1, {})
        pk_cities = set(cities_1x1.keys())

        pk_ids = ids & pk_cities
        sc_ids = ids - pk_cities

        cities = self.data.get(CITIES, {})
        pk_ids.update(self._process_stateful_cities(sc_ids, cities))

        refined_pk_ids = {}
        for pk_id in pk_ids:
            if cities_1x1[pk_id].get(KEY_MODE, MODE_DATA_ISSUE) == 0:
                refined_pk_ids[pk_id] = cities_1x1[pk_id]
            else:
                self.data[CITY_STATE][pk_id] = cities_1x1[pk_id]

        return refined_pk_ids

    def _lookup_states(self, identifiers: List[str]) -> Dict[str, Any]:
        """single-pass resolution with O(n) performance."""
        id_set = set(identifiers)

        sk_states = set(self.data[STATES].keys())
        pk_direct = id_set & sk_states

        sk_names = set(self.idx[IX_STATES_BY_NAME].keys())
        pk_name = {self.idx[IX_STATES_BY_NAME][name] for name in id_set & sk_names}

        pk_matches = pk_direct | pk_name

        return {pk: self.data[STATES][pk] for pk in pk_matches}

    def _lookup_geoname_city(
        self, city_name: str, state_code: str = None
    ) -> Dict[str, Any]:
        city_data = {}
        state_code = state_code or ""
        old_state = ""
        old_pop = 0

        reader = csv.DictReader(
            self._get_geoname_rows(city_name, state_code), delimiter="\t"
        )

        for irow, row in enumerate(reader, start=1):
            geoname = row["name"].strip()
            state = row["admin1_code"].strip()
            pop = to_type_int(row[KEY_POPULATION])

            if city_name in city_data:
                states_dict = city_data[city_name].get(STATES, {})
            else:
                states_dict = {}

            if not state_code or state == state_code:
                if state != old_state:
                    old_pop = 0

                if state not in states_dict or pop > old_pop:
                    airport_code, alias = self._extract_alternate_names(
                        geoname, row["alternatenames"]
                    )
                    state_dict = {
                        "geonameid": to_type_int(row["geonameid"]),
                        "population": pop,
                        "coordinates": {
                            "latitude": row["latitude"],
                            "longitude": row["longitude"],
                        },
                        "airport": airport_code,
                        "alias": "",
                        "updated": row["modification_date"],
                    }
                    states_dict[state] = state_dict
                    mode = (
                        MODE_NORMAL if len(states_dict) == 1 else MODE_MULTIPLE_STATES
                    )
                    city_data = {KEY_MODE: mode, STATES: states_dict}
                    old_pop = pop
                    old_state = state

        return city_data

    def _normalize_cities(self, cities: Optional[Union[Dict, List, str]]) -> Dict:
        """normalize cities input to dictionary format with improved string parsing"""

        if not cities:
            self.data[CITY_STATE] = self.idx[IX_CITY_STATE_1XN] or {}
            return self.idx[IX_CITY_STATE_1X1] or {}

        if isinstance(cities, str):
            if cities == CITIES:
                self.data[CITY_STATE] = self.idx[IX_CITY_STATE_1XN] or {}
                return self.idx[IX_CITY_STATE_1X1] or {}
            else:
                if ";" in cities:
                    cities = cities.split(";")  # city, state; ...
                else:
                    cities = cities.split(",")  # city, city, ...

        if isinstance(cities, list):
            if not cities:
                return {}

            if isinstance(cities[0], str):
                return self._lookup_cities(cities) or {}
            else:
                log.warning("Invalid item data type")
                return {}

        if isinstance(cities, dict):
            if cities and not isinstance(next(iter(cities.values())), dict):
                log.warning("Invalid dict structure: values must be dictionaries")
                return {}
            return cities

        log.warning("Invalid data type")
        return {}

    def _normalize_states(self, states: Optional[Union[Dict, List]]) -> Dict:
        """normalize states input to dictionary format"""

        if states is None or len(states) == 0:
            return self.get(STATES) or {}

        if isinstance(states, dict):
            if states and not isinstance(next(iter(states.values())), dict):
                log.warning("invalid dict structure: values must be dictionaries")
                return {}
            return states

        if isinstance(states, list):
            if not states:
                return {}

            if isinstance(states[0], str):
                return self._lookup_states(states) or {}
            else:
                log.warning("invalid item data type")
                return {}

        log.warning("invalid data type")
        return {}

    def _correct_city_mode(self) -> None:
        updates = []
        for city_name, data in self.data[CITIES].items():
            if data.get(KEY_MODE, 0) == 0 and len(data.get(STATES, {})) > 1:
                data[KEY_MODE] = MODE_MULTIPLE_STATES
                updates.append(city_name)

        if updates:
            file_path = f"{self.dir_data}{CITIES}.{FILE_EXT}"
            if not self.drive.write_or_false(file_path, json.dumps(self.data[CITIES])):
                log.warning(
                    f"failed to correct mode for city names in multiple states: {updates}"
                )

    def _get_geoname_rows(self, city_name: str, state_code: str = None) -> io.StringIO:
        # filter to cities AND search in one awk
        if state_code:
            cmd = f"""(head -1 {self.path_geonames}; awk -F'\\t' -v city="{city_name}" -v state="{state_code}" '
            NR>1 && $7=="P" && ($8=="PPL" || $8=="PPLA" || $8=="PPLA2" || $8=="PPLA3" || $8=="PPLA4" || $8=="PPLC") && $2==city && $11==state {{
                if ($15 > max_pop || max_pop=="") {{
                    max_pop = $15
                    best_row = $0
                }}
            }}
            END {{ if (best_row) print best_row }}' {self.path_geonames})"""
        else:
            cmd = f"""(head -1 {self.path_geonames}; awk -F'\\t' -v city="{city_name}" '
            NR>1 && $7=="P" && ($8=="PPL" || $8=="PPLA" || $8=="PPLA2" || $8=="PPLA3" || $8=="PPLA4" || $8=="PPLC") && $2==city {{
                if ($15 > max_pop[$11] || max_pop[$11]=="") {{
                    max_pop[$11] = $15
                    best_row[$11] = $0
                }}
            }}
            END {{ for (state in best_row) print best_row[state] }}' {self.path_geonames})"""

        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)

        if not result.stdout.strip():
            return io.StringIO("")

        return io.StringIO(result.stdout)

    def _process_stateful_cities(self, uncurated: set, cities: dict) -> set:
        promoted_ids = set()

        for uc_id in uncurated:
            uc_id = uc_id.strip()
            if "," in uc_id:
                uc_city, uc_state = uc_id.split(",", 1)
                uc_city = uc_city.strip()
                uc_state = uc_state.strip()

                if uc_city in cities:
                    mode = cities[uc_city].get(KEY_MODE, MODE_DATA_ISSUE)
                    states = cities[uc_city].get(STATES, {})
                    if uc_state in states:
                        if mode == MODE_NORMAL and len(states) == 1:
                            promoted_ids.add(uc_city)
                        else:
                            mode = MODE_DATA_ISSUE if mode == MODE_NORMAL else mode
                            self.data[CITY_STATE][uc_city] = {
                                KEY_MODE: mode,
                                STATES: {uc_state: {}},
                            }
                        continue

                extra = self._lookup_geoname_city(uc_id, uc_state)
                if not extra:
                    extra = {KEY_MODE: MODE_SUPPLEMENTAL, STATES: {uc_state: {}}}
                self.data[CITY_STATE][uc_id] = extra

            else:
                extra = self._lookup_geoname_city(uc_id)
                if not extra:
                    extra = {KEY_MODE: MODE_SUPPLEMENTAL, STATES: {}}
                self.data[CITY_STATE][uc_id] = extra

        return promoted_ids
