conversation_memory.py

"""
ConversationSummaryBufferMemory — Qdrant-backed.

Giới hạn buffer tính theo số thành viên room (n):
  MAX_BUFFER      = 10n
  SUMMARIZE_COUNT =  6n  (số tin cũ nhất được tóm tắt khi vượt ngưỡng)
  KEEP_RECENT     =  4n  (số tin giữ lại trong buffer sau khi tóm tắt)

Fallback n = 20 nếu không kết nối được Supabase hoặc không phải room.
"""

import logging
import uuid
from datetime import datetime, timezone, timedelta
from typing import Optional

from langchain_core.messages import HumanMessage, SystemMessage
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams

from src.config import QDRANT_API_KEY, QDRANT_URL, SUPABASE_SERVICE_ROLE_KEY, SUPABASE_URL
from src.llm import llm
from src.tools.utils import format_created_at

logger = logging.getLogger(__name__)

_DEFAULT_N    = 20
_COLLECTION   = "conversation_memory"
_DUMMY_VECTOR = [0.0]

_qdrant_client: Optional[QdrantClient] = None
_sb_client = None


# ── Qdrant client ─────────────────────────────────────────────────────────────

def _get_qdrant() -> QdrantClient:
    global _qdrant_client
    if _qdrant_client is None:
        _qdrant_client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
        _ensure_collection(_qdrant_client)
    return _qdrant_client


def _ensure_collection(client: QdrantClient) -> None:
    existing = {c.name for c in client.get_collections().collections}
    if _COLLECTION not in existing:
        client.create_collection(
            collection_name=_COLLECTION,
            vectors_config=VectorParams(size=1, distance=Distance.DOT),
        )
        logger.info("Qdrant: collection '%s' created.", _COLLECTION)


def _point_id(conversation_id: str) -> str:
    return str(uuid.uuid5(uuid.NAMESPACE_DNS, f"conv:{conversation_id}"))


# ── Supabase client ───────────────────────────────────────────────────────────

def _get_sb():
    global _sb_client
    if _sb_client is None and SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY:
        try:
            from supabase import create_client
            _sb_client = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)
        except Exception:
            logger.exception("[Memory] Không khởi tạo được Supabase client.")
    return _sb_client


# ── Dynamic limits ────────────────────────────────────────────────────────────

def _get_member_count(conversation_id: str) -> int:
    """Lấy số thành viên trong room từ Supabase. Fallback về _DEFAULT_N."""
    if not conversation_id.startswith("room-"):
        return _DEFAULT_N

    sb = _get_sb()
    if sb is None:
        return _DEFAULT_N

    room_id = conversation_id.removeprefix("room-")
    try:
        res = (
            sb.table("room_members")
            .select("user_id", count="exact")
            .eq("room_id", room_id)
            .execute()
        )
        n = res.count or 0
        return n if n > 0 else _DEFAULT_N
    except Exception:
        logger.exception("[Memory] Lỗi lấy số thành viên room '%s'", room_id)
        return _DEFAULT_N


def _get_limits(conversation_id: str) -> tuple[int, int, int]:
    """Trả về (max_buffer, summarize_count, keep_recent) theo số thành viên n."""
    n = _get_member_count(conversation_id)
    return 10 * n, 6 * n, 4 * n


# ── Load / Save ───────────────────────────────────────────────────────────────

def load(conversation_id: str) -> tuple[str, list[dict]]:
    """Trả về (summary, buffer). Trả về ('', []) nếu chưa có."""
    if not QDRANT_URL:
        return "", []
    try:
        results = _get_qdrant().retrieve(
            collection_name=_COLLECTION,
            ids=[_point_id(conversation_id)],
            with_payload=True,
        )
        if results:
            payload = results[0].payload or {}
            summary = payload.get("summary", "") or ""
            buffer  = payload.get("buffer", []) or []
            return summary, buffer
        return "", []
    except Exception:
        logger.exception("[Memory] Lỗi load conversation_id='%s'", conversation_id)
        return "", []


def save(conversation_id: str, summary: str, buffer: list[dict]) -> None:
    if not QDRANT_URL:
        return
    try:
        _get_qdrant().upsert(
            collection_name=_COLLECTION,
            points=[PointStruct(
                id=_point_id(conversation_id),
                vector=_DUMMY_VECTOR,
                payload={
                    "conversation_id": conversation_id,
                    "summary":         summary,
                    "buffer":          buffer,
                },
            )],
        )
    except Exception:
        logger.exception("[Memory] Lỗi save conversation_id='%s'", conversation_id)


# ── Summarization ─────────────────────────────────────────────────────────────

def _summarize(existing_summary: str, messages: list[dict]) -> str:
    """Tóm tắt danh sách tin nhắn, kết hợp với summary hiện có."""
    history_text = "\n".join(
        f"{'User' if m['role'] == 'user' else 'Assistant'}: {m['content']}"
        for m in messages
    )

    prior_block = (
        f"Tóm tắt trước đó:\n{existing_summary}\n\n"
        if existing_summary
        else ""
    )

    system = SystemMessage(content=(
        "Bạn là trợ lý tóm tắt hội thoại. "
        "Tóm tắt ngắn gọn, giữ nguyên các thông tin quan trọng, "
        "sự kiện, tên gọi, và quyết định đã được đề cập. "
        "Chỉ trả về đoạn tóm tắt, không giải thích thêm."
    ))
    human = HumanMessage(content=(
        f"{prior_block}"
        f"Hội thoại cần tóm tắt:\n{history_text}\n\n"
        "Viết tóm tắt:"
    ))

    try:
        return llm.invoke([system, human]).content.strip()
    except Exception:
        logger.exception("[Memory] Lỗi khi summarize")
        return existing_summary


# ── Public API ────────────────────────────────────────────────────────────────

def _now_vn() -> str:
    _VN_TZ = timezone(timedelta(hours=7))
    return datetime.now(timezone.utc).astimezone(_VN_TZ).strftime("%d/%m/%Y %H:%M:%S")


def add_turn(conversation_id: str, sender_id: str, user_msg: str, ai_msg: str) -> None:
    """Thêm 1 lượt user+assistant vào buffer, trigger summarize nếu cần."""
    summary, buffer = load(conversation_id)
    max_buffer, summarize_count, keep_recent = _get_limits(conversation_id)

    ts = _now_vn()
    buffer.append({"role": "user",      "content": f"[{ts}] {sender_id}: {user_msg}"})
    buffer.append({"role": "assistant", "content": f"[{ts}] Assistant: {ai_msg}"})

    if len(buffer) > max_buffer:
        to_summarize = buffer[:summarize_count]
        buffer       = buffer[summarize_count:]
        if len(buffer) > keep_recent:
            buffer = buffer[-keep_recent:]
        logger.info(
            "[Memory] Buffer vượt %d, tóm tắt %d tin → giữ %d tin.",
            max_buffer, len(to_summarize), len(buffer),
        )
        summary = _summarize(summary, to_summarize)

    save(conversation_id, summary, buffer)


def seed_room(conversation_id: str, messages: list[dict]) -> None:
    """
    Seed Qdrant buffer từ danh sách tin nhắn Redis thô.
    Mỗi message được chuyển thành role='user', content='[ts UTC+7] name: content'.
    Nếu vượt max_buffer thì tự động summarize trước khi lưu.
    """
    _NAME_FIELDS = ["sender_username", "username", "u_username", "name", "u_name",
                    "senderName", "displayName", "display_name", "fullName", "sender_id"]

    def _get_field(m: dict, fields: list[str]) -> str:
        for f in fields:
            v = m.get(f)
            if v and str(v).strip():
                return str(v).strip()
        return ""

    buffer: list[dict] = []
    for msg in messages:
        name    = _get_field(msg, _NAME_FIELDS) or "unknown"
        content = msg.get("content") or msg.get("text") or msg.get("msg", "")
        ts      = format_created_at(msg.get("created_at") or msg.get("timestamp", ""))
        buffer.append({"role": "user", "content": f"[{ts}] {name}: {content}"})

    max_buffer, summarize_count, keep_recent = _get_limits(conversation_id)
    summary = ""
    while len(buffer) > max_buffer:
        to_summarize = buffer[:summarize_count]
        buffer       = buffer[summarize_count:]
        if len(buffer) > keep_recent:
            buffer = buffer[-keep_recent:]
        logger.info(
            "[Memory] seed_room: tóm tắt %d tin → giữ %d tin còn lại",
            len(to_summarize), len(buffer),
        )
        summary = _summarize(summary, to_summarize)

    save(conversation_id, summary, buffer)
    logger.info(
        "[Memory] seed_room '%s': lưu buffer=%d tin, summary=%s",
        conversation_id, len(buffer), "có" if summary else "không",
    )


def get_context(conversation_id: str) -> str:
    """Trả về chuỗi context (summary + buffer) để đưa vào prompt."""
    summary, buffer = load(conversation_id)

    parts: list[str] = []

    if summary:
        parts.append(f"[Tóm tắt lịch sử trước đó]\n{summary}")

    if buffer:
        recent_lines = "\n".join(
            f"{'User' if m['role'] == 'user' else 'Assistant'}: {m['content']}"
            for m in buffer
        )
        parts.append(f"[Tin nhắn gần đây]\n{recent_lines}")

    return "\n\n".join(parts) if parts else "(Chưa có lịch sử trò chuyện)"