"""
Resolve Mauritian bus-stop names to coordinates so itineraries can be drawn
on a map. mauritius-buses.com only gives us stop *names*, so we resolve them in
three tiers, cheapest/most-accurate first:

  1. local cache (data/coords_cache.json)
  2. an OpenStreetMap index fetched once via Overpass — every highway=bus_stop
     node plus every named place (town/village/suburb) in Mauritius. Matched by
     normalised name, then by leading place token, then fuzzy.
  3. Nominatim geocoding ("<name>, Mauritius"), rate-limited and cached.

A small set of major hubs is hard-seeded so the map always has good anchors.
"""

import difflib
import json
import os
import re
import threading
import time
import unicodedata

import requests

DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
OSM_CACHE = os.path.join(DATA_DIR, "osm_index.json")
COORDS_CACHE = os.path.join(DATA_DIR, "coords_cache.json")
# Free-text address searches are cached too: addresses don't move, and this
# spares the public Photon/Nominatim endpoints (and the user) a live round-trip
# on every repeat search. Keyed by the normalised query; only non-empty results
# are stored (an empty result is usually transient — offline, or rate-limited).
SEARCH_CACHE = os.path.join(DATA_DIR, "search_cache.json")

UA = (
    "MauritiusBusPlanner/1.0 (personal route-planning demo; "
    "contact: syadone@gmail.com)"
)
OVERPASS = "https://overpass-api.de/api/interpreter"
NOMINATIM = "https://nominatim.openstreetmap.org/search"
# Photon is an OSM-based geocoder that, unlike the public Nominatim, permits
# server-side use. The public Nominatim blocks datacenter IPs, so it's
# unreachable from the HF Space — Photon is our primary online geocoder there,
# with Nominatim kept as a local/dev fallback.
PHOTON = "https://photon.komoot.io/api"
BBOX = (-20.6, 57.2, -19.9, 57.9)  # S, W, N, E — main island of Mauritius
PHOTON_BBOX = "57.0,-20.6,57.9,-19.9"  # minLon,minLat,maxLon,maxLat (Photon order)

# Accurate anchors for the busiest hubs (lat, lon).
SEED = {
    "port louis victoria square": (-20.1626, 57.4986),
    "port louis": (-20.1609, 57.5012),
    "port louis immigration square": (-20.1568, 57.4994),
    "port louis transportation centre": (-20.1606, 57.4990),
    "curepipe": (-20.3188, 57.5260),
    "curepipe ian palach north": (-20.3140, 57.5230),
    "quatre bornes": (-20.2654, 57.4791),
    "rose hill": (-20.2419, 57.4680),
    "rose hill place margeot": (-20.2410, 57.4669),
    "beau bassin": (-20.2289, 57.4660),
    "vacoas": (-20.2980, 57.4790),
    "phoenix": (-20.2870, 57.4980),
    "mahebourg": (-20.4081, 57.7000),
    "flacq": (-20.1900, 57.7150),
    "centre de flacq": (-20.1900, 57.7150),
    "goodlands": (-20.0380, 57.6500),
    "triolet": (-20.0560, 57.5470),
    "grand baie": (-20.0130, 57.5800),
    "souillac": (-20.5170, 57.5180),
    "rivière du rempart": (-20.1060, 57.6840),
    "pamplemousses": (-20.1040, 57.5700),
    "saint pierre": (-20.2200, 57.5210),
    "moka": (-20.2200, 57.4960),
    "ebene": (-20.2430, 57.4900),
    "bambous": (-20.2620, 57.4060),
    "albion": (-20.2180, 57.4010),
}

_lock = threading.Lock()
_osm_index = None  # {normalised_name: [lat, lon]}
_osm_keys = None  # list of keys for fuzzy matching
_coords_cache = None  # {original_name: [lat, lon] or None}
_search_cache = None  # {norm(query): [{name, coord}, ...]}
_last_nominatim = 0.0


# --------------------------------------------------------------------------- #
def _norm(s):
    s = unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode()
    s = s.lower()
    s = re.sub(r"\(.*?\)", " ", s)  # drop parenthetical detail
    s = re.sub(r"[^a-z0-9 ]", " ", s)
    s = re.sub(r"\s+", " ", s).strip()
    return s


def _load_coords_cache():
    global _coords_cache
    if _coords_cache is not None:
        return _coords_cache
    if os.path.exists(COORDS_CACHE):
        with open(COORDS_CACHE, encoding="utf-8") as fh:
            _coords_cache = json.load(fh)
    else:
        _coords_cache = {}
    return _coords_cache


def _save_coords_cache():
    os.makedirs(DATA_DIR, exist_ok=True)
    tmp = COORDS_CACHE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(_coords_cache, fh, ensure_ascii=False)
    os.replace(tmp, COORDS_CACHE)


def _load_search_cache():
    global _search_cache
    if _search_cache is not None:
        return _search_cache
    if os.path.exists(SEARCH_CACHE):
        with open(SEARCH_CACHE, encoding="utf-8") as fh:
            _search_cache = json.load(fh)
    else:
        _search_cache = {}
    return _search_cache


def _save_search_cache():
    os.makedirs(DATA_DIR, exist_ok=True)
    tmp = SEARCH_CACHE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(_search_cache, fh, ensure_ascii=False)
    os.replace(tmp, SEARCH_CACHE)


def _build_osm_index(force=False):
    """Fetch (once, cached) named bus stops + places in Mauritius from OSM."""
    global _osm_index, _osm_keys
    if _osm_index is not None and not force:
        return _osm_index

    raw = None
    if os.path.exists(OSM_CACHE) and not force:
        with open(OSM_CACHE, encoding="utf-8") as fh:
            raw = json.load(fh)
    else:
        s, w, n, e = BBOX
        query = f"""[out:json][timeout:90];
(
  node["highway"="bus_stop"]["name"]({s},{w},{n},{e});
  node["place"~"city|town|village|suburb|hamlet|neighbourhood|locality"]["name"]({s},{w},{n},{e});
);
out center;"""
        resp = requests.post(
            OVERPASS, data={"data": query}, headers={"User-Agent": UA}, timeout=120
        )
        resp.raise_for_status()
        raw = resp.json()
        os.makedirs(DATA_DIR, exist_ok=True)
        with open(OSM_CACHE, "w", encoding="utf-8") as fh:
            json.dump(raw, fh, ensure_ascii=False)

    index = {}
    for el in raw.get("elements", []):
        name = (el.get("tags") or {}).get("name")
        lat = el.get("lat") or (el.get("center") or {}).get("lat")
        lon = el.get("lon") or (el.get("center") or {}).get("lon")
        if not name or lat is None or lon is None:
            continue
        key = _norm(name)
        if key and key not in index:  # first writer wins (bus stops listed first)
            index[key] = [lat, lon]
    # Seed hubs override / fill in.
    for k, (la, lo) in SEED.items():
        index[k] = [la, lo]

    _osm_index = index
    _osm_keys = list(index.keys())
    return index


def _osm_lookup(name):
    index = _build_osm_index()
    key = _norm(name)
    if not key:
        return None
    if key in index:
        return index[key]
    # Leading place token: try progressively shorter prefixes.
    tokens = key.split()
    for cut in range(len(tokens), 0, -1):
        prefix = " ".join(tokens[:cut])
        if prefix in index:
            return index[prefix]
    # Fuzzy match.
    match = difflib.get_close_matches(key, _osm_keys, n=1, cutoff=0.86)
    if match:
        return index[match[0]]
    return None


def _nominatim_lookup(name):
    global _last_nominatim
    clean = re.sub(r"\(.*?\)", "", name).strip()
    params = {
        "q": f"{clean}, Mauritius",
        "format": "json",
        "limit": 1,
        "countrycodes": "mu",
    }
    with _lock:  # serialise to honour the 1 req/s policy
        wait = 1.1 - (time.time() - _last_nominatim)
        if wait > 0:
            time.sleep(wait)
        _last_nominatim = time.time()
    try:
        resp = requests.get(
            NOMINATIM, params=params, headers={"User-Agent": UA}, timeout=20
        )
        data = resp.json()
        if data:
            return [float(data[0]["lat"]), float(data[0]["lon"])]
    except Exception:
        pass
    return None


def _photon_results(q, limit=5):
    """Free-text Photon search -> [{name, coord}], constrained to Mauritius.
    Works from cloud servers (the public Nominatim does not)."""
    clean = re.sub(r"\(.*?\)", "", q or "").strip()
    if not clean:
        return []
    params = {"q": clean, "limit": limit, "lang": "en", "bbox": PHOTON_BBOX}
    out = []
    try:
        resp = requests.get(
            PHOTON, params=params, headers={"User-Agent": UA}, timeout=20
        )
        for f in resp.json().get("features", []):
            geom = (f.get("geometry") or {}).get("coordinates")
            if not geom or len(geom) < 2:
                continue
            p = f.get("properties") or {}
            # bbox already limits to the main island; double-check country.
            cc = (p.get("countrycode") or "").upper()
            if cc and cc != "MU":
                continue
            lon, lat = float(geom[0]), float(geom[1])
            label = ", ".join(
                x for x in (p.get("name"), p.get("city") or p.get("county")) if x
            )
            out.append({"name": label or p.get("name") or clean, "coord": [lat, lon]})
    except Exception:
        pass
    return out


def _photon_lookup(name):
    res = _photon_results(name, 1)
    return res[0]["coord"] if res else None


def _relaxed_variants(name):
    """Looser forms of a stop name OSM is more likely to know, e.g.
    'Brabant Street - Venus' -> 'Brabant Street'."""
    base = re.sub(r"\(.*?\)", "", name).strip()
    out = [base]
    if " - " in base:
        out.append(base.split(" - ")[0].strip())  # drop the "- sub-area" suffix
    return [v for v in dict.fromkeys(out) if v and v != name]


def resolve(name, allow_nominatim=True):
    """Return [lat, lon] for a stop name, or None. Non-null results are cached;
    nulls are re-attempted with relaxed variants so sparsely-named stops (which
    OSM doesn't know verbatim) still resolve."""
    cache = _load_coords_cache()
    if cache.get(name):  # trust only successful cached results
        return cache[name]

    variants = _relaxed_variants(name)
    coord = _osm_lookup(name)
    if coord is None:
        for v in variants:  # free: OSM index on looser names
            coord = _osm_lookup(v)
            if coord:
                break
    if coord is None and allow_nominatim:
        # Photon first (works from the Space); Nominatim as a local fallback.
        for q in [name] + variants:
            coord = _photon_lookup(q)
            if coord:
                break
        if coord is None:
            for q in [name] + variants:
                coord = _nominatim_lookup(q)
                if coord:
                    break

    cache[name] = coord
    _save_coords_cache()
    return coord


def resolve_many(names):
    """Resolve a list of names, returning {name: [lat, lon] or None}."""
    return {n: resolve(n) for n in dict.fromkeys(names)}


# --------------------------------------------------------------------------- #
# Free-text destination geocoding (with correction).
#
# OSM coverage in Mauritius is thin, so a literal address often fails. We relax
# the query — strip subdivision prefixes OSM rarely knows ("morcellement",
# "cité", ...) and progressively drop leading tokens — until something matches,
# e.g. "morcellement Black Rock, Tamarin" -> "Black Rock, Tamarin".
# --------------------------------------------------------------------------- #
NOISE_WORDS = {
    "morcellement",
    "morc",
    "cite",
    "residence",
    "resid",
    "camp",
    "hameau",
    "lotissement",
    "allee",
    "impasse",
    "ave",
    "av",
}


def _short_name(d):
    """A concise label from a Nominatim result's display_name (drops country)."""
    parts = [p.strip() for p in d.get("display_name", "").split(",") if p.strip()]
    if parts and _norm(parts[-1]) in ("mauritius", "maurice"):
        parts = parts[:-1]
    return ", ".join(parts[:3]) if parts else d.get("display_name", "")


def _nominatim_search(q, limit=5):
    """Free-text Nominatim search -> [{name, coord}], rate-limited."""
    global _last_nominatim
    clean = re.sub(r"\(.*?\)", "", q).strip()
    params = {
        "q": f"{clean}, Mauritius",
        "format": "json",
        "limit": limit,
        "countrycodes": "mu",
        "addressdetails": 1,
    }
    with _lock:
        wait = 1.1 - (time.time() - _last_nominatim)
        if wait > 0:
            time.sleep(wait)
        _last_nominatim = time.time()
    out = []
    try:
        resp = requests.get(
            NOMINATIM, params=params, headers={"User-Agent": UA}, timeout=20
        )
        for d in resp.json():
            out.append(
                {"name": _short_name(d), "coord": [float(d["lat"]), float(d["lon"])]}
            )
    except Exception:
        pass
    return out


def _query_variants(q):
    """Ordered, most-specific-first query relaxations for `q`."""
    q = q.strip()
    variants = [q]
    parts = [p.strip() for p in q.split(",") if p.strip()]
    # Strip noise words from the leading component (the subdivision name).
    if parts:
        head = [t for t in parts[0].split() if _norm(t) not in NOISE_WORDS]
        stripped = " ".join(head)
        if stripped and stripped != parts[0]:
            variants.append(", ".join([stripped] + parts[1:]))
        # Progressively drop leading tokens of the head component.
        toks = parts[0].split()
        for cut in range(1, len(toks)):
            tail_head = " ".join(toks[cut:])
            variants.append(", ".join([tail_head] + parts[1:]))
    # Finally, just the trailing components (e.g. the village/town).
    if len(parts) > 1:
        variants.append(", ".join(parts[1:]))
    # De-dupe, preserving order.
    seen, out = set(), []
    for v in variants:
        k = _norm(v)
        if k and k not in seen:
            seen.add(k)
            out.append(v)
    return out


def geocode_search(q, limit=5):
    """Resolve a free-text destination address to OSM-corrected candidates:
    [{name, coord}], best/most-specific first. Tries the literal query, then
    relaxes it until OSM returns matches."""
    if not (q or "").strip():
        return []

    cache = _load_search_cache()
    ckey = _norm(q)
    if ckey and cache.get(ckey):  # served from a previous successful search
        return cache[ckey][:limit]

    out, seen = [], set()

    def add(name, coord):
        key = (round(coord[0], 5), round(coord[1], 5))
        if key not in seen:
            seen.add(key)
            out.append({"name": name, "coord": coord})

    for variant in _query_variants(q):
        loc = _osm_lookup(variant)
        if loc:
            add(variant, loc)
        # Photon is the primary online geocoder (reachable from the Space and
        # rich in POIs like "Caudan Waterfront"); Nominatim is a local fallback.
        for cand in _photon_results(variant, limit):
            add(cand["name"], cand["coord"])
        if len(out) < limit:
            for cand in _nominatim_search(variant, limit):
                add(cand["name"], cand["coord"])
        if out:  # an earlier (more specific) variant matched
            break

    if ckey and out:  # cache only successful searches (addresses don't move)
        cache[ckey] = out
        _save_search_cache()
    return out[:limit]


if __name__ == "__main__":
    idx = _build_osm_index(force=True)
    print(f"OSM index: {len(idx)} named features")
    for t in [
        "Port Louis - Victoria Square",
        "Curepipe Ian Palach North",
        "Pailles (Grewals)",
        "Beau Bassin - Nid d'Hirondelle",
        "Gros Cailloux",
        "Trianon (Margarine Industry Limited)",
    ]:
        print(f"  {t!r:48} -> {resolve(t)}")