app/services/enrichment/normalize.py

import asyncio
import os
import logging
import json
import re
import time
import random
from app.config import GEMINI_API_KEY, GEMINI_MODEL

try:
    import google.genai as genai
    try:
        from google.genai import errors as genai_errors
    except Exception:
        genai_errors = None
except Exception:
    genai = None
    genai_errors = None
    logging.warning("[normalize_service] google.genai module not found; normalization disabled")

try:
    from google.api_core.exceptions import GoogleAPIError
except Exception:
    GoogleAPIError = Exception

gemini_client = None

if not genai:
    logging.warning("[normalize_service] google.genai not available, normalization will be disabled")
elif not GEMINI_API_KEY:
    logging.warning("[normalize_service] GEMINI_API_KEY is not set, normalization will be disabled")
else:
    try:
        gemini_client = genai.Client(api_key=GEMINI_API_KEY)
        logging.info(f"[normalize_service] Initialized google.genai client with model={GEMINI_MODEL}")
    except Exception as e:
        logging.exception(f"[normalize_service] Failed to init google.genai client: {e}")
        gemini_client = None


async def normalize_text(raw_text: str) -> str:
    if not raw_text:
        return raw_text

    prompt = f"""
Bạn là một hệ thống Xử lý Hậu kỳ NLP (NLP Post-Processing) Tiếng Việt.
Đầu vào là văn bản thô (raw transcript), có thể thiếu dấu câu, sai chính tả, Lặp từ, lặp cụm từ, hoặc lặp cả đoạn do lỗi nhận dạng giọng nói.

Nhiệm vụ:
- Sửa lỗi chính tả do ASR.
- Thêm dấu câu phù hợp.
- Viết hoa đúng chuẩn tiếng Việt (đầu câu, tên riêng nếu suy luận được).
- Loại bỏ hoàn toàn các phần bị lặp (từ, cụm từ, câu hoặc đoạn), chỉ giữ MỘT phiên bản hợp lý.
- Giữ nguyên nội dung và ý nghĩa gốc, không rút gọn, không thêm thông tin mới.

YÊU CẦU ĐẦU RA:
- Chỉ trả về văn bản đã chuẩn hóa
- KHÔNG JSON, KHÔNG giải thích, KHÔNG markdown, KHÔNG bọc trong ``` hoặc bất kỳ ký tự đặc biệt nào, chỉ trả về plain text thuần túy.

Văn bản đầu vào:
\"\"\"{raw_text}\"\"\"
"""

    loop = asyncio.get_event_loop()

    MAX_RETRIES = 3
    BASE_DELAY = 1.0

    def call():
        last_exc = None
        for attempt in range(1, MAX_RETRIES + 1):
            try:
                if gemini_client:
                    resp = gemini_client.models.generate_content(
                        model=GEMINI_MODEL,
                        contents=prompt,
                    )
                    return resp.text or ""
                else:
                    model = genai.GenerativeModel(GEMINI_MODEL) if genai else None
                    if model:
                        resp = model.generate_content(prompt)
                        return getattr(resp, "text", "") or ""
                    return ""
            except Exception as e:
                last_exc = e
                msg = str(e)
                if "503" in msg or "UNAVAILABLE" in msg:
                    if attempt < MAX_RETRIES:
                        delay = BASE_DELAY * (2 ** (attempt - 1))
                        delay += random.uniform(0, 0.5 * delay)
                        logging.warning(
                            f"[normalize_service] model overloaded "
                            f"(attempt {attempt}/{MAX_RETRIES}), retrying after {delay:.2f}s"
                        )
                        time.sleep(delay)
                        continue
                logging.exception(
                    f"[normalize_service] normalize call failed on attempt {attempt}: {e}"
                )
                break

        if last_exc:
            raise last_exc
        return ""

    try:
        raw = await loop.run_in_executor(None, call)
        def is_valid_normalized(text: str, raw: str) -> bool:
            if not text:
                return False
            if len(text) < 50:
                return False
            if text.strip() in {".", "…"}:
                return False
            # không quá ngắn so với raw
            if len(text) < 0.3 * len(raw):
                return False
            return True

        if raw:
            text = raw.strip()

            # remove markdown
            text = re.sub(r"^```.*?\n", "", text, flags=re.DOTALL)
            text = re.sub(r"```$", "", text)

            # remove meta-text Gemini
            text = re.sub(
                r"^the corrected text.*?is as follows:\s*",
                "",
                text,
                flags=re.IGNORECASE | re.DOTALL,
            )

            text = text.strip('"').strip("'").strip()

            if is_valid_normalized(text, raw_text):
                return text

            logging.warning(
                "[normalize_service] Gemini returned invalid normalized text, falling back"
            )

    except GoogleAPIError as e:
        logging.error(f"[normalize_service] Gemini API error: {e}")
    except Exception as e:
        logging.exception(f"[normalize_service] normalize_text failed: {e}")

    # ===== fallback: best-effort local normalization =====
    try:
        text = raw_text.strip()
        text = re.sub(r"\s+", " ", text)

        if not text or len(text) < 50:
            return raw_text.strip()

        if text[-1] not in ".!?":
            text += "."

        def cap_sentences(s: str) -> str:
            parts = re.split(r'([.!?]\s+)', s)
            out = ""
            for i in range(0, len(parts), 2):
                sentence = parts[i].strip()
                sep = parts[i + 1] if i + 1 < len(parts) else ""
                if sentence:
                    sentence = sentence[0].upper() + sentence[1:]
                out += sentence + sep
            return out

        return cap_sentences(text)

    except Exception:
        return raw_text