X-Financial/server/src/app/services/agent_foundation.py

from __future__ import annotations

import threading

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.core.logging import get_logger
from app.db.base import Base
from app.db.session import get_session_factory
from app.models.agent_asset import AgentAsset
from app.services.agent_foundation_asset_helpers import AgentFoundationAssetHelperMixin
from app.services.agent_foundation_asset_seed import AgentFoundationAssetSeedMixin
from app.services.agent_foundation_asset_topup import AgentFoundationAssetTopUpMixin
from app.services.agent_foundation_constants import (
    ATTACHMENT_RULE_ASSET_CODE,
    ATTACHMENT_RULE_RUNTIME_CONFIG,
    COMPANY_COMMUNICATION_RULE_SCENARIO_JSON,
    COMPANY_COMMUNICATION_RULE_VERSION,
    COMPANY_TRAVEL_RULE_SCENARIO_JSON,
    COMPANY_TRAVEL_RULE_VERSION,
    DEMO_EXPENSE_CLAIM_SIGNATURES,
    DEMO_PAYABLE_SIGNATURES,
    DEMO_RECEIVABLE_SIGNATURES,
    LEGACY_RULE_CODES,
    PLATFORM_DESTINATION_LOCATION_RULE_CODE,
    PLATFORM_DESTINATION_LOCATION_RULE_FILENAME,
)
from app.services.agent_foundation_financial_seed import AgentFoundationFinancialSeedMixin
from app.services.agent_foundation_markdown import AgentFoundationMarkdownMixin
from app.services.agent_foundation_risk_rules import AgentFoundationRiskRuleMixin
from app.services.agent_foundation_spreadsheets import AgentFoundationSpreadsheetMixin

logger = get_logger("app.services.agent_foundation")
_foundation_ready_lock = threading.RLock()
_foundation_ready_keys: set[str] = set()


def prepare_agent_foundation() -> None:
    settings = get_settings()
    if not settings.setup_completed:
        logger.info("Agent foundation bootstrap skipped because setup is incomplete")
        return

    session_factory = get_session_factory()
    with session_factory() as db:
        AgentFoundationService(db).ensure_foundation_ready()


class AgentFoundationService(
    AgentFoundationAssetSeedMixin,
    AgentFoundationFinancialSeedMixin,
    AgentFoundationAssetTopUpMixin,
    AgentFoundationSpreadsheetMixin,
    AgentFoundationAssetHelperMixin,
    AgentFoundationMarkdownMixin,
    AgentFoundationRiskRuleMixin,
):
    def __init__(self, db: Session) -> None:
        self.db = db

    def ensure_foundation_ready(self) -> None:
        cache_key = self._foundation_cache_key()
        if cache_key in _foundation_ready_keys:
            return

        with _foundation_ready_lock:
            if cache_key in _foundation_ready_keys:
                return
            self._prepare_foundation()
            _foundation_ready_keys.add(cache_key)

    def _prepare_foundation(self) -> None:
        try:
            Base.metadata.create_all(bind=self.db.get_bind())
            self._ensure_agent_asset_schema()
            self._seed_agent_assets()
            self._sync_demo_financial_records()
            self._seed_runs_and_logs()
            self.db.commit()
        except Exception:
            self.db.rollback()
            logger.exception("Failed to prepare agent foundation")
            raise

    def _foundation_cache_key(self) -> str:
        bind = self.db.get_bind()
        return str(getattr(bind, "url", "") or id(bind))

    def _sync_demo_financial_records(self) -> None:
        if get_settings().seed_demo_financial_records:
            self._seed_financial_records()
            return
        self._purge_demo_financial_records()