"""Small experimental QVikhr editor. Modes, literal edits and output checks are local."""
import re
from decimal import Decimal
from num2words import num2words
from prompts import split, TTS

PUNCTUATE = """Исправь регистр и пунктуацию. Начни с заглавной буквы. Расставь запятые. После просьбы произнести заданный текст поставь двоеточие. Сохрани каждое слово, код, число, знак, единицу и порядок слов. Сохрани незаконченный конец фразы. Верни всю реплику одной строкой."""
ALIASES = {
    "дабл ю эм эс": "WMS", "даблю эм эс": "WMS",
    "дабл ю эф эс": "WFS", "даблю эф эс": "WFS",
    "дабл ю джи эс восемьдесят четыре": "WGS84", "даблю джи эс восемьдесят четыре": "WGS84",
    "гео джейсон": "GeoJSON", "геоджейсон": "GeoJSON", "гео тифф": "GeoTIFF",
    "ю ю и и": "UUEE", "ю ю ди ди": "UUDD", "кью эн эйч": "QNH",
    "джи эф эс": "GFS", "ди би зи": "dBZ", "джи ар ай би": "GRIB",
    "джи пи икс": "GPX", "дабл ю эм ти эс": "WMTS", "ди и эм": "DEM",
    "кью джи ай эс": "QGIS", "и си эм дабл ю эф": "ECMWF",
    "ай эл эс": "ILS", "ай эф ар": "IFR", "джи эн эс эс": "GNSS",
    "эс ай ди": "SID", "эс ти эй ар": "STAR", "ар ви ар": "RVR",
}
TERMS = re.compile(r"(?<![\w-])(?:" + "|".join(map(re.escape, sorted(ALIASES, key=len, reverse=True))) + r")(?![\w-])", re.I)
TOKEN = re.compile(r"[+-−]?\d+(?:[.,:]\d+)*°?|[A-Za-z][A-Za-z0-9]*(?:[._/-][A-Za-z0-9]+)*|[^\W\d_]+|[^\s.,!?…:;«»\"„“()]", re.UNICODE)
NUM = re.compile(r"(?<![\w.])[+−-]?\d+(?::[0-5]\d|[.,]\d+)?°?(?![\w]|[.,]\d)")
METHOD = re.compile(r"\b[A-Za-z][\w]*\.[A-Za-z][\w]*(?:\.[A-Za-z][\w]*)*\b")
ASK = re.compile(r"^(?:что|где|когда|какой|какая|какие|сколько|куда|почему)\b", re.I)
UNFINISHED = re.compile(r"\b(?:если|когда|пока(?: не)?|потому что|но сначала)\s*[.!…]*$", re.I)


def literal_asr(text):
    text = TERMS.sub(lambda match: ALIASES[match[0].lower()], text)
    if re.search(r"\b(?:файл\w*|прогноз\w*|архив\w*)\b", text, re.I):
        text = re.sub(r"(\bформат(?:е|а)?\s+)гриб\b", r"\1GRIB", text, flags=re.I)
        text = re.sub(r"(\bформат(?:е|а)?\s+)айкон\b", r"\1ICON", text, flags=re.I)
    if re.search(r"\b(?:GFS|ECMWF|прогноз\w*|погод\w*)\b", text, re.I):
        text = re.sub(r"\bайкон\b", "ICON", text, flags=re.I)
    text = re.sub(r"(\bсхем\w*\s+вылета\s+)сид\b", r"\1SID", text, flags=re.I)
    return re.sub(r"(\bсхем\w*\s+прибытия\s+)стар\b", r"\1STAR", text, flags=re.I)


def signature(text):
    return [word.casefold() if re.fullmatch(r"[А-Яа-яЁё]+", word) else word
            for word in TOKEN.findall(text)]


def punctuate(source, candidate):
    accepted = bool(candidate and "\n" not in candidate and signature(source) == signature(candidate))
    text = candidate if accepted else source
    text = re.sub(r"(?<![,;:.!?])\s+(но)\s+", r", \1 ", text, flags=re.I)
    text = re.sub(r"(?<![,;:.!?])\s+(пока\s+(?:не$|\S+\s+не\s+станет\b)|(?<!\bно\s)только если$)", r", \1", text, flags=re.I)
    if re.fullmatch(r"(?i)(?:скажи|произнеси|повтори)\s+\w+\s+(?:готов\w*|открыт\w*|закрыт\w*)", source):
        text = re.sub(r"^([^\s:,]+)[\s:,]+", r"\1: ", text).rstrip(":")
    text = text[:1].upper() + text[1:]
    if UNFINISHED.search(source):
        return text.rstrip(".!…"), accepted
    if ASK.match(source):
        return text.rstrip(".!…?") + "?", accepted
    return text if text.endswith((".", "!", "?", "…")) else text.rstrip(":;,") + ".", accepted


def speak_numbers(text):
    def replace(match):
        raw = match[0]
        if re.search(r"[A-Za-z][A-Za-z0-9_.-]*:$", text[:match.start()]):
            return raw
        degrees = raw.endswith("°")
        value = raw.removesuffix("°").replace("−", "-").replace(",", ".")
        if ":" in value:
            hour, minute = map(int, value.split(":"))
            if not (0 <= hour <= 23):
                return raw
            hours = "час" if hour % 10 == 1 and hour != 11 else "часа" if hour % 10 in (2,3,4) and not 12 <= hour <= 14 else "часов"
            minutes = "минуту" if minute % 10 == 1 and minute != 11 else "минуты" if minute % 10 in (2,3,4) and not 12 <= minute <= 14 else "минут"
            return f"{num2words(hour, lang='ru')} {hours}" + (f" {num2words(minute, lang='ru', gender='f', case='a')} {minutes}" if minute else "")
        case = "g" if re.search(r"\b(?:больше|меньше|около|свыше|до|от)\s+$", text[:match.start()], re.I) else "n"
        number = Decimal(value)
        spoken = num2words(abs(number), lang="ru", case=case)
        if value.startswith(("-", "+")):
            spoken = ("минус " if value.startswith("-") else "плюс ") + spoken
        if degrees:
            count = int(abs(number))
            unit = "градуса" if "." in value or count % 10 in (2,3,4) and not 12 <= count % 100 <= 14 else "градус" if count % 10 == 1 and count % 100 != 11 else "градусов"
            spoken += " " + unit
        return spoken
    text = re.sub(r"\[([+−-]?\d+[.,]\d+)\s*;\s*([+−-]?\d+[.,]\d+)\]", r"\1 и \2", text)
    return NUM.sub(replace, text)


def spoken_answer(query, source, candidate=""):
    query = query or ""
    aloud = re.search(r"\b(?:вслух|озвучь|зачитай)\b", query, re.I)
    shown = re.search(r"\bпокажи\b.*\b(?:инструмент\w*|список)\b", query, re.I)
    tools = re.match(r"доступные инструменты", source, re.I)
    if shown and tools and not aloud:
        return "Все доступные инструменты показаны в чате.", "display-rule"
    if aloud and METHOD.search(source):
        descriptions = re.split(r"[,\n]\s*(?=(?:[-*]\s*)?[A-Za-z]+\.[A-Za-z])", source)
        if len(descriptions) > 1 and all(METHOD.match(re.sub(r"^[-*]\s*", "", item)) for item in descriptions):
            parts = [METHOD.sub("", item, count=1).strip(" -*—.") for item in descriptions]
            if all(parts):
                lead = ["Один инструмент ", "другой "] if len(parts) == 2 else ["Доступные функции: "] + [""]*(len(parts)-1)
                return speak_numbers(", ".join(prefix + part for prefix,part in zip(lead,parts)) + "."), "aloud-rule"

    # Preserve all source clauses whenever the model changes words or facts.
    accepted = bool(candidate and "\n" not in candidate and signature(source) == signature(candidate))
    text = candidate if accepted else source
    text = re.sub(r"\b(Команда|Проверка)\s+"+METHOD.pattern+r"\s+", r"\1 ", text)
    text = re.sub(r"\bПроверка завершилась по тайм-ауту\. Состояние", "Проверка завершилась по тайм-ауту, поэтому состояние", text)
    if re.search(r"\bпроверка\b[^.!?]*тайм-аут", text, re.I) and not re.search(r"неизвест|не подтверж", text, re.I):
        text = text.rstrip(".") + ", поэтому результат пока неизвестен."
    return speak_numbers(text), "model-punctuation" if accepted else "source-verified"


def model_input(text):
    mode, query, answer, body = split(text)
    if mode == "TTS":
        source = answer if answer is not None else body
        return TTS, f"Запрос: {query or ''}\nОтвет: {source}"
    return PUNCTUATE, literal_asr(query if query is not None else body)


def finish(text, model_output):
    mode, query, answer, body = split(text)
    if not body:
        return "", "empty"
    if mode == "TTS":
        return spoken_answer(query, answer if answer is not None else body, model_output)
    raw = literal_asr(query if query is not None else body)
    asr, accepted = punctuate(raw, model_output)
    if mode == "ASR":
        return asr, "model-punctuation" if accepted else "source-verified"
    spoken, policy = spoken_answer(query, answer if answer is not None else asr)
    return asr+"\n\n"+spoken, ("model-punctuation" if accepted else "source-verified") + "+" + policy