"""The proofreader's prompt protocol: the exact text the model was trained on.

The wordspace app sends two chat messages to the model, a system prompt and
the user's text, at temperature 0, and keeps only the "style" suggestions
whose "original" appears verbatim in the text (spelling and grammar come
from the operating system's checker). The model answers with one JSON object
and stops at its end-of-text token.

The wording must stay identical, character for character: the model has
only ever seen this text.
"""
import json

SYSTEM = """You are a proofreader. You reply ONLY with JSON.

Find mistakes in the user's text. For each one output an object with:
- "type": exactly one of "spelling", "grammar", "style"
- "original": the exact substring from the text, copied verbatim
- "replacement": the corrected substring

Rules, in order of importance:
1. MINIMAL SPAN. "original" must be the SHORTEST substring that has to
   change — usually ONE word. A misspelled word gives original="thatword"
   only. NEVER quote a whole sentence when a word is enough.
2. "original" MUST appear verbatim in the text, character for character.
   Never paraphrase it, never fix it inside "original".
3. One mistake = one object. Do not merge several mistakes into one.
4. "spelling" = the word does not exist. "grammar" = wrong agreement, tense
   or preposition. "style" = correct but wordy or clumsy.
5. Write in the SAME language as the text. Never translate.
6. Be conservative: if the text is fine, return an empty list.

Answer with this exact shape:
{"suggestions": [{"type": "...", "original": "...", "replacement": "..."}]}"""

# Language names as the app passes them (English names, capitalised); None
# when the text is too short to tell.
LANGUAGE_NAMES = {"it": "Italian", "en": "English", "es": "Spanish", "fr": "French",
                  "de": "German", "pt": "Portuguese"}


def user(text, language=None):
    """The user message: "Text (Italian):\\n..." or "Text:\\n..."."""
    message = "Text"
    if language:
        message += " (%s)" % language
    return message + ":\n" + text


def messages(text, language=None):
    """The two chat messages, for `tokenizer.apply_chat_template`."""
    return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user(text, language)}]


def first_json_object(answer):
    """The first JSON object with balanced braces, as the app reads it."""
    start = answer.find("{")
    while start >= 0:
        depth, in_string, escape = 0, False, False
        for i in range(start, len(answer)):
            ch = answer[i]
            if in_string:
                if escape:
                    escape = False
                elif ch == "\\":
                    escape = True
                elif ch == '"':
                    in_string = False
            elif ch == '"':
                in_string = True
            elif ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(answer[start:i + 1])
                    except ValueError:
                        break
        start = answer.find("{", start + 1)
    return None


def style_suggestions(answer, text):
    """[(original, replacement)]: style only, and only where `original` is in the text."""
    obj = first_json_object(answer)
    items = obj.get("suggestions") if isinstance(obj, dict) else None
    out = []
    for item in items if isinstance(items, list) else []:
        if not isinstance(item, dict):
            continue
        o, r = item.get("original"), item.get("replacement")
        if item.get("type") == "style" and isinstance(o, str) and o and isinstance(r, str) and o != r and o in text:
            out.append((o, r))
    return out