LangBot/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py

"""Embed widget routes - serve embeddable chat widget for external websites.

All user-facing URLs are keyed by **bot_uuid** (not pipeline_uuid) so that
internal pipeline identifiers are never exposed to end-users.  Each handler
resolves the bot_uuid to the owning ``web_page_bot`` RuntimeBot and extracts
the bound pipeline_uuid for internal routing.
"""

import asyncio
import datetime
import json
import logging
import uuid
import hmac
import hashlib
import time
import re
import httpx

import quart

from ... import group
from ......utils import httpclient, paths
from ......platform.sources.websocket_manager import WebSocketScope, is_valid_session_id, ws_connection_manager
from .websocket_chat import create_scoped_duplex_tasks, wait_for_duplex_tasks

logger = logging.getLogger(__name__)
_AUTH_TIMEOUT_SECONDS = 10.0

# Cache the widget template content
_widget_template_cache: str | None = None
_logo_bytes_cache: bytes | None = None


def _is_valid_uuid(s: str) -> bool:
    return bool(re.match(r'^[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}$', s))


def _get_widget_template() -> str:
    """Load and cache the widget JS template."""
    global _widget_template_cache
    if _widget_template_cache is None:
        template_path = paths.get_resource_path('templates/embed/widget.js')
        with open(template_path, 'r', encoding='utf-8') as f:
            _widget_template_cache = f.read()
    return _widget_template_cache


def _get_logo_bytes() -> bytes:
    """Load and cache the logo image."""
    global _logo_bytes_cache
    if _logo_bytes_cache is None:
        logo_path = paths.get_resource_path('templates/embed/logo.webp')
        with open(logo_path, 'rb') as f:
            _logo_bytes_cache = f.read()
    return _logo_bytes_cache


@group.group_class('embed', '/api/v1/embed')
class EmbedRouterGroup(group.RouterGroup):
    # -- helpers -------------------------------------------------------------

    async def _resolve_bot(self, bot_uuid: str):
        """Resolve *bot_uuid* to ``(runtime_bot, pipeline_uuid)``.

        Returns ``(None, None)`` when the bot does not exist, is not a
        ``web_page_bot``, is disabled, or has no pipeline bound.
        """
        bot = await self.ap.platform_mgr.resolve_public_bot(bot_uuid)
        if (
            bot is not None
            and bot.bot_entity.adapter == 'web_page_bot'
            and bot.bot_entity.enable
            and bot.bot_entity.use_pipeline_uuid
        ):
            return bot, bot.bot_entity.use_pipeline_uuid
        return None, None

    @staticmethod
    def _get_bot_config(runtime_bot) -> dict:
        return runtime_bot.bot_entity.adapter_config

    def _verify_session_token_value(self, token: str, runtime_bot) -> bool:
        config = self._get_bot_config(runtime_bot)
        secret = config.get('turnstile_secret_key', '')
        if not secret:
            return True
        try:
            ts_str, mac = token.split('.', 1)
            ts = float(ts_str)
            if time.time() - ts > 86400:
                return False
            expected_mac = hmac.new(secret.encode(), f'{ts_str}'.encode(), hashlib.sha256).hexdigest()
            return hmac.compare_digest(mac, expected_mac)
        except Exception:
            return False

    async def _verify_session_token(self, request, runtime_bot) -> bool:
        auth_header = request.headers.get('Authorization', '')
        token = auth_header[7:] if auth_header.startswith('Bearer ') else ''
        return self._verify_session_token_value(token, runtime_bot)

    async def _authenticate_websocket(self, runtime_bot) -> None:
        """Require the embed session token as the first WebSocket frame."""

        raw_message = await asyncio.wait_for(quart.websocket.receive(), timeout=_AUTH_TIMEOUT_SECONDS)
        payload = await asyncio.to_thread(json.loads, raw_message)
        if not isinstance(payload, dict) or payload.get('type') != 'authenticate':
            raise ValueError('Authentication is required')
        token = str(payload.get('token') or '')
        if not self._verify_session_token_value(token, runtime_bot):
            raise ValueError('Authentication is required')

    async def _assert_execution_active(self, runtime_bot) -> None:
        context = runtime_bot.execution_context
        await self.ap.workspace_service.get_execution_binding(
            context.workspace_uuid,
            expected_generation=context.placement_generation,
        )

    async def _resolve_connected_bot(self, owner_bot, pipeline_uuid: str):
        """Re-resolve mutable bot state before every public message."""
        current_bot, current_pipeline_uuid = await self._resolve_bot(owner_bot.bot_entity.uuid)
        if current_bot is None or current_pipeline_uuid != pipeline_uuid:
            raise RuntimeError('Bot is unavailable')

        owner_context = owner_bot.execution_context
        current_context = current_bot.execution_context
        if (
            current_context.instance_uuid,
            current_context.workspace_uuid,
            current_context.placement_generation,
        ) != (
            owner_context.instance_uuid,
            owner_context.workspace_uuid,
            owner_context.placement_generation,
        ):
            raise RuntimeError('Bot is unavailable')
        await self._assert_execution_active(current_bot)
        return current_bot

    # -- routes --------------------------------------------------------------

    async def initialize(self) -> None:
        @self.route('/<bot_uuid>/turnstile/verify', methods=['POST'], auth_type=group.AuthType.NONE)
        async def verify_turnstile(bot_uuid: str) -> str:
            if not _is_valid_uuid(bot_uuid):
                return self.http_status(400, -1, 'Invalid bot_uuid format')
            runtime_bot, pipeline_uuid = await self._resolve_bot(bot_uuid)
            if runtime_bot is None:
                return self.http_status(404, -1, 'Bot not found or not available')
            try:
                data = await quart.request.get_json()
                token = data.get('token')
                if not token:
                    return self.http_status(400, -1, 'Token is required')

                config = self._get_bot_config(runtime_bot)
                secret = config.get('turnstile_secret_key', '')
                if not secret:
                    ts = time.time()
                    return self.success(data={'token': f'{ts}.dummy'})

                async with httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks()) as client:
                    resp = await client.post(
                        'https://challenges.cloudflare.com/turnstile/v0/siteverify',
                        data={'secret': secret, 'response': token},
                    )
                    result = await httpclient.parse_json_response(resp)

                if not result.get('success'):
                    return self.http_status(403, -1, 'Turnstile verification failed')

                ts = time.time()
                mac = hmac.new(secret.encode(), f'{ts}'.encode(), hashlib.sha256).hexdigest()
                session_token = f'{ts}.{mac}'

                return self.success(data={'token': session_token})

            except Exception as e:
                logger.error(f'Turnstile verify failed: {e}', exc_info=True)
                return self.http_status(500, -1, 'Internal server error')

        @self.route('/<bot_uuid>/widget.js', methods=['GET'], auth_type=group.AuthType.NONE)
        async def serve_widget(bot_uuid: str) -> quart.Response:
            """Serve the embed widget JavaScript with injected configuration."""
            if not _is_valid_uuid(bot_uuid):
                return self.http_status(400, -1, 'Invalid bot_uuid format')
            runtime_bot, pipeline_uuid = await self._resolve_bot(bot_uuid)
            if runtime_bot is None:
                return quart.Response(
                    '// Bot not found or not available', status=404, content_type='application/javascript'
                )
            try:
                template = _get_widget_template()
            except FileNotFoundError:
                return quart.Response('// Widget template not found', status=404, content_type='application/javascript')

            base_url = quart.request.host_url.rstrip('/')
            webhook_prefix = self.ap.instance_config.data.get('api', {}).get('webhook_prefix', '')
            if webhook_prefix:
                base_url = webhook_prefix.rstrip('/')

            if not re.match(r'^https?://[a-zA-Z0-9._:/-]+$', base_url):
                base_url = quart.request.host_url.rstrip('/')

            config = self._get_bot_config(runtime_bot)
            site_key = config.get('turnstile_site_key', '')
            locale = config.get('language', 'en_US') or 'en_US'
            bubble_icon = config.get('bubble_icon', 'logo') or 'logo'
            widget_js = template.replace('__LANGBOT_TURNSTILE_SITE_KEY__', site_key)
            widget_js = widget_js.replace('__LANGBOT_BOT_UUID__', bot_uuid)
            widget_js = widget_js.replace('__LANGBOT_BASE_URL__', base_url)
            widget_js = widget_js.replace('__LANGBOT_LOCALE__', locale)
            widget_js = widget_js.replace('__LANGBOT_BUBBLE_ICON__', bubble_icon)

            response = quart.Response(widget_js, content_type='application/javascript; charset=utf-8')
            response.headers['Cache-Control'] = 'public, max-age=300'
            return response

        @self.route('/logo', methods=['GET'], auth_type=group.AuthType.NONE)
        async def serve_logo() -> quart.Response:
            """Serve the LangBot logo for the embed widget."""
            try:
                logo_data = _get_logo_bytes()
            except FileNotFoundError:
                return quart.Response('', status=404)

            response = quart.Response(logo_data, content_type='image/webp')
            response.headers['Cache-Control'] = 'public, max-age=86400'
            return response

        @self.route('/<bot_uuid>/messages/<session_type>', methods=['GET'], auth_type=group.AuthType.NONE)
        async def get_embed_messages(bot_uuid: str, session_type: str) -> str:
            if not _is_valid_uuid(bot_uuid):
                return self.http_status(400, -1, 'Invalid bot_uuid format')
            runtime_bot, pipeline_uuid = await self._resolve_bot(bot_uuid)
            if runtime_bot is None:
                return self.http_status(404, -1, 'Bot not found or not available')
            if not await self._verify_session_token(quart.request, runtime_bot):
                return self.http_status(403, -1, 'Unauthorized or session expired')
            try:
                if session_type not in ['person', 'group']:
                    return self.http_status(400, -1, 'session_type must be person or group')

                session_id = quart.request.args.get('session_id', '')
                if not is_valid_session_id(session_id):
                    return self.http_status(400, -1, 'Valid session_id is required')

                proxy_bot = await self.ap.platform_mgr.get_websocket_proxy_bot(runtime_bot.execution_context)
                websocket_adapter = proxy_bot.adapter
                if not websocket_adapter:
                    return self.http_status(404, -1, 'WebSocket adapter not found')

                messages = websocket_adapter.get_websocket_messages(pipeline_uuid, session_type, session_id)
                return self.success(data={'messages': messages})

            except Exception as e:
                logger.error(f'Failed to get embed messages: {e}', exc_info=True)
                return self.http_status(500, -1, 'Internal server error')

        @self.route('/<bot_uuid>/reset/<session_type>', methods=['POST'], auth_type=group.AuthType.NONE)
        async def reset_embed_session(bot_uuid: str, session_type: str) -> str:
            if not _is_valid_uuid(bot_uuid):
                return self.http_status(400, -1, 'Invalid bot_uuid format')
            runtime_bot, pipeline_uuid = await self._resolve_bot(bot_uuid)
            if runtime_bot is None:
                return self.http_status(404, -1, 'Bot not found or not available')
            if not await self._verify_session_token(quart.request, runtime_bot):
                return self.http_status(403, -1, 'Unauthorized or session expired')
            try:
                if session_type not in ['person', 'group']:
                    return self.http_status(400, -1, 'session_type must be person or group')

                session_id = quart.request.args.get('session_id', '')
                if not is_valid_session_id(session_id):
                    return self.http_status(400, -1, 'Valid session_id is required')

                proxy_bot = await self.ap.platform_mgr.get_websocket_proxy_bot(runtime_bot.execution_context)
                websocket_adapter = proxy_bot.adapter
                if not websocket_adapter:
                    return self.http_status(404, -1, 'WebSocket adapter not found')

                websocket_adapter.reset_session(pipeline_uuid, session_type, session_id)
                return self.success(data={'message': 'Session reset successfully'})

            except Exception as e:
                logger.error(f'Failed to reset embed session: {e}', exc_info=True)
                return self.http_status(500, -1, 'Internal server error')

        @self.route('/<bot_uuid>/feedback', methods=['POST'], auth_type=group.AuthType.NONE)
        async def submit_feedback(bot_uuid: str) -> str:
            if not _is_valid_uuid(bot_uuid):
                return self.http_status(400, -1, 'Invalid bot_uuid format')
            runtime_bot, pipeline_uuid = await self._resolve_bot(bot_uuid)
            if runtime_bot is None:
                return self.http_status(404, -1, 'Bot not found or not available')
            if not await self._verify_session_token(quart.request, runtime_bot):
                return self.http_status(403, -1, 'Unauthorized or session expired')
            try:
                data = await quart.request.get_json()
                message_id = data.get('message_id', '')
                feedback_type = data.get('feedback_type')

                if feedback_type not in (1, 2, 3):
                    return self.http_status(400, -1, 'feedback_type must be 1 (like), 2 (dislike), or 3 (cancel)')

                feedback_id = f'embed_{uuid.uuid4().hex[:12]}'

                await self.ap.monitoring_service.record_feedback(
                    runtime_bot.execution_context,
                    feedback_id=feedback_id,
                    feedback_type=feedback_type,
                    bot_id=runtime_bot.bot_entity.uuid,
                    bot_name=runtime_bot.bot_entity.name or bot_uuid,
                    pipeline_id=pipeline_uuid,
                    message_id=str(message_id),
                    platform='web_page_bot',
                )

                return self.success(data={'feedback_id': feedback_id})

            except Exception as e:
                logger.error(f'Failed to record feedback: {e}', exc_info=True)
                return self.http_status(500, -1, 'Internal server error')

        # -- Embed WebSocket endpoint ----------------------------------------

        @self.quart_app.websocket(self.path + '/<bot_uuid>/ws/connect')
        async def embed_websocket_connect(bot_uuid: str):
            """WebSocket connection for embed widget, keyed by bot_uuid."""
            await quart.websocket.accept()
            if not _is_valid_uuid(bot_uuid):
                await quart.websocket.send(json.dumps({'type': 'error', 'message': 'Invalid bot_uuid format'}))
                return

            runtime_bot, pipeline_uuid = await self._resolve_bot(bot_uuid)
            if runtime_bot is None:
                await quart.websocket.send(json.dumps({'type': 'error', 'message': 'Bot not found or not available'}))
                return

            session_type = quart.websocket.args.get('session_type', 'person')
            if session_type not in ['person', 'group']:
                await quart.websocket.send(
                    json.dumps({'type': 'error', 'message': 'session_type must be person or group'})
                )
                return

            session_id = quart.websocket.args.get('session_id', '')
            if not is_valid_session_id(session_id):
                await quart.websocket.send(json.dumps({'type': 'error', 'message': 'Valid session_id is required'}))
                return

            try:
                await self._authenticate_websocket(runtime_bot)
                await self._assert_execution_active(runtime_bot)
            except Exception:
                await quart.websocket.send(json.dumps({'type': 'error', 'message': 'Unauthorized'}))
                return

            try:
                proxy_bot = await self.ap.platform_mgr.get_websocket_proxy_bot(runtime_bot.execution_context)
                websocket_adapter = proxy_bot.adapter
                if not websocket_adapter:
                    await quart.websocket.send(json.dumps({'type': 'error', 'message': 'WebSocket adapter not found'}))
                    return

                connection = await ws_connection_manager.add_connection(
                    websocket=quart.websocket._get_current_object(),
                    scope=WebSocketScope.from_context(runtime_bot.execution_context),
                    pipeline_uuid=pipeline_uuid,
                    session_type=session_type,
                    session_id=session_id,
                    metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
                    send_queue_size=(
                        self.ap.instance_config.data.get('system', {})
                        .get('websocket_retention', {})
                        .get('send_queue_size', 100)
                    ),
                    max_connections=(
                        self.ap.instance_config.data.get('system', {})
                        .get('websocket_retention', {})
                        .get('max_connections', 1024)
                    ),
                    max_connections_per_workspace=(
                        self.ap.instance_config.data.get('system', {})
                        .get('websocket_retention', {})
                        .get('max_connections_per_workspace', 32)
                    ),
                )

                await quart.websocket.send(
                    json.dumps(
                        {
                            'type': 'connected',
                            'connection_id': connection.connection_id,
                            'bot_uuid': bot_uuid,
                            'session_type': session_type,
                            'timestamp': connection.created_at.isoformat(),
                        }
                    )
                )

                logger.debug(
                    f'Embed WebSocket connected: {connection.connection_id} '
                    f'(bot={bot_uuid}, pipeline={pipeline_uuid}, session_type={session_type})'
                )

                receive_task, send_task = create_scoped_duplex_tasks(
                    self._handle_receive(
                        connection,
                        websocket_adapter,
                        runtime_bot,
                        pipeline_uuid,
                    ),
                    self._handle_send(connection),
                    runtime_bot.execution_context.workspace_uuid,
                )

                try:
                    await wait_for_duplex_tasks(receive_task, send_task)
                except Exception as e:
                    logger.error(f'Embed WebSocket task error: {e}')
                finally:
                    await ws_connection_manager.remove_connection(connection.connection_id)

            except Exception as e:
                logger.error(f'Embed WebSocket connection error: {e}', exc_info=True)
                try:
                    await quart.websocket.send(json.dumps({'type': 'error', 'message': 'Internal server error'}))
                except Exception:
                    pass

    # -- WebSocket receive/send helpers --------------------------------------

    async def _handle_receive(self, connection, websocket_adapter, owner_bot, pipeline_uuid: str):
        try:
            while connection.is_active:
                message = await quart.websocket.receive()
                await ws_connection_manager.update_activity(connection.connection_id)

                try:
                    data = await asyncio.to_thread(json.loads, message)
                    message_type = data.get('type', 'message')

                    if message_type == 'ping':
                        await connection.send_queue.put(
                            {'type': 'pong', 'timestamp': datetime.datetime.now().isoformat()}
                        )
                    elif message_type == 'message':
                        try:
                            current_bot = await self._resolve_connected_bot(owner_bot, pipeline_uuid)
                        except Exception:
                            await connection.send_queue.put({'type': 'error', 'message': 'Bot is unavailable'})
                            break
                        await websocket_adapter.handle_websocket_message(connection, data, owner_bot=current_bot)
                    elif message_type == 'disconnect':
                        break

                except json.JSONDecodeError:
                    await connection.send_queue.put({'type': 'error', 'message': 'Invalid JSON format'})

        except Exception as e:
            logger.error(f'Embed receive error: {e}', exc_info=True)
        finally:
            connection.is_active = False
            try:
                connection.send_queue.put_nowait(None)
            except asyncio.QueueFull:
                pass

    async def _handle_send(self, connection):
        try:
            while connection.is_active or not connection.send_queue.empty():
                try:
                    message = await asyncio.wait_for(connection.send_queue.get(), timeout=1.0)
                    if message is None:
                        break
                    encoded = await asyncio.to_thread(json.dumps, message)
                    await quart.websocket.send(encoded)
                except asyncio.TimeoutError:
                    continue
        except Exception as e:
            logger.error(f'Embed send error: {e}', exc_info=True)
        finally:
            connection.is_active = False