"""Configuration builder for Digital Employee Agent.

This module provides the ConfigBuilder class that handles
building tool and MCP configurations from templates.

Authors:
    Immanuel Rhesa (immanuel.rhesa@gdplabs.id)

References:
    NONE
"""

from typing import Any, Callable

from glaip_sdk import MCP, Tool
from omegaconf import OmegaConf

from digital_employee_core.config_templates.loader import DEFAULTS_TEMPLATE, ConfigTemplateLoader
from digital_employee_core.digital_employee.config import DigitalEmployeeConfig
from digital_employee_core.digital_employee.dependencies import (
    DependencyBundle,
    MCPConfigBuilder,
    ToolConfigBuilder,
)
from digital_employee_core.escalation import EscalationManager


class ConfigBuilder:
    """Builds tool and MCP configurations from templates.

    This class handles configuration loading, merging from multiple loaders,
    and building configuration dictionaries for tools and MCPs.

    Attributes:
        config_loaders (list[ConfigTemplateLoader]): List of configuration loaders.
        configs (list[DigitalEmployeeConfig]): List of configuration objects.
        tools (list[Tool]): List of tools to build configs for.
        mcps (list[MCP]): List of MCPs to build configs for.
        escalation_manager (EscalationManager | None): Optional escalation manager for escalation dependencies.
    """

    def __init__(
        self,
        config_loaders: list[ConfigTemplateLoader],
        configs: list[DigitalEmployeeConfig],
        tools: list[Tool],
        mcps: list[MCP],
        escalation_manager: EscalationManager | None = None,
    ) -> None:
        """Initialize the ConfigBuilder.

        Args:
            config_loaders (list[ConfigTemplateLoader]): List of configuration loaders.
            configs (list[DigitalEmployeeConfig]): List of configuration objects.
            tools (list[Tool]): List of tools to build configs for.
            mcps (list[MCP]): List of MCPs to build configs for.
            escalation_manager (EscalationManager | None): Optional escalation manager.
        """
        self.config_loaders = config_loaders
        self.configs = configs
        self.tools = tools
        self.mcps = mcps
        self.escalation_manager = escalation_manager

    def build_tool_config(
        self, configs: list[DigitalEmployeeConfig], tools: list[Tool] | None = None
    ) -> dict[str, dict[str, Any]]:
        """Build tool configuration based on configurations.

        This method loads tool configuration templates from YAML and replaces
        placeholders with values from DigitalEmployeeConfig objects.

        Args:
            configs (list[DigitalEmployeeConfig]): List of configuration objects.
            tools (list[Tool] | None): Optional list of tools to build configs for.
                If None, uses self.tools.

        Returns:
            dict[str, dict[str, Any]]: Dictionary mapping tool names to their configuration dictionaries.
        """
        if tools is None:
            tools = self.tools
        return self._build_merged_configs(
            configs, lambda loader, configs, items: loader.load_tool_configs(configs, items), tools
        )

    def build_mcp_config(
        self, configs: list[DigitalEmployeeConfig], mcps: list[MCP] | None = None
    ) -> dict[str, dict[str, Any]]:
        """Build MCP configuration based on configurations.

        This method loads MCP configuration templates from YAML and replaces
        placeholders with values from DigitalEmployeeConfig objects.

        Args:
            configs (list[DigitalEmployeeConfig]): List of configuration objects.
            mcps (list[MCP] | None): Optional list of MCPs to build configs for.
                If None, uses self.mcps.

        Returns:
            dict[str, dict[str, Any]]: Dictionary mapping MCP names to their configuration dictionaries.
        """
        if mcps is None:
            mcps = self.mcps
        return self._build_merged_configs(
            configs, lambda loader, configs, items: loader.load_mcp_configs(configs, items), mcps
        )

    def prepare_tools_mcps_and_configs(
        self,
    ) -> tuple[list[Tool], list[MCP], dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
        """Prepare the final Tools/MCPs lists and their configurations.

        This combines base dependencies with escalation dependencies and merges their configurations.

        Returns:
            tuple[list[Tool], list[MCP], dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
                Tools, MCPs, tool config map, and MCP config map to be used for agent deployment.
        """
        base = self._build_base_dependency_bundle()
        escalation = self._build_escalation_dependency_bundle()
        merged = base.merged_with(
            escalation,
            merge_configs=bool(self.configs),
        )
        return merged.tools, merged.mcps, merged.tool_configs, merged.mcp_configs

    def _build_merged_configs(
        self,
        configs: list[DigitalEmployeeConfig],
        load_fn: Callable[
            [ConfigTemplateLoader, list[DigitalEmployeeConfig], list[Any] | None], dict[str, dict[str, Any]]
        ],
        filter_items: list[Any] | None = None,
    ) -> dict[str, dict[str, Any]]:
        """Build merged configurations from all config loaders.

        This is a helper method that reduces code duplication between
        build_tool_config() and build_mcp_config().

        Args:
            configs (list[DigitalEmployeeConfig]): List of configuration objects.
            load_fn (Callable): Function that takes (loader, configurations, filter_items) and returns configs.
            filter_items (list[Any] | None, optional): Items to filter (tools or MCPs). Defaults to None.

        Returns:
            dict[str, dict[str, Any]]: Dictionary mapping names to their configuration dictionaries.
        """
        if not self.config_loaders:
            return {}

        # First, merge all defaults from all loaders
        merged_defaults = OmegaConf.create({})
        for loader in self.config_loaders:
            defaults_path = loader.template_dir / DEFAULTS_TEMPLATE
            if defaults_path.exists():
                loader_defaults = OmegaConf.load(defaults_path)
                merged_defaults = OmegaConf.merge(merged_defaults, loader_defaults)

        # Convert merged defaults to DigitalEmployeeConfig list
        defaults_dict = OmegaConf.to_container(merged_defaults, resolve=False)
        defaults_config_list = [DigitalEmployeeConfig(key=str(k), value=v) for k, v in defaults_dict.items()]

        # Merge: defaults + user configurations (user takes precedence)
        all_configs = defaults_config_list + configs

        # Start with configs from the first loader (base) - using merged configs
        result = load_fn(self.config_loaders[0], all_configs, filter_items)

        # Merge configs from additional loaders (if any) - using merged configs
        for loader in self.config_loaders[1:]:
            additional_configs = load_fn(loader, all_configs, filter_items)
            result = ConfigTemplateLoader.merge_configs(result, additional_configs)

        return result

    def _build_tools_and_mcps_configs(
        self,
        tools: list[Tool],
        mcps: list[MCP],
        build_tool_configs: ToolConfigBuilder,
        build_mcp_configs: MCPConfigBuilder,
    ) -> tuple[dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
        """Build tool and MCP configuration maps for the given dependencies.

        This helper builds configuration dictionaries only when instance
        configurations are present (`self.configurations`).

        Args:
            tools (list[Tool]): Tools to build configurations for.
            mcps (list[MCP]): MCPs to build configurations for.
            build_tool_configs (ToolConfigBuilder): Callable that builds tool configs.
            build_mcp_configs (MCPConfigBuilder): Callable that builds MCP configs.

        Returns:
            tuple[dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
                Tool config map and MCP config map.
        """
        tool_configs = None
        mcp_configs = None
        if self.configs:
            tool_configs = build_tool_configs(self.configs, tools)
            mcp_configs = build_mcp_configs(self.configs, mcps)

        return tool_configs, mcp_configs

    def _build_base_tools_mcps_and_configs(
        self,
    ) -> tuple[list[Tool], list[MCP], dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
        """Build base Tools/MCPs and their configurations.

        Returns:
            tuple[list[Tool], list[MCP], dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
                Base tools, base MCPs, tool config map, and MCP config map.
        """
        tool_configs, mcp_configs = self._build_tools_and_mcps_configs(
            tools=self.tools,
            mcps=self.mcps,
            build_tool_configs=lambda configurations, _tools: self.build_tool_config(configurations),
            build_mcp_configs=lambda configurations, _mcps: self.build_mcp_config(configurations),
        )
        return self.tools, self.mcps, tool_configs, mcp_configs

    def _build_escalation_tools_mcps_and_configs(
        self,
    ) -> tuple[list[Tool], list[MCP], dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
        """Build escalation Tools/MCPs and their configurations.

        Returns:
            tuple[list[Tool], list[MCP], dict[str, dict[str, Any]] | None, dict[str, dict[str, Any]] | None]:
                Escalation tools, escalation MCPs, tool config map, and MCP config map.
        """
        if not self.escalation_manager:
            return [], [], None, None

        tools, mcps = self.escalation_manager.get_dependencies()
        loader = self.config_loaders[0] if self.config_loaders else ConfigTemplateLoader()
        tool_configs, mcp_configs = self._build_tools_and_mcps_configs(
            tools=tools,
            mcps=mcps,
            build_tool_configs=loader.load_tool_configs,
            build_mcp_configs=loader.load_mcp_configs,
        )
        return tools, mcps, tool_configs, mcp_configs

    def _build_base_dependency_bundle(self) -> DependencyBundle:
        """Build the base dependency bundle containing tools, MCPs, and their configurations.

        Returns:
            DependencyBundle: Bundle containing base tools, MCPs, and their config maps.
        """
        tools, mcps, tool_configs, mcp_configs = self._build_base_tools_mcps_and_configs()
        return DependencyBundle(tools=tools, mcps=mcps, tool_configs=tool_configs, mcp_configs=mcp_configs)

    def _build_escalation_dependency_bundle(self) -> DependencyBundle:
        """Build the escalation dependency bundle containing tools, MCPs, and their configurations.

        Returns:
            DependencyBundle: Bundle containing escalation tools, MCPs, and their config maps.
        """
        tools, mcps, tool_configs, mcp_configs = self._build_escalation_tools_mcps_and_configs()
        return DependencyBundle(tools=tools, mcps=mcps, tool_configs=tool_configs, mcp_configs=mcp_configs)
