"""Romansh logic: location -> idiom lookup and translation.

Uses the ZurichNLP Romansh NLLB-200 1.3B model (University of Zurich / Lia Rumantscha)
to translate a set of everyday phrases into the local Romansh idiom.
"""

import json
import math
import re
import threading
import unicodedata
from functools import lru_cache
from pathlib import Path

import requests

MODEL_ID = "ZurichNLP/romansh-nllb-1.3b-ct2"

IDIOMS = {
    "rm-sursilv": "Sursilvan",
    "rm-sutsilv": "Sutsilvan",
    "rm-surmiran": "Surmiran",
    "rm-puter": "Puter",
    "rm-vallader": "Vallader",
    "rm-rumgr": "Rumantsch Grischun",
}

# "Grüezi!" -> "Allegra!" in Pledari Grond (RG, Surmiran); understood in every idiom.
GREETING = "Allegra!"


# (names, lat, lon, idiom). The first name is used for display. None = Romansh is not the local language.
PLACES = [
    # Sursilvan
    (["Disentis/Mustér", "Disentis", "Mustér", "Muster"], 46.703, 8.853, "rm-sursilv"),
    (["Sedrun", "Tujetsch"], 46.680, 8.775, "rm-sursilv"),
    (["Sumvitg", "Somvix"], 46.728, 8.938, "rm-sursilv"),
    (["Trun"], 46.743, 8.988, "rm-sursilv"),
    (["Breil/Brigels", "Breil", "Brigels"], 46.770, 9.064, "rm-sursilv"),
    (["Ilanz/Glion", "Ilanz", "Glion"], 46.774, 9.205, "rm-sursilv"),
    (["Vella", "Lumnezia", "Lugnez"], 46.717, 9.170, "rm-sursilv"),
    (["Lumbrein"], 46.684, 9.137, "rm-sursilv"),
    (["Vrin"], 46.655, 9.100, "rm-sursilv"),
    (["Vignogn"], 46.698, 9.148, "rm-sursilv"),
    (["Waltensburg/Vuorz", "Vuorz", "Waltensburg"], 46.775, 9.110, "rm-sursilv"),
    (["Rueun", "Ruis"], 46.778, 9.150, "rm-sursilv"),
    (["Schluein"], 46.789, 9.226, "rm-sursilv"),
    (["Laax", "Lags"], 46.807, 9.258, "rm-sursilv"),
    (["Flims", "Flem"], 46.835, 9.284, "rm-sursilv"),
    (["Falera"], 46.801, 9.233, "rm-sursilv"),
    (["Trin"], 46.828, 9.363, "rm-sursilv"),
    (["Domat/Ems", "Domat", "Ems"], 46.835, 9.450, "rm-sursilv"),
    (["Obersaxen", "Sursaissa"], 46.747, 9.107, "rm-sursilv"),
    # Sutsilvan
    (["Andeer"], 46.603, 9.426, "rm-sutsilv"),
    (["Zillis/Ziràn", "Zillis", "Ziran"], 46.634, 9.444, "rm-sutsilv"),
    (["Donat"], 46.628, 9.430, "rm-sutsilv"),
    (["Casti-Wergenstein"], 46.629, 9.408, "rm-sutsilv"),
    (["Domleschg/Tumleastga", "Tomils", "Tumegl", "Domleschg"], 46.763, 9.442, "rm-sutsilv"),
    (["Scharans/Scharons", "Scharans"], 46.717, 9.467, "rm-sutsilv"),
    (["Thusis/Tusàn", "Thusis", "Tusan"], 46.697, 9.440, "rm-sutsilv"),
    # Surmiran
    (["Savognin"], 46.597, 9.598, "rm-surmiran"),
    (["Tiefencastel/Casti", "Tiefencastel", "Casti", "Albula/Alvra"], 46.660, 9.578, "rm-surmiran"),
    (["Lantsch/Lenz", "Lantsch", "Lenz"], 46.683, 9.563, "rm-surmiran"),
    (["Vaz/Obervaz", "Lenzerheide", "Lai", "Valbella"], 46.722, 9.556, "rm-surmiran"),
    (["Salouf"], 46.618, 9.576, "rm-surmiran"),
    (["Mulegns"], 46.523, 9.620, "rm-surmiran"),
    (["Bivio/Beiva", "Bivio"], 46.469, 9.651, "rm-surmiran"),
    (["Bergün/Bravuogn", "Bergün", "Bergun", "Bravuogn"], 46.630, 9.746, "rm-surmiran"),
    (["Surses", "Oberhalbstein"], 46.560, 9.610, "rm-surmiran"),
    # Puter
    (["St. Moritz/San Murezzan", "St. Moritz", "Sankt Moritz", "San Murezzan", "St Moritz"], 46.498, 9.839, "rm-puter"),
    (["Celerina/Schlarigna", "Celerina", "Schlarigna"], 46.512, 9.859, "rm-puter"),
    (["Samedan"], 46.534, 9.873, "rm-puter"),
    (["Pontresina"], 46.491, 9.902, "rm-puter"),
    (["Bever"], 46.551, 9.889, "rm-puter"),
    (["La Punt Chamues-ch", "La Punt"], 46.582, 9.927, "rm-puter"),
    (["Madulain"], 46.587, 9.939, "rm-puter"),
    (["Zuoz"], 46.602, 9.960, "rm-puter"),
    (["S-chanf"], 46.614, 9.987, "rm-puter"),
    (["Silvaplana"], 46.459, 9.795, "rm-puter"),
    (["Sils/Segl", "Sils", "Segl", "Sils Maria"], 46.430, 9.758, "rm-puter"),
    # Vallader
    (["Zernez"], 46.698, 10.093, "rm-vallader"),
    (["Susch"], 46.751, 10.080, "rm-vallader"),
    (["Lavin"], 46.770, 10.110, "rm-vallader"),
    (["Guarda"], 46.775, 10.151, "rm-vallader"),
    (["Ardez"], 46.776, 10.200, "rm-vallader"),
    (["Ftan"], 46.795, 10.249, "rm-vallader"),
    (["Scuol", "Schuls"], 46.797, 10.298, "rm-vallader"),
    (["Sent"], 46.817, 10.338, "rm-vallader"),
    (["Ramosch"], 46.834, 10.383, "rm-vallader"),
    (["Tschlin"], 46.869, 10.428, "rm-vallader"),
    (["Tschierv"], 46.624, 10.343, "rm-vallader"),
    (["Valchava"], 46.597, 10.406, "rm-vallader"),
    (["Sta. Maria Val Müstair", "Santa Maria", "Sta. Maria"], 46.602, 10.423, "rm-vallader"),
    (["Müstair", "Mustair", "Val Müstair"], 46.630, 10.447, "rm-vallader"),
    # Graubünden places where Romansh is not the main local language
    (["Chur/Cuira", "Chur", "Cuira", "Coira"], 46.849, 9.532, None),
    (["Landquart"], 46.967, 9.555, None),
    (["Davos/Tavau", "Davos", "Tavau"], 46.802, 9.836, None),
    (["Klosters"], 46.869, 9.879, None),
    (["Arosa"], 46.778, 9.677, None),
    (["Splügen"], 46.551, 9.322, None),
    (["Vals"], 46.616, 9.180, None),
    (["Samnaun"], 46.944, 10.360, None),
    (["Poschiavo/Puschlav", "Poschiavo", "Puschlav"], 46.325, 10.058, None),
    (["Vicosoprano", "Bregaglia", "Bergell"], 46.354, 9.626, None),
    (["Mesocco/Misox", "Mesocco", "Misox"], 46.394, 9.233, None),
    (["Roveredo"], 46.237, 9.127, None),
]

GR_BBOX = (46.15, 47.07, 8.65, 10.50)  # lat_min, lat_max, lon_min, lon_max


def _norm(s: str) -> str:
    s = unicodedata.normalize("NFKD", s or "").encode("ascii", "ignore").decode()
    return " ".join(s.lower().replace("-", " ").replace(".", " ").split())


_NAME_INDEX = {_norm(n): p for p in PLACES for n in p[0]}


def _haversine(lat1, lon1, lat2, lon2):
    r = 6371.0
    p1, p2 = math.radians(lat1), math.radians(lat2)
    dp, dl = p2 - p1, math.radians(lon2 - lon1)
    a = math.sin(dp / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2
    return 2 * r * math.asin(math.sqrt(a))


def nearest_place(lat, lon):
    return min(((p, _haversine(lat, lon, p[1], p[2])) for p in PLACES), key=lambda x: x[1])


def idioms_by_distance(lat, lon, first=None):
    """The five idioms with `first` (the local one) at the front, the others ordered by how close their nearest
    village is; Rumantsch Grischun, the written standard, comes last."""
    regional = [code for code in IDIOMS if code != "rm-rumgr"]
    if lat is not None and lon is not None:
        dist = {code: min((_haversine(lat, lon, p[1], p[2]) for p in PLACES if p[3] == code), default=float("inf"))
                for code in regional}
        regional.sort(key=dist.get)
    if first in regional:
        regional.remove(first)
        regional.insert(0, first)
    return [*regional, "rm-rumgr"]


SEARCH_URL = "https://api3.geo.admin.ch/rest/services/api/SearchServer"
# Municipalities first, then postcodes, then other named places (localities before e.g. passes or buildings)
_ORIGIN_RANK = {"gg25": 0, "zipcode": 1, "gazetteer": 2}
_UMLAUTS = {"ae": "ä", "oe": "ö", "ue": "ü"}


def _search_name(label):
    """Place name from a swisstopo search label like '<i>Ort</i> <b>Ftan</b> (GR) - Scuol' or '<b>Chur (GR)</b>'."""
    m = re.search(r"<b>([^<]*)", label)
    name = m.group(1) if m else re.sub(r"<[^>]+>", "", label)
    return re.sub(r"\s*\([A-Z]{2}\)$", "", name).strip()


def _swisstopo_search(text):
    r = requests.get(SEARCH_URL, params={"searchText": text, "type": "locations", "origins": "gg25,zipcode,gazetteer",
                                         "limit": 15, "sr": 4326}, timeout=10)
    r.raise_for_status()
    return r.json().get("results", [])


@lru_cache(maxsize=512)
def geocode(city: str):
    """Look up a Swiss place name or postcode with the swisstopo search service (geo.admin.ch).

    The service also returns loose matches, so only a hit whose name is the typed name (or starts with it) counts;
    exact names win over longer ones, then municipalities over postcodes over other places.
    Returns (lat, lon, name) or None.
    """
    q = _norm(city)
    if not q:
        return None
    queries = [city]
    if any(k in q for k in _UMLAUTS):  # "Zuerich" -> also try "Zürich"
        queries.append(re.sub("ae|oe|ue", lambda m: _UMLAUTS[m.group(0)], city.lower()))
    best = None
    try:
        for text in queries:
            for hit in _swisstopo_search(text):
                a = hit["attrs"]
                name = _search_name(a.get("label", ""))
                names = [_norm(n) for n in re.split(r"/| - ", name)]
                if a.get("origin") == "zipcode":
                    names.append(_norm(re.sub(r"^\d{4}\s*-?\s*", "", name)))
                names = [n for n in names if n]
                variants = {q, _norm(text)}
                exact = any(n in variants for n in names)
                if not exact and not any(n.startswith(v + " ") for n in names for v in variants):
                    continue
                locality = a.get("origin") != "gazetteer" or "<i>Ort</i>" in a.get("label", "")
                key = (not exact, _ORIGIN_RANK.get(a.get("origin"), 9), not locality)
                if best is None or key < best[0]:
                    best = (key, float(a["lat"]), float(a["lon"]), name)
            if best:
                break
    except Exception:
        return None
    return best[1:] if best else None


def in_graubuenden(lat, lon):
    return GR_BBOX[0] <= lat <= GR_BBOX[1] and GR_BBOX[2] <= lon <= GR_BBOX[3]


def resolve_location(lat, lon, city):
    """Return (idiom_code, info). info holds language-neutral facts that i18n turns into text:

    src:    "town" | "geo" | "device" | None   where the location came from
    status: "local" | "non_rm" | "outside" | "notfound" | "noloc"
    place, near, city, label: details for the message
    """
    city = (city or "").strip()
    info = {"src": None, "city": city}
    if city:
        place = _NAME_INDEX.get(_norm(city))
        if place:
            return _idiom_for(place, 0.0, {**info, "src": "town"})
        hit = geocode(city)
        if not hit:
            return "rm-rumgr", {**info, "status": "notfound"}
        lat, lon, label = hit
        info.update(src="geo", label=label.split(",")[0], lat=lat, lon=lon)
    elif lat is not None and lon is not None:
        info.update(src="device", lat=lat, lon=lon)
    else:
        return "rm-rumgr", {**info, "status": "noloc"}

    if not in_graubuenden(lat, lon):
        return "rm-rumgr", {**info, "status": "outside", "lat": lat, "lon": lon}
    place, dist = nearest_place(lat, lon)
    return _idiom_for(place, dist, info)


def _idiom_for(place, dist, info):
    idiom = place[3]
    info = {**info, "place": place[0][0], "near": dist > 1.5,
            "lat": info.get("lat", place[1]), "lon": info.get("lon", place[2])}
    if idiom is None:
        return "rm-rumgr", {**info, "status": "non_rm"}
    return idiom, {**info, "status": "local"}


# ---------------------------------------------------------------- phrases

# Built by tools/build_phrases.py from Pledari Grond (Lia Rumantscha) dictionary entries, ranked by
# how often each form is used as a greeting in ZurichNLP's Romansh corpora. Puter and Vallader
# (no queryable dictionary) come from the ZurichNLP model, ranked the same way.
_PHRASES = json.loads((Path(__file__).parent / "phrases.json").read_text(encoding="utf-8"))


def phrases_for(idiom):
    """The 9 phrases after the greeting, as dicts with 'text' and 'source' ('pledarigrond' | 'model')."""
    return _PHRASES[idiom]


# ---------------------------------------------------------------- model (free-text translation only)

_model = None
_lock = threading.Lock()


def _load_model():
    global _model
    if _model is None:
        import ctranslate2
        from huggingface_hub import snapshot_download
        from transformers import NllbTokenizer

        path = snapshot_download(MODEL_ID)
        _model = (NllbTokenizer.from_pretrained(path), ctranslate2.Translator(path, device="cpu", compute_type="int8"))
    return _model


def translate(sentences, src, tgt, beam_size=4):
    with _lock:  # loading and the tokenizer's source language are shared state
        tokenizer, translator = _load_model()
        tokenizer.src_lang = src
        tokenizer.set_src_lang_special_tokens(src)
        batch = [tokenizer.convert_ids_to_tokens(tokenizer.encode(s)) for s in sentences]
        results = translator.translate_batch(
            batch, target_prefix=[[tgt]] * len(batch), beam_size=beam_size, max_decoding_length=256
        )
    return [
        tokenizer.decode(tokenizer.convert_tokens_to_ids(r.hypotheses[0][1:]), skip_special_tokens=True)
        for r in results
    ]


def idiom_code(name):
    return next(k for k, v in IDIOMS.items() if v == name)


def phrase_rows(idiom):
    """The 10 phrases for an idiom as (romansh, verified, key); meanings live in i18n.MEANINGS."""
    return [(GREETING, True, "hello"),
            *((p["text"], p["source"] == "pledarigrond", p["key"]) for p in phrases_for(idiom))]