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

from __future__ import annotations

import threading

from sqlalchemy import inspect, select, text
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._ensure_financial_record_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()
        url = str(getattr(bind, "url", "") or "")
        if url.endswith("/:memory:"):
            return f"{url}:{id(bind)}"
        return url or str(id(bind))

    def _ensure_financial_record_schema(self) -> None:
        bind = self.db.get_bind()
        inspector = inspect(bind)
        if "expense_claims" not in inspector.get_table_names():
            return

        column_names = {column["name"] for column in inspector.get_columns("expense_claims")}
        dialect_name = bind.dialect.name
        timestamp_type = "TIMESTAMP WITH TIME ZONE" if dialect_name == "postgresql" else "DATETIME"
        boolean_default = "FALSE" if dialect_name == "postgresql" else "0"

        if "hermes_scanned_at" not in column_names:
            self.db.execute(
                text(f"ALTER TABLE expense_claims ADD COLUMN hermes_scanned_at {timestamp_type}")
            )
        if "hermes_risk_flag" not in column_names:
            self.db.execute(
                text(
                    "ALTER TABLE expense_claims "
                    f"ADD COLUMN hermes_risk_flag BOOLEAN DEFAULT {boolean_default} NOT NULL"
                )
            )
            self.db.execute(
                text(
                    "CREATE INDEX IF NOT EXISTS ix_expense_claims_hermes_risk_flag "
                    "ON expense_claims (hermes_risk_flag)"
                )
            )
        self.db.flush()

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