llm-in-text/backend/prompt.py

from datetime import datetime, timedelta, timezone
import re
from typing import Protocol, Tuple, runtime_checkable


@runtime_checkable
class UserPreferences(Protocol):
    language: str
    currency: str
    timezone: str


def _get_current_datetime(timezone_pref: str = "auto") -> str:
    # Default to UTC+8 if auto or not specified.
    offset = 8
    tz_info = " (UTC+8)"

    if timezone_pref and timezone_pref != "auto":
        # Parse values like "UTC+8" or "GMT-5".
        match = re.search(r"([+-])(\d+)", timezone_pref)
        if match:
            sign = match.group(1)
            hours = int(match.group(2))
            offset = hours if sign == "+" else -hours
            tz_info = f" ({timezone_pref})"
        else:
            tz_info = f" ({timezone_pref})"

    now = datetime.now(timezone(timedelta(hours=offset)))
    weekdays = [
        "Monday",
        "Tuesday",
        "Wednesday",
        "Thursday",
        "Friday",
        "Saturday",
        "Sunday",
    ]
    weekday = weekdays[now.weekday()]
    return (
        f"{now.year}-{now.month:02d}-{now.day:02d} "
        f"{weekday} {now.hour:02d}:{now.minute:02d}:{now.second:02d}{tz_info}"
    )


def _sanitize_language_id(language_id: str) -> str:
    if not language_id:
        return "markdown"
    allowed = []
    for ch in language_id.strip():
        if ch.isalnum() or ch in "-_+.":
            allowed.append(ch)
    value = "".join(allowed)[:32]
    return value or "markdown"


def _normalize_newlines(text: str) -> str:
    return (text or "").replace("\r\n", "\n").replace("\r", "\n")


def _prepare_context(prefix: str, suffix: str) -> Tuple[str, str]:
    """
    Prepare prefix/suffix for model completion context.
    Filter out potential web-scraping or legacy artifacts like <br>, <br/>, <br\\>.
    """
    br_pattern = re.compile(r"<br\s*/?\s*\\?>", re.IGNORECASE)
    clean_prefix = br_pattern.sub("", prefix or "")
    clean_suffix = br_pattern.sub("", suffix or "")
    return clean_prefix, clean_suffix


FENCE_LINE_RE = re.compile(r"^[ \t]*```.*$")
FENCE_INFO_RE = re.compile(r"^[ \t]*```[ \t]*(.*)$")
MERMAID_CONTEXT_RE = re.compile(
    r"```[ \t]*mermaid\b|"
    r"\b(flowchart|sequencediagram|classdiagram|statediagram(?:-v2)?|"
    r"erdiagram|journey|gantt|pie|mindmap|timeline|gitgraph|quadrantchart|xychart-beta)\b|"
    r"\bgraph[ \t]+(TD|TB|BT|RL|LR)\b",
    re.IGNORECASE,
)


def _cursor_in_fenced_code_block(prefix: str) -> bool:
    """
    Determine whether the cursor is currently inside a fenced code block.
    The state is computed by toggling on each markdown fence line that matches:
    ^[ \t]*```.*$
    """
    return _active_fence_language(prefix) != "none"


def _active_fence_language(prefix: str) -> str:
    """
    Return active fence language at cursor based on prefix.
    - "none": cursor is outside fenced code block
    - "unknown": cursor is inside a fence without language tag
    - "<language>": cursor is inside a fenced block with language tag
    """
    normalized = _normalize_newlines(prefix)
    in_fence = False
    active_language = "none"
    for line in normalized.split("\n"):
        if FENCE_LINE_RE.match(line):
            if in_fence:
                in_fence = False
                active_language = "none"
            else:
                info_match = FENCE_INFO_RE.match(line)
                info = info_match.group(1).strip() if info_match else ""
                if not info:
                    active_language = "unknown"
                else:
                    first_token = info.split()[0]
                    lang_chars = []
                    for ch in first_token.strip():
                        if ch.isalnum() or ch in "-_+.":
                            lang_chars.append(ch)
                    active_language = "".join(lang_chars)[:32].lower() or "unknown"
                in_fence = True
    return active_language if in_fence else "none"


def _is_mermaid_context(prefix: str, suffix: str, cursor_fence_language: str) -> bool:
    if cursor_fence_language == "mermaid":
        return True

    prefix_tail = (prefix or "")[-1200:]
    suffix_head = (suffix or "")[:400]
    combined = f"{prefix_tail}\n{suffix_head}"
    return MERMAID_CONTEXT_RE.search(combined) is not None


def prepare_prompt_context(prefix: str, suffix: str) -> Tuple[str, str]:
    return _prepare_context(prefix, suffix)


LANGUAGE_SYNONYMS = {
    "md": "markdown",
    "markdown": "markdown",
    "txt": "text",
    "text": "text",
    "plain": "text",
    "plaintext": "text",
    "py": "python",
    "python": "python",
    "js": "javascript",
    "javascript": "javascript",
    "jsx": "javascript",
    "node": "javascript",
    "ts": "typescript",
    "tsx": "typescript",
    "typescript": "typescript",
    "json": "json",
    "jsonc": "json",
    "json5": "json",
    "yaml": "yaml",
    "yml": "yaml",
    "toml": "toml",
    "ini": "ini",
    "cfg": "ini",
    "bash": "bash",
    "shell": "bash",
    "sh": "bash",
    "zsh": "bash",
    "fish": "bash",
    "ps": "powershell",
    "ps1": "powershell",
    "powershell": "powershell",
    "sql": "sql",
    "postgres": "sql",
    "postgresql": "sql",
    "mysql": "sql",
    "sqlite": "sql",
    "html": "html",
    "xml": "xml",
    "svg": "xml",
    "css": "css",
    "scss": "css",
    "less": "css",
    "latex": "latex",
    "tex": "latex",
    "katex": "latex",
    "mermaid": "mermaid",
    "c": "c",
    "c++": "cpp",
    "cpp": "cpp",
    "cxx": "cpp",
    "h": "c",
    "hpp": "cpp",
    "c#": "csharp",
    "cs": "csharp",
    "csharp": "csharp",
    "go": "go",
    "golang": "go",
    "rust": "rust",
    "rs": "rust",
    "java": "java",
    "kotlin": "kotlin",
    "swift": "swift",
    "ruby": "ruby",
    "rb": "ruby",
    "php": "php",
    "lua": "lua",
    "r": "r",
    "matlab": "matlab",
    "dart": "dart",
    "docker": "dockerfile",
    "dockerfile": "dockerfile",
    "make": "makefile",
    "makefile": "makefile",
    "diff": "diff",
    "patch": "diff",
    "regex": "regex",
}


def _canonical_language_id(language_id: str) -> str:
    safe = _sanitize_language_id(language_id).lower()
    if not safe:
        return "markdown"
    return LANGUAGE_SYNONYMS.get(safe, safe)


_JS_LANGS = {"javascript", "typescript"}
_CODE_LANGS = {"python", "go", "rust", "java", "kotlin", "swift", "ruby", "php", "lua", "c", "cpp", "csharp", "r", "matlab", "dart"}

_LANG_GUIDANCE = {
    "mermaid": """
Language-specific guidance (mermaid):
- Output valid Mermaid syntax only.
- Prefer concise, syntactically correct diagram statements.
- Avoid prose unless the user prompt explicitly requires it.""",
    "latex": """
Language-specific guidance (latex):
- Output LaTeX math content only when completing LaTeX.
- If CURSOR_IN_FENCED_CODE_BLOCK=true and CURSOR_FENCE_LANGUAGE is latex/tex/katex:
  - Output raw LaTeX lines only.
  - Do not wrap with $ or $$.""",
    "json": """
Language-specific guidance (json):
- Output strict JSON only (no comments, no trailing commas).
- Ensure valid quotes and braces.""",
    "yaml": """
Language-specific guidance (yaml):
- Output valid YAML only.
- Use consistent indentation and avoid tabs.""",
    "toml": """
Language-specific guidance (toml):
- Output valid TOML only.
- Keep key types consistent.""",
    "ini": """
Language-specific guidance (ini):
- Output valid INI only.
- Keep section headers and key=value pairs consistent.""",
    "sql": """
Language-specific guidance (sql):
- Output a single, valid SQL statement unless context requires multiple.
- Prefer ANSI SQL when dialect is unclear.""",
    "bash": """
Language-specific guidance (bash):
- Output POSIX-compatible shell when possible.
- Avoid interactive prompts or destructive commands unless requested.""",
    "powershell": """
Language-specific guidance (powershell):
- Output valid PowerShell commands.
- Avoid destructive commands unless explicitly requested.""",
    "html": """
Language-specific guidance (html):
- Output valid HTML only.
- Keep markup minimal and well-formed.""",
    "css": """
Language-specific guidance (css):
- Output valid CSS only.
- Use concise, readable selectors.""",
    "diff": """
Language-specific guidance (diff):
- Output a unified diff only.
- Ensure @@ hunk headers and +/- lines are consistent.""",
    "regex": """
Language-specific guidance (regex):
- Output the regex pattern only.
- Avoid delimiters unless explicitly requested.""",
    "text": """
Language-specific guidance (text):
- Output plain text only.
- Avoid markdown formatting unless explicitly asked.""",
    "xml": """
Language-specific guidance (xml):
- Output well-formed XML only.
- Ensure matching tags and proper escaping.""",
    "dockerfile": """
Language-specific guidance (dockerfile):
- Output valid Dockerfile instructions only.
- Keep layers minimal and ordered logically.""",
    "makefile": """
Language-specific guidance (makefile):
- Output valid Makefile syntax only.
- Use tabs for recipe lines.""",
}

_GENERIC_CODE = """
Language-specific guidance ({lang}):
- Output valid {lang} code.
- Avoid prose unless context clearly expects comments or docstrings."""

_JS_CODE = """
Language-specific guidance ({lang}):
- Output valid {lang} code.
- Prefer modern syntax and avoid prose unless comments are needed."""


def _language_guidance(language_id: str) -> str:
    canonical = _canonical_language_id(language_id)
    if canonical == "markdown":
        return ""
    guidance = _LANG_GUIDANCE.get(canonical)
    if guidance:
        return guidance
    if canonical in _JS_LANGS:
        return _JS_CODE.format(lang=canonical)
    if canonical in _CODE_LANGS:
        return _GENERIC_CODE.format(lang=canonical)
    return _GENERIC_CODE.format(lang=canonical)


def build_inline_system_prompt(language_id: str = "markdown") -> str:
    safe_language_id = _canonical_language_id(language_id)
    language_guidance = _language_guidance(safe_language_id)

    system_prompt = f"""You are an inline completion engine for a {safe_language_id} editor with ghost-text suggestions.

Return only the insertion text that should be placed between PREFIX and SUFFIX.

CORE PRINCIPLE: Output insertion text only. No explanations, no meta labels, no wrapper quotes.

PRIORITY 1: CONTEXT AWARENESS (Read these flags from user prompt)
- CURSOR_IN_FENCED_CODE_BLOCK: Are you inside a code fence?
- CURSOR_FENCE_LANGUAGE: What language is the current fence?
- PREFIX_ENDS_WITH_NEWLINE: Does prefix end with newline?
- SUFFIX_STARTS_WITH_NEWLINE: Does suffix start with newline?
- MERMAID_CONTEXT: Is this a Mermaid diagram context?

PRIORITY 2: SPECIALIZED CONTENT RULES

2.1 Code Block Handling:
If CURSOR_IN_FENCED_CODE_BLOCK=true:
  - You are inside a code fence
  - Output code lines ONLY (no triple backticks)
  - Use single \\n for code line separation

If CURSOR_IN_FENCED_CODE_BLOCK=false and code needed:
  - Wrap code in fenced block with language tag:
    ```{{language}}
    code here
    ```
  - Never use inline backticks for code snippets

2.2 Math Formatting (KaTeX):
- Inline math: wrap with $...$
- Block math: wrap with $$...$$
- Never output bare formulas
- Exception: inside latex/tex/katex fence, output raw LaTeX

2.3 Mermaid Diagrams:
If CURSOR_FENCE_LANGUAGE=mermaid:
  - Output Mermaid syntax ONLY
  - No backticks, no explanations

If MERMAID_CONTEXT=true and outside fence:
  - Output complete fenced block:
    ```mermaid
    diagram syntax
    ```

PRIORITY 3: MARKDOWN STRUCTURE

3.1 Newline Semantics:
- Single \\n: soft break (same paragraph, renders as space or <br>)
- Double \\n\\n: hard break (new paragraph/block)
- Use \\n\\n for: new paragraphs, before headings, starting lists/tables
- Use \\n for: continuation within blocks (list items, table cells)
- Exception: inside code blocks, use \\n freely for code lines

3.2 Boundary Management:
Check PREFIX_ENDS_WITH_NEWLINE and SUFFIX_STARTS_WITH_NEWLINE:
- If PREFIX lacks needed newline: start OUTPUT with \\n
- If SUFFIX lacks needed newline: end OUTPUT with \\n
- Common cases requiring leading \\n:
  * Starting a list after "Steps:"
  * Creating new paragraph after text
  * Adding heading after paragraph
- Common cases requiring trailing \\n:
  * Before new heading
  * End of section

3.3 Context Stitching:
- Never repeat text from SUFFIX beginning
- Match PREFIX tone, style, indentation
- Continue structures: lists, tables, quotes, headings

PRIORITY 4: HIDDEN CONTEXT
- OCR metadata like <OCR:...> is hidden context
- Never copy OCR tags to output
- Use OCR content as semantic hint only
"""

    if language_guidance:
        system_prompt = f"{system_prompt.rstrip()}\\n{language_guidance.strip()}"

    return system_prompt.strip()


INLINE_EXAMPLES = """=== CATEGORY A: PROSE CONTINUATION ===

[EX01] Simple prose continuation
<PREFIX>The quick brown fox </PREFIX>
<SUFFIX>jumps over the lazy dog.</SUFFIX>
Expected OUTPUT:
moved quietly and then

[EX02] Avoid repeating suffix
<PREFIX>Our launch plan starts with </PREFIX>
<SUFFIX>phase one, followed by phase two.</SUFFIX>
Expected OUTPUT:
careful internal testing before
WRONG: phase one starts with  (repeats suffix)

=== CATEGORY B: MARKDOWN STRUCTURES ===

[EX03] Continue checklist
<PREFIX>## TODO
- [ ] Buy milk
- [ ] </PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
Write release notes and share draft with team

[EX04] Start list after header (PREFIX lacks newline)
PREFIX_ENDS_WITH_NEWLINE=false
<PREFIX>Deployment steps:</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:

- Build artifact
- Deploy service

[EX05] Continue table row
<PREFIX>| Name | Score |
| --- | --- |
| Alice | 92 |
| Bob | </PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
88 |

[EX06] Start new paragraph
<PREFIX>First paragraph ends.</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:

Second paragraph starts.
WRONG: Second paragraph starts.  (missing leading \\n\\n)

[EX07] Add newline before heading
PREFIX_ENDS_WITH_NEWLINE=false
<PREFIX>End of previous section.</PREFIX>
<SUFFIX>## Next Heading</SUFFIX>
Expected OUTPUT:

WRONG:  (would join with heading without separation)

=== CATEGORY C: CODE BLOCKS ===

[EX08] Outside fence: wrap code in fence
CURSOR_IN_FENCED_CODE_BLOCK=false
<PREFIX>Parse this JSON payload in Python:</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
```python
import json
data = json.loads(payload)
```
WRONG: import json\\ndata = json.loads(payload)  (no fence)

[EX09] Inside fence: output code only
CURSOR_IN_FENCED_CODE_BLOCK=true
<PREFIX>```python
def add(a, b):
    return </PREFIX>
<SUFFIX>
```</SUFFIX>
Expected OUTPUT:
a + b
WRONG: ```python\\nreturn a + b\\n```  (duplicate fences)

[EX10] Code inside fence uses single newline
CURSOR_IN_FENCED_CODE_BLOCK=true
<PREFIX>```python
def hello():</PREFIX>
<SUFFIX>
```</SUFFIX>
Expected OUTPUT:
    print("Hello")
    return True
(Note: single \\n between code lines, no markdown rules)

=== CATEGORY D: MATH ===

[EX11] Inline math
<PREFIX>The derivative of x^2 is </PREFIX>
<SUFFIX>.</SUFFIX>
Expected OUTPUT:
$2x$
WRONG: 2x  (bare formula)

[EX12] Block math
<PREFIX>We can write the Gaussian integral as:</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
$$
\\int_{-\\infty}^{\\infty} e^{-x^2}\\,dx = \\sqrt{\\pi}
$$
WRONG: \\int... (bare formula without $$)

=== CATEGORY E: MERMAID ===

[EX13] Inside mermaid fence
CURSOR_FENCE_LANGUAGE=mermaid
CURSOR_IN_FENCED_CODE_BLOCK=true
<PREFIX>```mermaid
flowchart TD
    A[Start] --> </PREFIX>
<SUFFIX>
```</SUFFIX>
Expected OUTPUT:
B{Valid?}
B -->|Yes| C[Done]
WRONG: ```mermaid\\nB{Valid?}...  (duplicate fence)

[EX14] Outside fence with mermaid context
CURSOR_IN_FENCED_CODE_BLOCK=false
MERMAID_CONTEXT=true
<PREFIX>Please provide a simple release pipeline diagram.</PREFIX>
<SUFFIX></SUFFIX>
Expected OUTPUT:
```mermaid
flowchart LR
    Build --> Test --> Deploy
```

=== CATEGORY F: OCR METADATA ===

[EX15] Use OCR as context, never output
<PREFIX>![whiteboard](img.png) <OCR:equation y = mx + b>
The relationship is </PREFIX>
<SUFFIX>.</SUFFIX>
Expected OUTPUT:
$y = mx + b$
WRONG: <OCR:equation y = mx + b>  (OCR tag in output)"""


def build_completion_prompts(
    prefix: str,
    suffix: str,
    language_id: str = "markdown",
    location: str = "",
    thinking_level: str = "low",
    preferences: UserPreferences | None = None,
) -> Tuple[str, str]:
    safe_language_id = _canonical_language_id(language_id)
    recent_prefix, recent_suffix = _prepare_context(prefix, suffix)
    recent_prefix = _normalize_newlines(recent_prefix)
    recent_suffix = _normalize_newlines(recent_suffix)

    cursor_fence_language = _active_fence_language(recent_prefix)
    cursor_in_fenced_code_block = cursor_fence_language != "none"
    mermaid_context = _is_mermaid_context(
        recent_prefix, recent_suffix, cursor_fence_language
    )
    prefix_ends_with_newline = recent_prefix.endswith("\n")
    suffix_starts_with_newline = recent_suffix.startswith("\n")

    tz_pref = preferences.timezone if preferences else "auto"
    current_time = _get_current_datetime(tz_pref)
    location_info = f"\nUser location: {location}" if location else ""

    pref_info = []
    if preferences:
        if preferences.language and preferences.language != "auto":
            pref_info.append(f"Preferred language: {preferences.language}")
        if preferences.currency and preferences.currency != "auto":
            pref_info.append(f"Preferred currency: {preferences.currency}")

    preferences_instruction = "\n".join(pref_info)
    if preferences_instruction:
        preferences_instruction = f"\nUser Preferences:\n{preferences_instruction}"

    user_prompt = f"""Current time: {current_time}{location_info}{preferences_instruction}
Reasoning level: {thinking_level}
Editor language: {safe_language_id}

=== STATE FLAGS ===
- CURSOR_IN_FENCED_CODE_BLOCK: {"true" if cursor_in_fenced_code_block else "false"}
- CURSOR_FENCE_LANGUAGE: {cursor_fence_language}
- MERMAID_CONTEXT: {"true" if mermaid_context else "false"}
- PREFIX_ENDS_WITH_NEWLINE: {"true" if prefix_ends_with_newline else "false"}
- SUFFIX_STARTS_WITH_NEWLINE: {"true" if suffix_starts_with_newline else "false"}

=== TASK ===
Produce the best insertion text between PREFIX and SUFFIX.
Requirements:
- Non-empty and meaningful
- Concise unless structure needs more
- Follows markdown rules in system prompt

=== BOUNDARY DECISION GUIDE ===

Step 1: Check PREFIX_ENDS_WITH_NEWLINE
If false, ask: "Does output need to start on a new line?"
  - YES if PREFIX ends with: ":", "steps:", "items:", heading text, or complete sentence before heading
  - If YES: start output with \\n

Step 2: Check SUFFIX_STARTS_WITH_NEWLINE  
If false, ask: "Does output need to end with a newline?"
  - YES if SUFFIX starts with: heading (##), new paragraph, or list marker
  - If YES: end output with \\n

Step 3: Choose newline type
  - Use \\n\\n for: new paragraphs, before headings, starting lists
  - Use \\n for: continuing within blocks, list items, table cells
  - Exception: inside code fences, use \\n freely

=== CONTEXT NOTES ===
- OCR metadata (e.g., <OCR:description>) is hidden context, never copy to output
- Match PREFIX tone, style, and indentation
- Do not repeat text from SUFFIX beginning

=== EXAMPLES BY CATEGORY ===
{INLINE_EXAMPLES}

=== NOW COMPLETE THE TASK ===

<PREFIX>
{recent_prefix}
</PREFIX>

<SUFFIX>
{recent_suffix}
</SUFFIX>

Output:"""

    system_prompt = build_inline_system_prompt(safe_language_id)
    return system_prompt.strip(), user_prompt.strip()


def build_prompt(
    prefix: str,
    suffix: str,
    language_id: str = "markdown",
    location: str = "",
    thinking_level: str = "low",
    preferences: UserPreferences | None = None,
) -> str:
    """
    Backward-compatible helper. Returns only the user prompt body.
    """
    _, user_prompt = build_completion_prompts(
        prefix=prefix,
        suffix=suffix,
        language_id=language_id,
        location=location,
        thinking_level=thinking_level,
        preferences=preferences,
    )
    return user_prompt