app/services/enrichment/mindmap.py

import asyncio
import re
import logging
from random import random
from time import time
import google.genai as genai
import json
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("[mindmap_service] google.genai module not found; mindmap generation disabled")

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

gemini_client = None

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

async def generate_mindmap(text: str) -> dict:
    if not text:
        return {}

    prompt = f"""
Bạn là chuyên gia tạo Sơ đồ tư duy. Hãy phân tích văn bản sau và tạo CẤU TRÚC JSON Mindmap.
Yêu cầu:
1. Xác định Ý chính làm Root.
2. Phân tách ý phụ thành nhánh con (tối đa 3 cấp).
3. Nhãn (label) ngắn gọn (< 7 từ).
4. Màu sắc (colorHex):
   - Root: "#6200EE"
   - Các nhánh con: sử dụng một trong các màu: "#F59E2B", "#2ECF9A", "#2F9BFF"
5. CHỈ TRẢ VỀ JSON, không giải thích thêm.
Cấu trúc JSON bắt buộc:
{{
  "root": {{
    "label": "Chủ đề",
    "colorHex": "#6200EE",
    "children": [
      {{
        "label": "Ý 1",
        "colorHex": "#F59E2B",
        "children": []
      }}
    ]
  }}
}}
Văn bản:
\"\"\"{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:
                resp = gemini_client.models.generate_content(
                    model=GEMINI_MODEL,
                    contents=prompt,
                )
                return resp.text or ""
            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:
                    is_server_error = False

                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"[mindmap_service] model overloaded (attempt {attempt}/{MAX_RETRIES}), retrying after {delay:.2f}s")
                        time.sleep(delay)
                        continue
                logging.exception(f"[mindmap_service] generate_mindmap call failed on attempt {attempt}: {e}")
                break

        if last_exc:
            raise last_exc
        return ""

    try:
        raw = await loop.run_in_executor(None, call)
        start = raw.find("{")
        end = raw.rfind("}")
        if start != -1 and end != -1:
            try:
                return json.loads(raw[start:end + 1])
            except Exception as e:
                logging.warning(f"[mindmap_service] Failed to parse mindmap JSON: {e}")
        else:
            logging.warning("[mindmap_service] Mindmap response has no JSON block")
    except GoogleAPIError as e:
        logging.error(f"[mindmap_service] Gemini API error: {e}")
    except Exception as e:
        logging.exception(f"[mindmap_service] generate_mindmap failed: {e}")

    # fallback: build a minimal mindmap from the text (best-effort)
    try:
        # use first sentence as root label (shortened)
        sentences = re.split(r"(?<=[.!?])\s+", text.strip())
        root_label = sentences[0] if sentences else "Topic"
        # shorten to ~8 words
        root_label = " ".join(root_label.split()[:8]).strip()

        children = []
        # take up to 3 next sentences as child labels
        for s in sentences[1:4]:
            label = " ".join(s.split()[:6]).strip()
            if label:
                children.append({
                    "label": label,
                    "colorHex": "#F59E2B",
                    "children": []
                })

        fallback = {
            "root": {
                "label": root_label or "Topic",
                "colorHex": "#6200EE",
                "children": children,
            }
        }
        logging.info("[mindmap_service] Returning fallback mindmap after errors")
        return fallback
    except Exception:
        return {}