ecm_sync_system/sync_state_machine/pipeline/multi_project_pipeline.py

from __future__ import annotations

import logging
import copy
from pathlib import Path
from typing import Dict, Optional, List, TypedDict, Any

from ..config import (
    MultiProjectItemConfig,
    MultiProjectRunConfig,
    PipelineRunConfig,
    apply_config_overrides,
    build_config_from_file,
)
from ..logging import get_logger
from .factory import create_pipeline_from_config


logger = get_logger(__name__)


class _SharedStateSnapshot(TypedDict):
    local_nodes: List[Any]
    remote_nodes: List[Any]
    bindings_by_type: Dict[str, Dict[str, Optional[str]]]


def _resolve_embedded_config_path(raw_path: str, *, project_root: Path, profile_path: Path) -> Path:
    candidate = Path(raw_path)
    if candidate.is_absolute():
        return candidate

    profile_relative = (profile_path.parent / candidate).resolve()
    if profile_relative.exists():
        return profile_relative
    return (project_root / candidate).resolve()


def _project_scope_overrides(item: MultiProjectItemConfig, *, node_types: list[str]) -> Dict[str, object]:
    del node_types
    local_handler_configs = {
        "project": {"data_id_filter": [item.local_project_id]},
    }
    remote_handler_configs = {
        "project": {"data_id_filter": [item.remote_project_id]},
    }
    return {
        "local_datasource": {
            "handler_configs": local_handler_configs,
        },
        "remote_datasource": {
            "handler_configs": remote_handler_configs,
        },
    }


def _resolve_project_bootstrap_node_types(
    *,
    config: MultiProjectRunConfig,
    shared_node_types: list[str],
) -> list[str]:
    configured = [node_type for node_type in config.project_bootstrap_node_types if node_type]
    if configured:
        return list(dict.fromkeys(configured))
    return list(dict.fromkeys(shared_node_types))


class MultiProjectPipeline:
    def __init__(self, *, project_root: Path, config_path: Path, config: MultiProjectRunConfig):
        self.project_root = project_root
        self.config_path = config_path
        self.config = config

    def _load_pipeline_config(self, raw_path: str) -> PipelineRunConfig:
        resolved = _resolve_embedded_config_path(raw_path, project_root=self.project_root, profile_path=self.config_path)
        return build_config_from_file(project_root=self.project_root, file_path=resolved)

    def _prepare_project_config(
        self,
        *,
        shared_persist: Dict[str, object],
        item: MultiProjectItemConfig,
        base_path: Optional[str],
    ) -> PipelineRunConfig:
        raw_path = base_path or self.config.default_project_config
        cfg = self._load_pipeline_config(raw_path)
        cfg = apply_config_overrides(cfg, {"persist": {**shared_persist, "wipe_on_start": False}})
        return apply_config_overrides(cfg, _project_scope_overrides(item, node_types=list(cfg.node_types)))

    async def _snapshot_shared_state(self, *, pipeline, node_types: List[str]) -> _SharedStateSnapshot:
        local_nodes = [
            copy.deepcopy(node)
            for node_type in node_types
            for node in pipeline.local_collection.filter(node_type=node_type)
        ]
        remote_nodes = [
            copy.deepcopy(node)
            for node_type in node_types
            for node in pipeline.remote_collection.filter(node_type=node_type)
        ]
        bindings_by_type: Dict[str, Dict[str, Optional[str]]] = {}
        for node_type in node_types:
            bindings = await pipeline.binding_manager.get_all_bindings(node_type)
            bindings_by_type[node_type] = {local_id: remote_id for local_id, remote_id in bindings}
        return {
            "local_nodes": local_nodes,
            "remote_nodes": remote_nodes,
            "bindings_by_type": bindings_by_type,
        }

    async def run(self) -> Dict[str, Dict[str, object]]:
        results: Dict[str, Dict[str, object]] = {}

        global_config = self._load_pipeline_config(self.config.global_config)
        shared_persist = global_config.persist.model_dump(mode="json")
        shared_node_types = list(global_config.node_types)
        project_bootstrap_node_types = _resolve_project_bootstrap_node_types(
            config=self.config,
            shared_node_types=shared_node_types,
        )
        logger.info("🚀 Multi-project pipeline start: projects=%d", len(self.config.projects))
        global_pipeline = await create_pipeline_from_config(
            global_config,
            scope="global",
            pipeline_name="multi-project pipeline [global]",
        )
        try:
            results["global"] = await global_pipeline.run()
            shared_state = await self._snapshot_shared_state(
                pipeline=global_pipeline,
                node_types=project_bootstrap_node_types,
            )
        finally:
            await global_pipeline.close()

        for item in self.config.projects:
            project_config = self._prepare_project_config(
                shared_persist=shared_persist,
                item=item,
                base_path=item.config_path,
            )
            scope = item.scope
            title = f"multi-project pipeline [{scope}]"
            logger.info(
                "🚀 Project pipeline start: scope=%s local_project_id=%s remote_project_id=%s",
                scope,
                item.local_project_id,
                item.remote_project_id,
            )
            project_pipeline = await create_pipeline_from_config(
                project_config,
                scope=scope,
                pipeline_name=title,
                bootstrap_binding_node_types=project_bootstrap_node_types,
                ephemeral_node_types=project_bootstrap_node_types,
            )
            try:
                project_pipeline.preload_shared_state(
                    local_nodes=shared_state["local_nodes"],
                    remote_nodes=shared_state["remote_nodes"],
                    bindings_by_type=shared_state["bindings_by_type"],
                )
                results[scope] = await project_pipeline.run()
            finally:
                await project_pipeline.close()

        return results