ForcePilot/backend/package/yuxi/channels/adapters/zalo_oa/session.py

from __future__ import annotations

import time
from typing import Any

from yuxi.channels.models import ChatType

DEFAULT_SESSION_TTL_SEC = 3600
DEFAULT_BIND_TTL_SEC = 86400


class SessionRouter:
    def __init__(self, ttl_sec: int = DEFAULT_SESSION_TTL_SEC):
        self._ttl_sec = ttl_sec
        self._sessions: dict[str, dict[str, Any]] = {}
        self._bindings: dict[str, dict[str, Any]] = {}

    def resolve_thread_key(self, agent_id: str, follower_id: str, account_id: str = "default") -> str:
        return f"agent:{agent_id}:zalo_oa:{account_id}:{ChatType.DIRECT.value}:{follower_id}"

    def resolve_bind_key(self, agent_id: str, account_id: str = "default") -> str:
        return f"agent:{agent_id}:zalo_oa:{account_id}"

    def resolve_session_params(
        self,
        source_id: str,
        sender_name: str = "",
        avatar: str = "",
        dm_policy: str = "",
        account_id: str = "default",
    ) -> dict[str, Any]:
        params: dict[str, Any] = {
            "chat_type": ChatType.DIRECT,
            "channel_chat_id": source_id,
        }
        if sender_name:
            params["sender_name"] = sender_name
        if avatar:
            params["avatar"] = avatar
        if dm_policy:
            params["dm_policy"] = dm_policy
        if account_id:
            params["account_id"] = account_id
        return params

    def get_session(self, thread_key: str) -> dict[str, Any] | None:
        entry = self._sessions.get(thread_key)
        if entry is None:
            return None
        if time.time() - entry.get("created_at", 0) > self._ttl_sec:
            del self._sessions[thread_key]
            return None
        return entry.get("data")

    def set_session(self, thread_key: str, data: dict[str, Any]):
        self._sessions[thread_key] = {
            "data": data,
            "created_at": time.time(),
        }

    def clear_session(self, thread_key: str):
        self._sessions.pop(thread_key, None)

    def bind_agent_account(
        self,
        agent_id: str,
        account_id: str,
        bind_data: dict[str, Any] | None = None,
    ):
        bind_key = self.resolve_bind_key(agent_id, account_id)
        self._bindings[bind_key] = {
            "agent_id": agent_id,
            "account_id": account_id,
            "data": bind_data or {},
            "created_at": time.time(),
            "ttl_sec": DEFAULT_BIND_TTL_SEC,
        }

    def unbind_agent_account(self, agent_id: str, account_id: str):
        bind_key = self.resolve_bind_key(agent_id, account_id)
        self._bindings.pop(bind_key, None)

    def is_bound(self, agent_id: str, account_id: str) -> bool:
        bind_key = self.resolve_bind_key(agent_id, account_id)
        entry = self._bindings.get(bind_key)
        if entry is None:
            return False
        ttl = entry.get("ttl_sec", DEFAULT_BIND_TTL_SEC)
        if time.time() - entry.get("created_at", 0) > ttl:
            del self._bindings[bind_key]
            return False
        return True

    def get_binding(self, agent_id: str, account_id: str) -> dict[str, Any] | None:
        bind_key = self.resolve_bind_key(agent_id, account_id)
        entry = self._bindings.get(bind_key)
        if entry is None:
            return None
        ttl = entry.get("ttl_sec", DEFAULT_BIND_TTL_SEC)
        if time.time() - entry.get("created_at", 0) > ttl:
            del self._bindings[bind_key]
            return None
        return entry

    def get_bound_sessions(self, agent_id: str, account_id: str) -> list[dict[str, Any]]:
        bind_key = self.resolve_bind_key(agent_id, account_id)
        prefix = f"{bind_key}:{ChatType.DIRECT.value}:"
        sessions = []
        now = time.time()
        for key, entry in self._sessions.items():
            if key.startswith(prefix):
                if now - entry.get("created_at", 0) <= self._ttl_sec:
                    session_data = entry.get("data", {})
                    if session_data:
                        sessions.append(session_data)
        return sessions

    def cleanup_expired(self) -> int:
        now = time.time()
        expired_sessions = [k for k, v in self._sessions.items() if now - v.get("created_at", 0) > self._ttl_sec]
        for k in expired_sessions:
            del self._sessions[k]
        expired_bindings = [
            k
            for k, v in self._bindings.items()
            if now - v.get("created_at", 0) > v.get("ttl_sec", DEFAULT_BIND_TTL_SEC)
        ]
        for k in expired_bindings:
            del self._bindings[k]
        return len(expired_sessions) + len(expired_bindings)

    @property
    def session_count(self) -> int:
        return len(self._sessions)

    @property
    def binding_count(self) -> int:
        return len(self._bindings)


_router = SessionRouter()


def resolve_thread_key(agent_id: str, follower_id: str, account_id: str = "default") -> str:
    return _router.resolve_thread_key(agent_id, follower_id, account_id)


def resolve_session_params(
    source_id: str,
    sender_name: str = "",
    avatar: str = "",
    dm_policy: str = "",
    account_id: str = "default",
) -> dict[str, Any]:
    return _router.resolve_session_params(source_id, sender_name, avatar, dm_policy, account_id)


def bind_agent_account(agent_id: str, account_id: str, bind_data: dict[str, Any] | None = None):
    _router.bind_agent_account(agent_id, account_id, bind_data)


def unbind_agent_account(agent_id: str, account_id: str):
    _router.unbind_agent_account(agent_id, account_id)


def is_bound(agent_id: str, account_id: str) -> bool:
    return _router.is_bound(agent_id, account_id)


def get_session_router() -> SessionRouter:
    return _router