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

from __future__ import annotations

import asyncio
import datetime
import functools
import os
import re
from pathlib import Path
from typing import Any

import sqlalchemy

from ....core import app
from ....entity.persistence import bstorage as persistence_bstorage
from ....entity.persistence import monitoring as persistence_monitoring
from ..authz import WorkspaceRequiredError
from ..context import ExecutionContext
from .tenant import TenantContext, require_workspace_uuid


LOG_FILE_PATTERN = re.compile(r'^langbot-(\d{4}-\d{2}-\d{2})\.log(?:\.\d+)?$')
DEFAULT_UPLOAD_FILE_RETENTION_DAYS = 7
DEFAULT_LOG_RETENTION_DAYS = 3
DEFAULT_MAX_FILES_PER_RUN = 1000
HARD_MAX_FILES_PER_RUN = 10000
UPLOAD_OWNER_TYPES = ('upload_image', 'upload_document', 'upload')


def _workspace_scope(method):
    """Bind maintenance work to a Workspace without spanning external I/O."""

    @functools.wraps(method)
    async def wrapped(self, context, *args, **kwargs):
        workspace_uuid = require_workspace_uuid(context)
        persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
        tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
        cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
        if cloud_runtime:
            if not callable(tenant_scope):
                raise RuntimeError('Cloud maintenance requires an explicit tenant scope')
            async with tenant_scope(workspace_uuid):
                return await method(self, context, *args, **kwargs)
        return await method(self, context, *args, **kwargs)

    return wrapped


class MaintenanceService:
    """Storage maintenance and diagnostics."""

    ap: app.Application

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

    def _max_files_per_run(self) -> int:
        cleanup_cfg = (
            getattr(getattr(self.ap, 'instance_config', None), 'data', {}).get('storage', {}).get('cleanup', {})
        )
        value = self._positive_int(
            cleanup_cfg.get('max_files_per_run', DEFAULT_MAX_FILES_PER_RUN),
            DEFAULT_MAX_FILES_PER_RUN,
            'storage.cleanup.max_files_per_run',
        )
        return min(value, HARD_MAX_FILES_PER_RUN)

    @_workspace_scope
    async def cleanup_expired_files(self, context: ExecutionContext) -> dict[str, int]:
        if not isinstance(context, ExecutionContext):
            raise WorkspaceRequiredError('Storage cleanup requires an ExecutionContext')
        require_workspace_uuid(context)
        cleanup_cfg = self.ap.instance_config.data.get('storage', {}).get('cleanup', {})
        upload_retention_days = self._positive_int(
            cleanup_cfg.get('uploaded_file_retention_days'),
            DEFAULT_UPLOAD_FILE_RETENTION_DAYS,
            'storage.cleanup.uploaded_file_retention_days',
        )
        log_retention_days = self._positive_int(
            cleanup_cfg.get('log_retention_days'),
            DEFAULT_LOG_RETENTION_DAYS,
            'storage.cleanup.log_retention_days',
        )

        return {
            'uploaded_files': await self._cleanup_expired_uploaded_files(context, upload_retention_days),
            'log_files': await asyncio.to_thread(
                self._cleanup_expired_log_files,
                log_retention_days,
            )
            if await self._is_oss_singleton(context)
            else 0,
        }

    async def get_storage_analysis(self, context: TenantContext) -> dict[str, Any]:
        require_workspace_uuid(context)
        cleanup_cfg = self.ap.instance_config.data.get('storage', {}).get('cleanup', {})
        upload_retention_days = self._positive_int(
            cleanup_cfg.get('uploaded_file_retention_days'),
            DEFAULT_UPLOAD_FILE_RETENTION_DAYS,
            'storage.cleanup.uploaded_file_retention_days',
        )
        log_retention_days = self._positive_int(
            cleanup_cfg.get('log_retention_days'),
            DEFAULT_LOG_RETENTION_DAYS,
            'storage.cleanup.log_retention_days',
        )

        database_cfg = self.ap.instance_config.data.get('database', {})
        database_type = database_cfg.get('use', 'sqlite')
        database_path = (
            Path(database_cfg.get('sqlite', {}).get('path', 'data/langbot.db')) if database_type == 'sqlite' else None
        )
        is_oss_singleton = await self._is_oss_singleton(context)
        if is_oss_singleton:
            roots: list[tuple[str, Path | None]] = [
                ('database', database_path),
                ('logs', Path('data/logs')),
                ('storage', Path('data/storage')),
                ('vector_store', Path('data/chroma')),
                ('plugins', Path('data/plugins')),
                ('mcp', Path('data/mcp')),
                ('temp', Path('data/temp')),
            ]
        else:
            scoped_storage_path = Path('data/storage') / self.ap.storage_mgr.scoped_prefix(context)
            roots = [('storage', scoped_storage_path)]

        sections = await asyncio.to_thread(self._collect_sections, roots)

        monitoring_counts = await self._monitoring_counts(context)
        binary_storage = await self._binary_storage_stats(context)
        upload_candidates = await self._expired_uploaded_candidates(context, upload_retention_days)
        log_candidates = (
            await asyncio.to_thread(
                self._expired_log_candidates,
                log_retention_days,
            )
            if is_oss_singleton
            else []
        )

        return {
            'generated_at': datetime.datetime.now(datetime.timezone.utc).isoformat(),
            'cleanup_policy': {
                'uploaded_file_retention_days': upload_retention_days,
                'log_retention_days': log_retention_days,
            },
            'sections': sections,
            'database': {
                'type': database_type,
                'monitoring_counts': monitoring_counts,
                'binary_storage': binary_storage,
            },
            'cleanup_candidates': {
                'uploaded_files': upload_candidates,
                'log_files': log_candidates,
            },
            'tasks': self.ap.task_mgr.get_stats() if is_oss_singleton and self.ap.task_mgr else {},
        }

    def _collect_sections(
        self,
        roots: list[tuple[str, Path | None]],
    ) -> list[dict[str, Any]]:
        sections = []
        for key, path in roots:
            sections.append(
                {
                    'key': key,
                    'path': str(path) if path else '',
                    'exists': path.exists() if path else False,
                    'size_bytes': self._path_size(path) if path else 0,
                    'file_count': self._file_count(path) if path else 0,
                }
            )
        return sections

    async def _is_oss_singleton(self, context: TenantContext) -> bool:
        try:
            await self.ap.workspace_service.get_local_execution_binding(
                require_workspace_uuid(context),
                expected_generation=getattr(context, 'placement_generation', None),
            )
        except Exception:
            return False
        return True

    async def _cleanup_expired_uploaded_files(
        self,
        context: ExecutionContext,
        retention_days: int,
    ) -> int:
        provider = self.ap.storage_mgr.storage_provider
        provider_name = provider.__class__.__name__
        if provider_name == 'LocalStorageProvider':
            candidates = await asyncio.to_thread(
                self._expired_local_upload_candidates,
                context,
                retention_days,
                True,
            )
            return await asyncio.to_thread(
                self._delete_local_candidates,
                candidates,
            )

        if provider_name == 'S3StorageProvider':
            return await self._cleanup_expired_s3_uploaded_files(context, retention_days)

        return 0

    async def _expired_uploaded_candidates(
        self,
        context: TenantContext,
        retention_days: int,
    ) -> list[dict[str, Any]]:
        provider_name = self.ap.storage_mgr.storage_provider.__class__.__name__
        if provider_name == 'LocalStorageProvider':
            return await asyncio.to_thread(
                self._expired_local_upload_candidates,
                context,
                retention_days,
            )
        if provider_name == 'S3StorageProvider':
            return await self._expired_s3_upload_candidates(context, retention_days)
        return []

    async def _cleanup_expired_s3_uploaded_files(
        self,
        context: ExecutionContext,
        retention_days: int,
    ) -> int:
        provider = self.ap.storage_mgr.storage_provider
        candidates = await self._expired_s3_upload_candidates(context, retention_days)
        deleted = 0
        for item in candidates:
            await provider.delete(item['key'])
            deleted += 1
        return deleted

    async def _expired_s3_upload_candidates(
        self,
        context: TenantContext,
        retention_days: int,
    ) -> list[dict[str, Any]]:
        provider = self.ap.storage_mgr.storage_provider
        run_io = getattr(provider, '_run_io', None)
        if callable(run_io):
            return await run_io(
                self._expired_s3_upload_candidates_sync,
                context,
                retention_days,
            )
        return await asyncio.to_thread(
            self._expired_s3_upload_candidates_sync,
            context,
            retention_days,
        )

    def _expired_s3_upload_candidates_sync(
        self,
        context: TenantContext,
        retention_days: int,
    ) -> list[dict[str, Any]]:
        provider = self.ap.storage_mgr.storage_provider
        cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=retention_days)
        candidates = []
        max_candidates = self._max_files_per_run()
        paginator = provider.s3_client.get_paginator('list_objects_v2')

        seen_prefixes: set[str] = set()
        for owner_type in UPLOAD_OWNER_TYPES:
            prefix = self.ap.storage_mgr.scoped_prefix(context, owner_type=owner_type)
            if prefix in seen_prefixes:
                continue
            seen_prefixes.add(prefix)
            for page in paginator.paginate(Bucket=provider.bucket_name, Prefix=prefix):
                for obj in page.get('Contents', []):
                    key = obj.get('Key', '')
                    last_modified = obj.get('LastModified')
                    if not self._is_uploaded_file_key(context, key):
                        continue
                    if last_modified and last_modified < cutoff:
                        candidates.append(
                            {
                                'key': key,
                                'size_bytes': obj.get('Size', 0),
                                'modified_at': last_modified.isoformat(),
                            }
                        )
                        if len(candidates) >= max_candidates:
                            return candidates

        return candidates

    def _delete_local_candidates(self, candidates: list[dict[str, Any]]) -> int:
        deleted = 0
        for item in candidates:
            try:
                os.remove(item['path'])
                deleted += 1
            except FileNotFoundError:
                pass
            except Exception as e:
                self.ap.logger.warning(f'Failed to delete expired uploaded file {item["key"]}: {e}')
        return deleted

    def _cleanup_expired_log_files(self, retention_days: int) -> int:
        deleted = 0
        for item in self._expired_log_candidates(retention_days, include_paths=True):
            try:
                os.remove(item['path'])
                deleted += 1
            except FileNotFoundError:
                pass
            except Exception as e:
                self.ap.logger.warning(f'Failed to delete expired log file {item["name"]}: {e}')
        return deleted

    def _expired_local_upload_candidates(
        self,
        context: TenantContext,
        retention_days: int,
        include_paths: bool = False,
    ) -> list[dict[str, Any]]:
        storage_root = Path('data/storage')
        cutoff = datetime.datetime.now().timestamp() - retention_days * 86400
        candidates = []
        max_candidates = self._max_files_per_run()
        seen_roots: set[Path] = set()
        for owner_type in UPLOAD_OWNER_TYPES:
            scoped_root = storage_root / self.ap.storage_mgr.scoped_prefix(context, owner_type=owner_type)
            if scoped_root in seen_roots:
                continue
            seen_roots.add(scoped_root)
            if not scoped_root.exists():
                continue
            for entry in scoped_root.rglob('*'):
                if not entry.is_file():
                    continue
                stat = entry.stat()
                if stat.st_mtime >= cutoff:
                    continue
                item = {
                    'key': entry.relative_to(storage_root).as_posix(),
                    'size_bytes': stat.st_size,
                    'modified_at': datetime.datetime.fromtimestamp(
                        stat.st_mtime,
                        datetime.timezone.utc,
                    ).isoformat(),
                }
                if include_paths:
                    item['path'] = str(entry)
                candidates.append(item)
                if len(candidates) >= max_candidates:
                    return candidates
        return candidates

    def _expired_log_candidates(self, retention_days: int, include_paths: bool = False) -> list[dict[str, Any]]:
        log_root = Path('data/logs')
        if not log_root.exists():
            return []

        cutoff_date = datetime.date.today() - datetime.timedelta(days=retention_days - 1)
        candidates = []
        for entry in log_root.iterdir():
            if not entry.is_file():
                continue
            match = LOG_FILE_PATTERN.match(entry.name)
            if not match:
                continue
            try:
                file_date = datetime.date.fromisoformat(match.group(1))
            except ValueError:
                continue
            if file_date >= cutoff_date:
                continue
            stat = entry.stat()
            item = {
                'name': entry.name,
                'date': file_date.isoformat(),
                'size_bytes': stat.st_size,
            }
            if include_paths:
                item['path'] = str(entry)
            candidates.append(item)
        return candidates

    def _is_uploaded_file_key(self, context: TenantContext, key: str) -> bool:
        return any(
            key.startswith(self.ap.storage_mgr.scoped_prefix(context, owner_type=owner_type))
            and self.ap.storage_mgr.is_scoped_object_key(key, expected_owner_type=owner_type)
            for owner_type in UPLOAD_OWNER_TYPES
        )

    async def _monitoring_counts(self, context: TenantContext) -> dict[str, int]:
        workspace_uuid = require_workspace_uuid(context)
        tables = {
            'messages': (persistence_monitoring.MonitoringMessage, persistence_monitoring.MonitoringMessage.id),
            'llm_calls': (persistence_monitoring.MonitoringLLMCall, persistence_monitoring.MonitoringLLMCall.id),
            'tool_calls': (persistence_monitoring.MonitoringToolCall, persistence_monitoring.MonitoringToolCall.id),
            'embedding_calls': (
                persistence_monitoring.MonitoringEmbeddingCall,
                persistence_monitoring.MonitoringEmbeddingCall.id,
            ),
            'errors': (persistence_monitoring.MonitoringError, persistence_monitoring.MonitoringError.id),
            'sessions': (
                persistence_monitoring.MonitoringSession,
                persistence_monitoring.MonitoringSession.session_id,
            ),
            'feedback': (persistence_monitoring.MonitoringFeedback, persistence_monitoring.MonitoringFeedback.id),
        }
        counts: dict[str, int] = {}
        for key, (model, column) in tables.items():
            result = await self.ap.persistence_mgr.execute_async(
                sqlalchemy.select(sqlalchemy.func.count(column)).where(model.workspace_uuid == workspace_uuid)
            )
            counts[key] = result.scalar() or 0
        return counts

    async def _binary_storage_stats(self, context: TenantContext) -> dict[str, Any]:
        workspace_uuid = require_workspace_uuid(context)
        count_result = await self.ap.persistence_mgr.execute_async(
            sqlalchemy.select(sqlalchemy.func.count(persistence_bstorage.BinaryStorage.unique_key)).where(
                persistence_bstorage.BinaryStorage.workspace_uuid == workspace_uuid
            )
        )
        size_bytes = None
        try:
            size_result = await self.ap.persistence_mgr.execute_async(
                sqlalchemy.select(
                    sqlalchemy.func.sum(sqlalchemy.func.length(persistence_bstorage.BinaryStorage.value))
                ).where(persistence_bstorage.BinaryStorage.workspace_uuid == workspace_uuid)
            )
            size_bytes = size_result.scalar() or 0
        except Exception as e:
            self.ap.logger.warning(f'Failed to estimate binary storage size: {e}')

        return {
            'count': count_result.scalar() or 0,
            'size_bytes': size_bytes,
        }

    def _path_size(self, path: Path) -> int:
        if not path.exists():
            return 0
        if path.is_file():
            return path.stat().st_size
        total = 0
        for root, _, files in os.walk(path):
            for file_name in files:
                file_path = Path(root) / file_name
                try:
                    total += file_path.stat().st_size
                except FileNotFoundError:
                    pass
        return total

    def _file_count(self, path: Path) -> int:
        if not path.exists():
            return 0
        if path.is_file():
            return 1
        count = 0
        for _, _, files in os.walk(path):
            count += len(files)
        return count

    def _positive_int(self, value: Any, default: int, name: str) -> int:
        try:
            parsed = int(value)
        except (TypeError, ValueError):
            self.ap.logger.warning(f'Invalid {name}: {value!r}, using {default}')
            return default
        if parsed < 1:
            self.ap.logger.warning(f'Invalid {name}: {value!r}, using {default}')
            return default
        return parsed