"""Centralized runtime loading for Beaver agents."""

from __future__ import annotations

import asyncio
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from beaver.coordinator.registry import AgentRegistry
from beaver.engine.context import ContextBuilder
from beaver.engine.session import SessionManager
from beaver.foundation.config import BeaverConfig, load_config
from beaver.integrations.mcp import MCPConnectionManager
from beaver.memory.curated.store import MemoryStore
from beaver.memory.runs import RunMemoryStore
from beaver.memory.skills import SkillLearningStore
from beaver.services.memory_service import MemoryService
from beaver.skills.drafts import DraftService
from beaver.skills.learning import EvidenceSelector, SkillDraftSynthesizer, SkillLearningPipelineService, SkillLearningService
from beaver.skills.learning.safety import SkillDraftSafetyChecker
from beaver.skills.learning.eval import SkillDraftEvaluator
from beaver.skills.publisher import SkillPublisher
from beaver.skills.reviews import ReviewService
from beaver.skills.specs import SkillSpecStore
from beaver.tasks import TaskExecutionPlanner, TaskService
from beaver.tasks.skill_resolver import TaskSkillResolver
from beaver.skills import SkillAssembler, SkillsLoader
from beaver.tools import ObjectBackedTool, ToolAssembler, ToolExecutor, ToolRegistry
from beaver.tools.builtins import (
    ClarifyTool,
    CronTool,
    DelegateTool,
    EchoTool,
    ExecuteCodeTool,
    ListDirectoryTool,
    MemoryTool,
    PatchFileTool,
    ProcessTool,
    ReadFileTool,
    SearchFilesTool,
    SendMessageTool,
    SpawnTool,
    SessionSearchTool,
    SkillManageTool,
    SkillViewTool,
    SkillsListTool,
    TerminalTool,
    TodoTool,
    UserFilesCopyToWorkspaceTool,
    UserFilesListTool,
    UserFilesMkdirTool,
    UserFilesPublishOutputTool,
    UserFilesReadTool,
    UserFilesWriteTool,
    WebFetchTool,
    WebSearchTool,
    WriteFileTool,
)


@dataclass(slots=True)
class EngineLoadResult:
    """描述当前 agent runtime 已经装好的依赖。

    这里同时保留两类字段：
    1. `tools/skills/memory_stores/permissions`
       - 便于做状态展示、调试、轻量测试
    2. `session_manager/tool_registry/...`
       - 供真正的运行时主链直接使用
    """

    workspace: Path
    config: BeaverConfig = field(default_factory=BeaverConfig)
    tools: list[str] = field(default_factory=list)
    skills: list[str] = field(default_factory=list)
    memory_stores: list[str] = field(default_factory=list)
    permissions: list[str] = field(default_factory=list)
    session_manager: SessionManager | None = None
    curated_memory_store: MemoryStore | None = None
    memory_service: MemoryService | None = None
    run_memory_store: RunMemoryStore | None = None
    skill_learning_store: SkillLearningStore | None = None
    tool_registry: ToolRegistry | None = None
    tool_assembler: ToolAssembler | None = None
    tool_executor: ToolExecutor | None = None
    context_builder: ContextBuilder | None = None
    skills_loader: SkillsLoader | None = None
    skill_assembler: SkillAssembler | None = None
    skill_spec_store: SkillSpecStore | None = None
    draft_service: DraftService | None = None
    review_service: ReviewService | None = None
    skill_publisher: SkillPublisher | None = None
    skill_learning_service: SkillLearningService | None = None
    skill_learning_pipeline: SkillLearningPipelineService | None = None
    agent_registry: AgentRegistry | None = None
    task_skill_resolver: TaskSkillResolver | None = None
    task_service: TaskService | None = None
    task_execution_planner: TaskExecutionPlanner | None = None
    mcp_manager: MCPConnectionManager | None = None
    mcp_report: dict[str, dict] = field(default_factory=dict)
    closeables: list[tuple[str, Callable[[], None]]] = field(default_factory=list, repr=False)
    closed: bool = False

    def register_closeable(self, name: str, close_fn: Callable[[], None]) -> None:
        """登记一个由 runtime 统一关闭的资源。"""

        self.closeables.append((name, close_fn))

    def close(self) -> None:
        """按后进先出顺序关闭 runtime 资源。

        这一步先保持同步、最小、可组合：
        1. 只管理已经明确需要关闭的资源
        2. 暂不引入 async shutdown 协议
        3. 为后续 Web/Gateway lifespan 留统一入口
        """

        if self.closed:
            return

        errors: list[tuple[str, BaseException]] = []
        for name, close_fn in reversed(self.closeables):
            try:
                close_fn()
            except BaseException as exc:  # pragma: no cover - defensive cleanup path
                errors.append((name, exc))
        self.closed = True

        if errors:
            parts = ", ".join(f"{name}: {exc}" for name, exc in errors)
            raise RuntimeError(f"Runtime shutdown failed for {parts}")


class EngineLoader:
    """为任意 Beaver agent 装载共享 runtime 能力。

    当前先做“最小可运行主链”需要的装配：
    - session manager
    - curated memory store
    - context builder
    - built-in tools
    - tool executor

    等主链跑稳后，再把 skills、权限、MCP、delegation 逐步加进来。
    """

    def __init__(
        self,
        *,
        workspace: str | Path | None = None,
        config_path: str | Path | None = None,
        config: BeaverConfig | None = None,
        session_manager: SessionManager | None = None,
        curated_memory_store: MemoryStore | None = None,
        memory_service: MemoryService | None = None,
        run_memory_store: RunMemoryStore | None = None,
        skill_learning_store: SkillLearningStore | None = None,
        tool_registry: ToolRegistry | None = None,
        tool_assembler: ToolAssembler | None = None,
        context_builder: ContextBuilder | None = None,
        skills_loader: SkillsLoader | None = None,
        skill_assembler: SkillAssembler | None = None,
        skill_spec_store: SkillSpecStore | None = None,
        draft_service: DraftService | None = None,
        review_service: ReviewService | None = None,
        skill_publisher: SkillPublisher | None = None,
        skill_learning_service: SkillLearningService | None = None,
        skill_learning_pipeline: SkillLearningPipelineService | None = None,
        agent_registry: AgentRegistry | None = None,
        task_skill_resolver: TaskSkillResolver | None = None,
        task_service: TaskService | None = None,
        task_execution_planner: TaskExecutionPlanner | None = None,
    ) -> None:
        self.config = config or load_config(workspace=workspace, config_path=config_path)
        configured_workspace = self.config.agents_defaults.workspace
        env_workspace = os.getenv("BEAVER_WORKSPACE")
        self.workspace = Path(workspace or configured_workspace or env_workspace or Path.cwd())
        self._session_manager = session_manager
        self._curated_memory_store = curated_memory_store
        self._memory_service = memory_service
        self._run_memory_store = run_memory_store
        self._skill_learning_store = skill_learning_store
        self._tool_registry = tool_registry
        self._tool_assembler = tool_assembler
        self._context_builder = context_builder
        self._skills_loader = skills_loader
        self._skill_assembler = skill_assembler
        self._skill_spec_store = skill_spec_store
        self._draft_service = draft_service
        self._review_service = review_service
        self._skill_publisher = skill_publisher
        self._skill_learning_service = skill_learning_service
        self._skill_learning_pipeline = skill_learning_pipeline
        self._agent_registry = agent_registry
        self._task_skill_resolver = task_skill_resolver
        self._task_service = task_service
        self._task_execution_planner = task_execution_planner

    def load(self) -> EngineLoadResult:
        """装配当前主链需要的最小 runtime 对象。"""

        workspace = self.workspace
        session_manager = self._session_manager or SessionManager(workspace)

        curated_root = workspace / "memory" / "curated"
        curated_memory_store = self._curated_memory_store or MemoryStore(curated_root)
        memory_service = self._memory_service or MemoryService(curated_root, store=curated_memory_store)
        memory_service.initialize()
        run_memory_store = self._run_memory_store or RunMemoryStore(workspace / "memory" / "runs")
        skill_learning_store = self._skill_learning_store or SkillLearningStore(workspace / "memory" / "skills")

        tool_registry = self._tool_registry or ToolRegistry()
        skill_spec_store = self._skill_spec_store or SkillSpecStore(workspace)
        skills_loader = self._skills_loader or SkillsLoader(workspace, skill_store=skill_spec_store)
        if self._tool_registry is None:
            # 这里先注册最小工具集，满足主链的 tool loop。
            tool_registry.register_many(
                [
                    ObjectBackedTool(EchoTool()),
                    ObjectBackedTool(MemoryTool(store=memory_service.get_store())),
                    ObjectBackedTool(SessionSearchTool(db=session_manager)),
                    ObjectBackedTool(ListDirectoryTool()),
                    ObjectBackedTool(ReadFileTool()),
                    ObjectBackedTool(SearchFilesTool()),
                    ObjectBackedTool(WriteFileTool()),
                    ObjectBackedTool(PatchFileTool()),
                    ObjectBackedTool(UserFilesListTool()),
                    ObjectBackedTool(UserFilesReadTool()),
                    ObjectBackedTool(UserFilesWriteTool()),
                    ObjectBackedTool(UserFilesMkdirTool()),
                    ObjectBackedTool(UserFilesCopyToWorkspaceTool()),
                    ObjectBackedTool(UserFilesPublishOutputTool()),
                    ObjectBackedTool(WebFetchTool()),
                    ObjectBackedTool(WebSearchTool()),
                    ObjectBackedTool(TerminalTool()),
                    ObjectBackedTool(ProcessTool()),
                    ObjectBackedTool(ExecuteCodeTool()),
                    ObjectBackedTool(TodoTool()),
                    ObjectBackedTool(ClarifyTool()),
                    ObjectBackedTool(SendMessageTool()),
                    ObjectBackedTool(DelegateTool()),
                    ObjectBackedTool(SpawnTool()),
                    SkillsListTool(),
                    ObjectBackedTool(SkillViewTool(loader=skills_loader)),
                    SkillManageTool(),
                    CronTool(),
                ]
            )

        context_builder = self._context_builder or ContextBuilder()
        tool_assembler = self._tool_assembler or ToolAssembler()
        tool_executor = ToolExecutor(tool_registry)
        skill_assembler = self._skill_assembler or SkillAssembler(skills_loader)
        draft_service = self._draft_service or DraftService(skill_spec_store)
        review_service = self._review_service or ReviewService(skill_spec_store)
        skill_publisher = self._skill_publisher or SkillPublisher(skill_spec_store)
        evidence_selector = EvidenceSelector(run_memory_store, session_manager=session_manager)
        skill_learning_service = self._skill_learning_service or SkillLearningService(
            run_store=run_memory_store,
            learning_store=skill_learning_store,
            draft_service=draft_service,
            evidence_selector=evidence_selector,
            synthesizer=SkillDraftSynthesizer(),
        )
        skill_learning_pipeline = self._skill_learning_pipeline or SkillLearningPipelineService(
            learning_store=skill_learning_store,
            learning_service=skill_learning_service,
            draft_service=draft_service,
            review_service=review_service,
            publisher=skill_publisher,
            safety_checker=SkillDraftSafetyChecker(
                allowed_tool_names={spec.name for spec in tool_registry.list_specs()},
                allowed_tool_prefixes={
                    f"mcp_{server_id}_"
                    for server_id in self.config.tools.mcp_servers
                    if str(server_id).strip()
                },
            ),
            evaluator=SkillDraftEvaluator(run_memory_store),
        )
        agent_registry = self._agent_registry or AgentRegistry(workspace)
        task_skill_resolver = self._task_skill_resolver or TaskSkillResolver(
            skills_loader=skills_loader,
            draft_service=draft_service,
        )
        task_service = self._task_service or TaskService(workspace / "tasks")
        task_execution_planner = self._task_execution_planner or TaskExecutionPlanner(task_skill_resolver=task_skill_resolver)
        mcp_manager = MCPConnectionManager(
            self.config.tools.mcp_servers,
            authz_config=self.config.authz,
            backend_identity=self.config.backend_identity,
        )

        result = EngineLoadResult(
            workspace=workspace,
            config=self.config,
            tools=[spec.name for spec in tool_registry.list_specs()],
            skills=[record.name for record in skills_loader.list_skills(filter_unavailable=False)],
            memory_stores=["curated"],
            permissions=[],
            session_manager=session_manager,
            curated_memory_store=memory_service.get_store(),
            memory_service=memory_service,
            run_memory_store=run_memory_store,
            skill_learning_store=skill_learning_store,
            tool_registry=tool_registry,
            tool_assembler=tool_assembler,
            tool_executor=tool_executor,
            context_builder=context_builder,
            skills_loader=skills_loader,
            skill_assembler=skill_assembler,
            skill_spec_store=skill_spec_store,
            draft_service=draft_service,
            review_service=review_service,
            skill_publisher=skill_publisher,
            skill_learning_service=skill_learning_service,
            skill_learning_pipeline=skill_learning_pipeline,
            agent_registry=agent_registry,
            task_skill_resolver=task_skill_resolver,
            task_service=task_service,
            task_execution_planner=task_execution_planner,
            mcp_manager=mcp_manager,
        )
        if self._session_manager is None:
            result.register_closeable("session_manager", session_manager.close)
        result.register_closeable("mcp_manager", lambda: _close_mcp_manager(mcp_manager))
        return result


def _close_mcp_manager(manager: MCPConnectionManager) -> None:
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        asyncio.run(manager.close())
        return
    loop.create_task(manager.close())
