LangBot/src/langbot/pkg/api/http/service/space.py

from __future__ import annotations

from collections import OrderedDict

from langbot.pkg.utils import httpclient
import typing
import datetime
import time
import sqlalchemy

from ....core import app
from ....entity.persistence import user
from ....entity.dto.space_model import SpaceModel


_CREDITS_CACHE_TTL_SECONDS = 60
_CREDITS_CACHE_MAX_ENTRIES = 4096


class SpaceService:
    """Service for interacting with LangBot Space API"""

    ap: app.Application
    _credits_cache: typing.Dict[str, typing.Tuple[int, float]]  # {user_email: (credits, timestamp)}

    def __init__(self, ap: app.Application) -> None:
        self.ap = ap
        self._credits_cache = OrderedDict()

    def _ordered_credits_cache(
        self,
    ) -> OrderedDict[str, tuple[int, float]]:
        if not isinstance(self._credits_cache, OrderedDict):
            # Preserve compatibility with tests and callers that seed the cache.
            self._credits_cache = OrderedDict(self._credits_cache)
        return self._credits_cache

    def _prune_credits_cache(self, now: float) -> None:
        cache = self._ordered_credits_cache()
        while cache:
            email = next(iter(cache))
            _, cached_at = cache[email]
            if now - cached_at < _CREDITS_CACHE_TTL_SECONDS:
                break
            cache.pop(email, None)

    def _get_space_config(self) -> typing.Dict[str, str]:
        """Get Space configuration from config file"""
        space_config = self.ap.instance_config.data.get('space', {})
        return {
            'url': space_config.get('url', 'https://space.langbot.app'),
            'oauth_authorize_url': space_config.get('oauth_authorize_url', 'https://space.langbot.app/auth/authorize'),
        }

    async def _get_user_by_email(self, user_email: str) -> user.User | None:
        result = await self.ap.persistence_mgr.execute_async(
            sqlalchemy.select(user.User).where(user.User.user == user_email)
        )
        result_list = result.all()
        return result_list[0] if result_list else None

    async def _ensure_valid_token(self, user_email: str) -> str | None:
        """Ensure access token is valid, refresh if expired. Returns valid access_token or None."""
        user_obj = await self._get_user_by_email(user_email)
        if not user_obj or user_obj.account_type != 'space':
            return None

        if not user_obj.space_access_token:
            return None

        # Check if token is expired (with 60s buffer)
        if user_obj.space_access_token_expires_at:
            if datetime.datetime.now() >= user_obj.space_access_token_expires_at - datetime.timedelta(seconds=60):
                # Token expired, try to refresh
                if user_obj.space_refresh_token:
                    try:
                        new_token = await self._refresh_and_save_token(user_obj)
                        return new_token
                    except Exception:
                        return None
                return None

        return user_obj.space_access_token

    async def _refresh_and_save_token(self, user_obj: user.User) -> str:
        """Refresh token and save to database"""
        token_data = await self.refresh_token(user_obj.space_refresh_token)
        access_token = token_data.get('access_token')
        expires_in = token_data.get('expires_in', 0)

        if not access_token:
            raise ValueError('Failed to refresh token')

        expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in) if expires_in > 0 else None

        await self.ap.persistence_mgr.execute_async(
            sqlalchemy.update(user.User)
            .where(user.User.user == user_obj.user)
            .values(
                space_access_token=access_token,
                space_access_token_expires_at=expires_at,
            )
        )

        return access_token

    # === Raw API calls (no token validation) ===

    def get_oauth_authorize_url(self, redirect_uri: str, state: str = '') -> str:
        """Get the Space OAuth authorization URL for redirect"""
        from urllib.parse import urlencode

        space_config = self._get_space_config()
        authorize_url = space_config['oauth_authorize_url']
        params = {'redirect_uri': redirect_uri}
        if state:
            params['state'] = state
        return f'{authorize_url}?{urlencode(params)}'

    async def exchange_oauth_code(self, code: str) -> typing.Dict:
        """Exchange OAuth authorization code for tokens"""
        from langbot.pkg.utils import constants

        space_config = self._get_space_config()
        space_url = space_config['url']

        session = httpclient.get_session()
        async with session.post(
            f'{space_url}/api/v1/accounts/oauth/token',
            json={'code': code, 'instance_id': constants.instance_id},
        ) as response:
            if response.status != 200:
                error = await httpclient.read_text_limited(response)
                raise ValueError(f'Failed to exchange OAuth code: {error}')
            data = await httpclient.read_json_limited(response)
            if data.get('code') != 0:
                raise ValueError(f'Failed to exchange OAuth code: {data.get("msg")}')
            return data.get('data', {})

    async def refresh_token(self, refresh_token: str) -> typing.Dict:
        """Refresh Space access token"""
        space_config = self._get_space_config()
        space_url = space_config['url']

        session = httpclient.get_session()
        async with session.post(
            f'{space_url}/api/v1/accounts/token/refresh', json={'refresh_token': refresh_token}
        ) as response:
            if response.status != 200:
                error = await httpclient.read_text_limited(response)
                raise ValueError(f'Failed to refresh token: {error}')
            data = await httpclient.read_json_limited(response)
            if data.get('code') != 0:
                raise ValueError(f'Failed to refresh token: {data.get("msg")}')
            return data.get('data', {})

    async def get_user_info_raw(self, access_token: str) -> typing.Dict:
        """Get user info from Space using access token (no validation)"""
        space_config = self._get_space_config()
        space_url = space_config['url']

        session = httpclient.get_session()
        async with session.get(
            f'{space_url}/api/v1/accounts/me', headers={'Authorization': f'Bearer {access_token}'}
        ) as response:
            if response.status != 200:
                error = await httpclient.read_text_limited(response)
                raise ValueError(f'Failed to get user info: {error}')
            data = await httpclient.read_json_limited(response)
            if data.get('code') != 0:
                raise ValueError(f'Failed to get user info: {data.get("msg")}')
            return data.get('data', {})

    # === API calls with token validation ===

    async def get_user_info(self, user_email: str) -> typing.Dict | None:
        """Get user info from Space (with token validation)"""
        access_token = await self._ensure_valid_token(user_email)
        if not access_token:
            return None
        return await self.get_user_info_raw(access_token)

    async def get_credits(self, user_email: str, force_refresh: bool = False) -> int | None:
        """Get Space credits for user with caching (60s TTL)"""
        now = time.time()
        cached_fallback = self._credits_cache.get(user_email)
        self._prune_credits_cache(now)

        if not force_refresh and user_email in self._credits_cache:
            credits, ts = self._credits_cache[user_email]
            if now - ts < _CREDITS_CACHE_TTL_SECONDS:
                return credits

        try:
            info = await self.get_user_info(user_email)
            if info is None:
                return None
            credits = info.get('credits')
            if credits is not None:
                cache = self._ordered_credits_cache()
                cache.pop(user_email, None)
                if len(cache) >= _CREDITS_CACHE_MAX_ENTRIES:
                    cache.popitem(last=False)
                cache[user_email] = (credits, time.time())
            return credits
        except Exception:
            return cached_fallback[0] if cached_fallback is not None else None

    async def get_models(self) -> typing.List[SpaceModel]:
        """Get models from Space"""

        space_config = self._get_space_config()
        space_url = space_config['url']

        session = httpclient.get_session()
        async with session.get(f'{space_url}/api/v1/models', params={'page_size': 100}) as response:
            if response.status != 200:
                error = await httpclient.read_text_limited(response)
                raise ValueError(f'Failed to get models: {error}')
            data = await httpclient.read_json_limited(response)
            if data.get('code') != 0:
                raise ValueError(f'Failed to get models: {data.get("msg")}')
            models_data = data.get('data', {}).get('models', [])
            return [SpaceModel.model_validate(model_dict) for model_dict in models_data]