utils/editor.py

"""
Post-generation Editor for AI Writer.
Performs strict review and correction of generated text based on requirements.
"""

import re
from typing import Dict, List, Tuple, Optional


# Currency symbols by country
CURRENCY_MAP = {
    "Индия": "₹",
    "Аргентина": "$",
    "Бангладеш": "৳",
    "Испания": "€",
    "Германия": "€",
    "Польша": "zł",
    "Франция": "€",
    "США": "$",
    "Великобритания": "£",
    "Китай": "¥",
    "Япония": "¥",
    "Южная Корея": "₩",
    "Бразилия": "R$",
    "Мексика": "$",
    "Канада": "$",
    "Австралия": "$",
    "Италия": "€",
    "Нидерланды": "€",
    "Бельгия": "€",
    "Швейцария": "CHF",
    "Турция": "₺",
    "Россия": "₽",
    "Украина": "₴",
    "Казахстан": "₸",
    "Беларусь": "Br",
    "ОАЭ": "د.إ",
    "Саудовская Аравия": "﷼",
    "Египет": "E£",
    "ЮАР": "R",
    "Нигерия": "₦",
    "Кения": "KSh",
    "Таиланд": "฿",
    "Вьетнам": "₫",
    "Индонезия": "Rp",
    "Филиппины": "₱",
    "Малайзия": "RM",
    "Сингапур": "$",
    "Израиль": "₪",
    "Швеция": "kr",
    "Норвегия": "kr",
    "Дания": "kr",
    "Финляндия": "€",
    "Португалия": "€",
    "Греция": "€",
    "Чехия": "Kč",
    "Румыния": "lei",
    "Венгрия": "Ft",
    "Австрия": "€",
    "Ирландия": "€",
}

# Currency letter codes that should be replaced with symbols
CURRENCY_CODE_MAP = {
    "USD": "$",
    "EUR": "€",
    "GBP": "£",
    "RUB": "₽",
    "UAH": "₴",
    "INR": "₹",
    "JPY": "¥",
    "CNY": "¥",
    "KRW": "₩",
    "BRL": "R$",
    "CAD": "$",
    "AUD": "$",
    "CHF": "CHF",
    "TRY": "₺",
    "MXN": "$",
    "SGD": "$",
    "HKD": "$",
    "NZD": "$",
    "SEK": "kr",
    "NOK": "kr",
    "DKK": "kr",
    "PLN": "zł",
    "CZK": "Kč",
    "HUF": "Ft",
    "THB": "฿",
    "VND": "₫",
    "PHP": "₱",
    "MYR": "RM",
    "IDR": "Rp",
    "BDT": "৳",
    "KZT": "₸",
    "BYN": "Br",
    "ARS": "$",
    "EGP": "E£",
    "ZAR": "R",
    "NGN": "₦",
    "KES": "KSh",
    "AED": "د.إ",
    "SAR": "﷼",
    "ILS": "₪",
    "RON": "lei",
    "BGN": "лв",
    "HRK": "kn",
    "RSD": "дин.",
    "GEL": "₾",
    "AMD": "֏",
    "AZN": "₼",
    "UZS": "сўм",
}


def get_currency_symbol(country: str) -> str:
    """Get the currency symbol for a given country."""
    return CURRENCY_MAP.get(country, "$")


def remove_dashes(text: str) -> str:
    """Remove all types of dashes from text, replacing with appropriate punctuation."""
    # Replace em-dash and en-dash with comma or period depending on context
    # Em dash
    text = re.sub(r'\s*—\s*', ', ', text)
    # En dash
    text = re.sub(r'\s*–\s*', ', ', text)
    # Regular dash surrounded by spaces (used as punctuation, not hyphen)
    text = re.sub(r'\s+-\s+', ', ', text)
    # Clean up double commas
    text = text.replace(',, ', ', ')
    text = text.replace(', , ', ', ')
    # Clean up comma after period
    text = text.replace('., ', '. ')
    text = text.replace('. , ', '. ')
    return text


def remove_quotation_marks(text: str) -> str:
    """Remove all types of quotation marks from text."""
    # Remove curly/smart quotes
    text = text.replace('\u201c', '')  # left double
    text = text.replace('\u201d', '')  # right double
    text = text.replace('\u2018', '')  # left single
    text = text.replace('\u2019', '')  # right single
    # Remove guillemets (Russian quotes)
    text = text.replace('\u00ab', '')  # «
    text = text.replace('\u00bb', '')  # »
    # Remove straight quotes
    text = text.replace('"', '')
    text = text.replace('"', '')
    text = text.replace('"', '')
    # Remove single quotes used as quotation marks (but keep apostrophes in contractions)
    # This is tricky - we only remove standalone single quotes, not apostrophes
    text = re.sub(r'(?<=\s)\'|\'(?=\s)', '', text)
    # Clean up extra spaces left after removal
    text = re.sub(r'\s{2,}', ' ', text)
    return text


def remove_personal_pronouns(text: str) -> str:
    """Remove personal pronouns and replace with impersonal constructions."""
    # Russian first person plural
    replacements = {
        # мы (we)
        r'\bмы\b': '',
        r'\bнаш\b': '',
        r'\bнаша\b': '',
        r'\bнаше\b': '',
        r'\bнаши\b': '',
        r'\bнашего\b': '',
        r'\bнашей\b': '',
        r'\bнашему\b': '',
        r'\bнашим\b': '',
        r'\bнашими\b': '',
        r'\bнаших\b': '',
        r'\bнашу\b': '',
        r'\bнам\b': '',
        r'\bнами\b': '',
        r'\bнас\b': '',
        # ты/вы (you)
        r'\bты\b': '',
        r'\bтвой\b': '',
        r'\bтвоя\b': '',
        r'\bтвоё\b': '',
        r'\bтвои\b': '',
        r'\bтвоего\b': '',
        r'\bтвоей\b': '',
        r'\bтвоему\b': '',
        r'\bтвоим\b': '',
        r'\bтвоими\b': '',
        r'\bтвоих\b': '',
        r'\bтвою\b': '',
        r'\bтебе\b': '',
        r'\bтобой\b': '',
        r'\bтобою\b': '',
        r'\bтебя\b': '',
        r'\bвы\b': '',
        r'\bваш\b': '',
        r'\bваша\b': '',
        r'\bваше\b': '',
        r'\bваши\b': '',
        r'\bвашего\b': '',
        r'\bвашей\b': '',
        r'\bвашему\b': '',
        r'\bвашим\b': '',
        r'\bвашими\b': '',
        r'\bваших\b': '',
        r'\bвашу\b': '',
        r'\bвам\b': '',
        r'\bвами\b': '',
        r'\bвас\b': '',
    }
    for pattern, replacement in replacements.items():
        text = re.sub(pattern, replacement, text, flags=re.IGNORECASE)
    # Clean up double spaces
    text = re.sub(r'\s{2,}', ' ', text)
    # Clean up sentences starting with leftover particles
    text = re.sub(r'^\s*(же|ли|бы|же)\b', '', text, flags=re.MULTILINE)
    return text


def replace_currency_codes_with_symbols(text: str) -> str:
    """Replace currency letter codes (USD, EUR etc.) with their symbols."""
    for code, symbol in CURRENCY_CODE_MAP.items():
        # Pattern: number followed by space and currency code, or currency code followed by number
        text = re.sub(
            rf'(\d[\d\s]*)\s*{code}\b',
            rf'{symbol}\1',
            text,
            flags=re.IGNORECASE
        )
        text = re.sub(
            rf'\b{code}\s*(\d[\d\s]*)',
            rf'{symbol}\1',
            text,
            flags=re.IGNORECASE
        )
    return text


def ensure_currency_symbol(text: str, country: str) -> str:
    """Ensure the correct currency symbol is used for the specified country."""
    symbol = get_currency_symbol(country)
    # First replace any currency codes with symbols
    text = replace_currency_codes_with_symbols(text)
    return text


def count_keyword_occurrences(text: str, keyword: str) -> int:
    """Count how many times a keyword appears in the text (case-insensitive)."""
    return len(re.findall(re.escape(keyword), text, re.IGNORECASE))


def check_text_length(text: str) -> int:
    """Count words in the text."""
    return len(text.split())


def validate_keywords(text: str, keywords: List[Dict]) -> Dict[str, int]:
    """Check keyword counts against required counts. Returns dict of keyword -> actual count."""
    results = {}
    for kw in keywords:
        word = kw.get("word", "")
        if word:
            results[word] = count_keyword_occurrences(text, word)
    return results


def run_editor(
    text: str,
    country: str,
    keywords: List[Dict],
    lsi_keywords: List[Dict],
    min_words: int,
    max_words: int,
) -> Tuple[str, str]:
    """
    Run the automated editor on generated text.
    Returns (edited_text, editor_report).
    """
    report_lines = ["=== EDITOR REPORT ===\n"]
    changes_made = []

    # 1. Remove dashes
    original = text
    text = remove_dashes(text)
    if text != original:
        changes_made.append("Removed dashes")

    # 2. Remove quotation marks
    original = text
    text = remove_quotation_marks(text)
    if text != original:
        changes_made.append("Removed quotation marks")

    # 3. Remove personal pronouns
    original = text
    text = remove_personal_pronouns(text)
    if text != original:
        changes_made.append("Removed personal pronouns (we/our/us, you/your)")

    # 4. Ensure currency symbols
    original = text
    text = ensure_currency_symbol(text, country)
    if text != original:
        changes_made.append(f"Replaced currency codes with symbols for {country}")

    # 5. Clean up extra whitespace
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = re.sub(r'[ \t]+', ' ', text)

    # 6. Check text length
    word_count = check_text_length(text)
    report_lines.append(f"Word count: {word_count} (min: {min_words}, max: {max_words})")
    if word_count < min_words:
        report_lines.append(f"WARNING: Text is {min_words - word_count} words SHORT of minimum.")
    elif word_count > max_words:
        report_lines.append(f"WARNING: Text is {word_count - max_words} words OVER maximum.")
    else:
        report_lines.append("Text length is within the specified range.")

    # 7. Check keywords
    report_lines.append("\n--- Keywords ---")
    if keywords:
        for kw in keywords:
            word = kw.get("word", "")
            required = kw.get("count", 1)
            actual = count_keyword_occurrences(text, word)
            status = "✓" if actual >= required else "✗"
            report_lines.append(f"  {status} '{word}': {actual}/{required} occurrences")
    else:
        report_lines.append("  No keywords specified.")

    # 8. Check LSI keywords
    report_lines.append("\n--- LSI Keywords ---")
    if lsi_keywords:
        for kw in lsi_keywords:
            word = kw.get("word", "")
            required = kw.get("count", 1)
            actual = count_keyword_occurrences(text, word)
            status = "✓" if actual >= required else "✗"
            report_lines.append(f"  {status} '{word}': {actual}/{required} occurrences")
    else:
        report_lines.append("  No LSI keywords specified.")

    # 9. Check for remaining dashes
    dash_check = bool(re.search(r'[—–]', text))
    if dash_check:
        report_lines.append("\nWARNING: Some dashes may still remain in the text.")
    else:
        report_lines.append("\nNo dashes found in text. ✓")

    # 10. Check for remaining quotation marks
    quote_check = bool(re.search(r'["\u201c\u201d\u00ab\u00bb]', text))
    if quote_check:
        report_lines.append("WARNING: Some quotation marks may still remain in the text.")
    else:
        report_lines.append("No quotation marks found in text. ✓")

    # 11. Check for personal pronouns
    pronoun_pattern = r'\b(мы|наш|наша|наше|наши|нам|нами|нас|нашу|ты|твой|твоя|твоё|тебе|тебя|вы|ваш|ваша|ваше|ваши|вам|вами|вас)\b'
    pronoun_check = bool(re.search(pronoun_pattern, text, re.IGNORECASE))
    if pronoun_check:
        report_lines.append("WARNING: Some personal pronouns may still remain in the text.")
    else:
        report_lines.append("No personal pronouns found in text. ✓")

    # 12. Currency symbol check
    currency_symbol = get_currency_symbol(country)
    report_lines.append(f"\nCurrency symbol for {country}: {currency_symbol}")

    # Summary
    report_lines.append("\n=== CHANGES MADE ===")
    if changes_made:
        for change in changes_made:
            report_lines.append(f"  • {change}")
    else:
        report_lines.append("  No automatic changes were needed.")

    report = "\n".join(report_lines)
    return text, report


def build_editor_prompt(
    country: str,
    keywords: List[Dict],
    lsi_keywords: List[Dict],
    min_words: int,
    max_words: int,
) -> str:
    """Build the prompt for the LLM-based editor pass."""
    currency_symbol = get_currency_symbol(country)

    keyword_list = ""
    if keywords:
        for kw in keywords:
            keyword_list += f"  - '{kw['word']}' (x{kw.get('count', 1)})\n"
    if lsi_keywords:
        for kw in lsi_keywords:
            keyword_list += f"  - '{kw['word']}' (x{kw.get('count', 1)}) [LSI]\n"

    prompt = f"""You are a professional text editor with eight years of experience. Your job is to review the written text and then compare it to the specified requirements. If all requirements are met, you write a report on the number of keywords and the length of the text.

If the text doesn't meet the requirements, then:

Without asking the user, you take all necessary measures to ensure that the text 100% complies with the assignment.
Check that all keywords and additional keywords are included in the text.
If necessary, reduce the text size if it exceeds the assignment.
Remove all dashes and quotation marks; we are not allowed to use them.
Check the order and consistency of the headings so that they are identical to those specified in the assignment.
Be sure to check whether the currency of the country specified in the assignment is used (this should be an icon, not a letter. If not, use it everywhere and remove the letter. The icon should precede the amount, and the amount itself should be converted according to the exchange rate).

REQUIREMENTS:
- Country: {country}
- Currency symbol: {currency_symbol}
- Minimum words: {min_words}
- Maximum words: {max_words}
- Keywords and LSI words with required counts:
{keyword_list if keyword_list else "  None specified"}
- NO dashes (—, –, -) allowed
- NO quotation marks (" ", « », ' ') allowed
- NO personal pronouns (мы, наш, нам, нами, нас, вы, ваш, вам, вами, вас, ты, твой, тебе, тебя) allowed
- Text must be in Russian, but subheadings and keywords can be in the language specified in the assignment
- Currency symbol {currency_symbol} must be used before amounts, not letter codes

Output the corrected text followed by a brief editor report."""
    return prompt