app/models/llm/postprocessor.py

import re
import logging

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Patterns for output cleanup
# ---------------------------------------------------------------------------

# Strips raw [THINK]...[/THINK] blocks from visible output
_THINK_PATTERN = re.compile(r'\[THINK\](.*?)\[/THINK\]', re.DOTALL)

# Strips the internal [RAG_EMPTY] signal tag if it leaks into output
_RAG_EMPTY_PATTERN = re.compile(r'\[RAG_EMPTY\].*?\[/RAG_EMPTY\]', re.DOTALL)

# Matches "N/A", "Unknown", "Missing", "None" (case-insensitive) when used as values
_NA_VALUE_PATTERN = re.compile(
    r'(?:rating|price|cost|cuisine|review|stars?|score)\s*[:：]\s*(?:N/A|Unknown|Missing|None|-)\b',
    re.IGNORECASE
)

# Matches standalone known POI markers that are NOT yet wrapped in [[...]]
# We rely on the model to wrap them; this is a safety net for bare uppercase-starting words
# preceded by known location/attraction indicators.
_POI_INDICATOR = re.compile(
    r'(?<!\[)\[(?!\[)([A-ZÁÀẢÃẠĂẮẰẲẴẶÂẤẦẨẪẬÉÈẺẼẸÊẾỀỂỄỆÍÌỈĨỊÓÒỎÕỌÔỐỒỔỖỘƠỚỜỞỠỢÚÙỦŨỤƯỨỪỬỮỰÝỲỶỸỴĐ][^\[\]]{2,60})\](?!\])',
    re.UNICODE
)

# "Based on general travel information..." prefix for empty-RAG responses
_GENERAL_INFO_PREFIX = {
    "vi": "Dựa trên thông tin du lịch chung",
    "en": "Based on general travel information",
    "ja": "一般的な旅行情報に基づくと",
    "ko": "일반적인 여행 정보에 따르면",
}


def parse_llm_response(raw_output: str, language: str = "en", rag_was_empty: bool = False) -> dict:
    """Parse and post-process LLM output.

    Steps:
    1. Extract [THINK] reasoning blocks (hidden from user).
    2. Remove internal control tags.
    3. Strip N/A / Unknown / Missing value strings.
    4. Enforce [[POI]] bracket format on any single-bracketed references.
    5. Prepend "Based on general travel information..." when RAG was empty.
    6. Clean up whitespace.

    Args:
        raw_output: Raw text from the LLM.
        language: Response language code (vi/en/ja/ko).
        rag_was_empty: Whether the RAG context was empty for this query.

    Returns:
        dict with keys:
            - text: User-facing response text
            - reasoning: Chain-of-thought reasoning (or None)
            - rag_was_empty: Passed-through flag for downstream use
    """
    if not raw_output or not raw_output.strip():
        return {"text": "", "reasoning": None, "rag_was_empty": rag_was_empty}

    # 1. Extract [THINK] blocks
    think_matches = _THINK_PATTERN.findall(raw_output)
    reasoning = "\n".join(m.strip() for m in think_matches) if think_matches else None

    # 2. Remove [THINK] and [RAG_EMPTY] tags
    visible_text = _THINK_PATTERN.sub('', raw_output)
    visible_text = _RAG_EMPTY_PATTERN.sub('', visible_text)

    # 3. Remove N/A value patterns
    visible_text = _NA_VALUE_PATTERN.sub('', visible_text)

    # 4. Fix single-bracket POI references → [[POI]]
    #    e.g. [Hội An Ancient Town] → [[Hội An Ancient Town]]
    visible_text = _POI_INDICATOR.sub(r'[[\1]]', visible_text)

    # 5. Prepend general-info prefix when RAG was empty and model hasn't already added it
    if rag_was_empty:
        prefix = _GENERAL_INFO_PREFIX.get(language, _GENERAL_INFO_PREFIX["en"])
        if prefix.lower() not in visible_text.lower():
            visible_text = f"{prefix} —\n\n{visible_text}"

    # 6. Normalise whitespace
    visible_text = re.sub(r'\n{3,}', '\n\n', visible_text)
    visible_text = visible_text.strip()

    return {
        "text": visible_text,
        "reasoning": reasoning,
        "rag_was_empty": rag_was_empty,
    }


def extract_quick_replies_from_text(text: str, language: str = "vi") -> list[str]:
    """Attempt to extract suggested follow-up questions from LLM output.

    If the LLM includes numbered suggestions or bullet points at the end,
    extract them as quick replies.
    """
    lines = text.strip().split('\n')
    suggestions = []

    # Look for patterns like "1. ...", "- ...", "• ..." at the end of text
    for line in reversed(lines):
        line = line.strip()
        match = re.match(r'^(?:\d+[.)]\s*|[-•]\s*)(.+)$', line)
        if match and len(match.group(1)) < 60:
            suggestions.insert(0, match.group(1).strip())
        elif suggestions:
            break

    # Only return if we found a reasonable number
    if 2 <= len(suggestions) <= 6:
        return suggestions

    return []


def validate_response(text: str, intent: str, min_length: int = 10) -> bool:
    """Basic validation of LLM response quality.

    Args:
        text: Generated response text
        intent: Expected intent
        min_length: Minimum response length

    Returns:
        True if response passes validation
    """
    if not text or len(text.strip()) < min_length:
        return False

    # Check for obvious garbage (repeated characters, mostly special chars)
    if len(set(text)) < 5:
        return False

    # Check repetition ratio
    words = text.split()
    if len(words) > 10:
        unique_ratio = len(set(words)) / len(words)
        if unique_ratio < 0.2:
            return False

    return True