app/services/enrichment/title_keywords.py

import os
import logging
import json
import asyncio
import time
import random
import re
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("[keywords_service] google.genai module not found; keyword extraction disabled")

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

gemini_client = None

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


async def extract_title_and_keywords(text: str) -> tuple[str | None, list[str]]:
    if not text or not text.strip():
        return None, []

    if not gemini_client and not genai:
        # AI not available → safe fallback
        return None, []

    prompt = f"""
Bạn là một hệ thống Xử lý Hậu kỳ NLP (NLP Post-Processing) Tiếng Việt.

Nhiệm vụ:
1. Sinh **tiêu đề (title)** ngắn gọn phản ánh đúng chủ đề chính của văn bản:
   - Độ dài tối đa **10 từ**
   - Mang tính mô tả, trung tính, phù hợp làm tiêu đề ghi chú (note)
   - KHÔNG giật tít, KHÔNG suy diễn quá mức

2. Rút trích các **từ khóa quan trọng** phản ánh đúng **chủ đề và nội dung chính** của văn bản.
- Mỗi từ khóa dài từ **1–4 từ**.
- Ưu tiên danh từ, cụm danh từ, thuật ngữ, khái niệm chính.
- Loại bỏ từ chung chung, từ đệm, từ cảm thán, từ lặp nghĩa.
- KHÔNG diễn giải, KHÔNG tóm tắt, KHÔNG chuẩn hóa lại văn bản.
- KHÔNG tạo từ khóa không xuất hiện hoặc không suy luận hợp lý từ văn bản.

Quy tắc:
- Số lượng từ khóa: 3–10 (tùy độ dài và nội dung văn bản).
- Giữ nguyên chữ thường/hoa theo cách viết phổ biến.
- KHÔNG trùng lặp từ khóa.
- KHÔNG sắp xếp theo bảng chữ cái; ưu tiên theo mức độ quan trọng.

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

YÊU CẦU ĐẦU RA:
- Chỉ trả về **JSON hợp lệ**
- KHÔNG giải thích
- KHÔNG markdown
- KHÔNG thêm trường khác ngoài schema dưới đây

Cấu trúc JSON bắt buộc:
{{
    "title": "Tiêu đề ngắn gọn",
    "keywords": ["Từ khóa 1", "Từ khóa 2", "..."]
}}
"""

    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 getattr(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
                is_server_error = False
                try:
                    if genai_errors and isinstance(e, genai_errors.ServerError):
                        is_server_error = True
                except Exception:
                    pass

                msg = str(e)
                if "503" in msg or "UNAVAILABLE" in msg or is_server_error:
                    if attempt < MAX_RETRIES:
                        delay = BASE_DELAY * (2 ** (attempt - 1))
                        delay = delay + random.uniform(0, 0.5 * delay)
                        logging.warning(f"[keywords_service] model overloaded (attempt {attempt}/{MAX_RETRIES}), retrying after {delay:.2f}s")
                        time.sleep(delay)
                        continue
                logging.exception(f"[keywords_service] extract_keywords call failed on attempt {attempt}: {e}")
                break

        if last_exc:
            raise last_exc
        return ""

    try:
        raw = await loop.run_in_executor(None, call)
        title, keywords = _parse_response(raw)
        return title, keywords
    except GoogleAPIError as e:
        logging.error("[title_keywords_service] Gemini API error: %s", e)
    except Exception as e:
        logging.exception("[title_keywords_service] extract failed: %s", e)

    return None, []

def _parse_response(raw: str) -> tuple[str | None, list[str]]:
    if not raw:
        return None, []

    raw = raw.strip()

    # Try extracting JSON block
    start = raw.find("{")
    end = raw.rfind("}")

    if start != -1 and end != -1 and end > start:
        raw_json = raw[start : end + 1]
    else:
        raw_json = raw

    try:
        parsed = json.loads(raw_json)
    except Exception as e:
        logging.warning(
            "[title_keywords_service] Failed to parse JSON: %s | raw=%r",
            e,
            raw[:300],
        )
        return None, []

    title = parsed.get("title")
    keywords = parsed.get("keywords")

    # Validate title
    if not isinstance(title, str) or not title.strip():
        title = None
    else:
        title = title.strip()

    # Validate keywords
    if not isinstance(keywords, list):
        keywords = []
    else:
        keywords = [
            k.strip()
            for k in keywords
            if isinstance(k, str) and k.strip()
        ]

    return title, keywords