"""Core digital employee orchestrator.

This module provides the base DigitalEmployee class that manages agents,
tools, and MCPs. It implements the ConfigBuilder interface to
provide configuration building capabilities.

Authors:
    Immanuel Rhesa (immanuel.rhesa@gdplabs.id)
    Vio Albert Ferdinand (vio.a.ferdinand@gdplabs.id)

References:
    NONE
"""

from pathlib import Path
from typing import Any

from glaip_sdk import MCP, Agent, Tool
from omegaconf import OmegaConf

from digital_employee_core.config_templates.loader import ConfigTemplateLoader
from digital_employee_core.configuration.agent_configuration import AgentConfigKeys, StepLimitConfig
from digital_employee_core.configuration.configuration import DigitalEmployeeConfiguration
from digital_employee_core.constants import DEFAULT_MODEL_NAME
from digital_employee_core.identity.identity import DigitalEmployeeIdentity


class DigitalEmployee:
    """Core digital employee orchestrator.

    This class manages the lifecycle of a digital employee, including
    initialization, deployment, and execution of agents with tools and MCPs.
    Provides methods for building tool and MCP configurations from YAML templates.

    Attributes:
        identity (DigitalEmployeeIdentity): The digital employee's identity.
        tools (list[Tool]): List of tools the digital employee can use.
        sub_agents (list[Agent]): List of sub-agents (for future use).
        mcps (list[MCP]): List of MCPs the digital employee can use.
        configurations (list[DigitalEmployeeConfiguration]): List of configuration objects for tools and MCPs.
        model (str): Model identifier to use for the agent.
        agent_config (dict[str, Any] | Agent._UNSET): Agent execution configuration (e.g., execution_timeout).
        agent (Agent | None): The main agent instance.
        config_loader (ConfigTemplateLoader): The configuration template loader instance.
    """

    def __init__(  # noqa: PLR0913
        self,
        identity: DigitalEmployeeIdentity,
        tools: list[Tool] | None = None,
        sub_agents: list[Agent] | None = None,
        mcps: list[MCP] | None = None,
        configurations: list[DigitalEmployeeConfiguration] | None = None,
        model: str = DEFAULT_MODEL_NAME,
        agent_config: dict[str, Any] | None = None,
    ):
        """Initialize the digital employee.

        Args:
            identity (DigitalEmployeeIdentity): The digital employee's identity.
            tools (list[Tool] | None, optional): List of tools the digital employee can use. Defaults to None.
            sub_agents (list[Agent] | None, optional): List of sub-agents (for future use). Defaults to None.
            mcps (list[MCP] | None, optional): List of MCPs the digital employee can use. Defaults to None.
            configurations (list[DigitalEmployeeConfiguration] | None, optional): List of configuration objects
                for tools and MCPs. Defaults to None.
            model (str, optional): Model identifier to use for the agent. Defaults to DEFAULT_MODEL_NAME.
            agent_config (dict[str, Any] | None, optional): Agent execution configuration (e.g., execution_timeout).
                Defaults to None. If None, will be converted to Agent._UNSET internally.
        """
        self.identity = identity
        self.tools = tools or []
        self.sub_agents = sub_agents or []
        self.mcps = mcps or []
        self.configurations = configurations or []
        self.model = model
        # TODO: make agent config strong typing
        self.agent_config = self._ensure_valid_agent_config(agent_config) if agent_config is not None else Agent._UNSET
        self.agent: Agent | None = None
        self.config_loader = ConfigTemplateLoader()

    def add_tools(self, tools: list[Tool]) -> None:
        """Add tools to the digital employee.

        Args:
            tools (list[Tool]): List of tools to add.
        """
        self.tools.extend(tools)

    def remove_tools(self, tools: list[Tool]) -> None:
        """Remove tools from the digital employee.

        Args:
            tools (list[Tool]): List of tools to remove.
        """
        tools_to_remove = set(tools)
        self.tools = [tool for tool in self.tools if tool not in tools_to_remove]

    def add_mcps(self, mcps: list[MCP]) -> None:
        """Add MCPs to the digital employee.

        Args:
            mcps (list[MCP]): List of MCPs to add.
        """
        self.mcps.extend(mcps)

    def remove_mcps(self, mcps: list[MCP]) -> None:
        """Remove MCPs from the digital employee.

        Args:
            mcps (list[MCP]): List of MCPs to remove.
        """
        mcps_to_remove = set(mcps)
        self.mcps = [mcp for mcp in self.mcps if mcp not in mcps_to_remove]

    def add_sub_agents(self, sub_agents: list[Agent]) -> None:
        """Add sub-agents to the digital employee.

        Args:
            sub_agents (list[Agent]): List of sub-agents to add.
        """
        self.sub_agents.extend(sub_agents)

    def remove_sub_agents(self, sub_agents: list[Agent]) -> None:
        """Remove sub-agents from the digital employee by name.

        Args:
            sub_agents (list[Agent]): List of sub-agents to remove.
        """
        agent_names_to_remove = {agent.name for agent in sub_agents if agent.name}
        self.sub_agents = [agent for agent in self.sub_agents if agent.name not in agent_names_to_remove]

    def build_prompt(self) -> str:
        """Build the prompt for the agent.

        This method can be overridden by subclasses to provide custom
        prompts based on the digital employee's identity and job.
        Always includes the digital employee's name, email, and language preferences
        in the prompt to ensure identity awareness.

        Returns:
            str: The prompt string for the agent.
        """
        # Build identity information with language
        language_names = [lang.value for lang in self.identity.languages]

        if len(language_names) > 1:
            language_list = ", ".join(language_names[:-1]) + f" and {language_names[-1]}"
            language_instruction = (
                f"You must respond only in the following languages: {language_list}. "
                f"Choose the most appropriate language based on the user's input or context."
            )
        else:
            language_instruction = f"You must respond only in {language_names[0]} language."

        identity_prefix = (
            f"You are {self.identity.name} ({self.identity.email}), {self.identity.job.title}. {language_instruction}"
        )

        return f"{identity_prefix}\n\n{self.identity.job.instruction}"

    def deploy(self) -> None:
        """Deploy the digital employee.

        This method initializes the agent with tools and MCPs,
        builds configurations, and deploys it to the AI platform.
        """
        # Build configurations if provided
        tool_configs = None
        mcp_configs = None
        if self.configurations:
            tool_configs = self.build_tool_config(self.configurations)
            mcp_configs = self.build_mcp_config(self.configurations)

        self.agent = Agent(
            name=self.identity.name,
            description=self.identity.job.description,
            instruction=self.build_prompt(),
            tools=self.tools,
            agents=self.sub_agents,
            mcps=self.mcps,
            tool_configs=tool_configs,
            mcp_configs=mcp_configs,
            model=self.model,
            agent_config=self.agent_config,
        )

        self.agent.deploy()

    def run(
        self,
        message: str,
        configurations: list[DigitalEmployeeConfiguration] | None = None,
        override_agent_config: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> str:
        """Run the digital employee with a message.

        Args:
            message (str): The message/prompt to send to the digital employee.
            configurations (list[DigitalEmployeeConfiguration] | None, optional): List of configuration objects.
                Defaults to None.
            override_agent_config (dict[str, Any] | None, optional): Dictionary of agent config to be overridden.
                Defaults to None.
            **kwargs (Any): Additional keyword arguments to pass to agent.run.

        Returns:
            str: The agent's response as a string.

        Raises:
            RuntimeError: If the agent has not been deployed.
        """
        if self.agent is None:
            raise RuntimeError("Agent has not been deployed. Call deploy() before run().")

        runtime_config = {}

        if configurations:
            runtime_config.update(self._build_runtime_configs(configurations))

        if override_agent_config:
            validated_agent_config = self._ensure_valid_agent_config(override_agent_config)
            if validated_agent_config:
                runtime_config.update({"agent_config": validated_agent_config})

        return self.agent.run(message=message, runtime_config=runtime_config or None, **kwargs)

    def build_tool_config(self, configurations: list[DigitalEmployeeConfiguration]) -> dict[str, dict[str, Any]]:
        """Build tool configuration based on identity and configurations.

        This method loads tool configuration templates from YAML and replaces
        placeholders with values from DigitalEmployeeConfiguration objects. Subclasses can
        override this to provide custom logic and merge with parent configs.

        Example for subclasses:
            def build_tool_config(self, configurations):
                # Get base configs from parent
                base_configs = super().build_tool_config(configurations)

                # Load additional configs for this subclass
                additional = self.config_loader.load_tool_configs(configurations)

                # Merge them
                return self.config_loader.merge_configs(base_configs, additional)

        Args:
            configurations (list[DigitalEmployeeConfiguration]): List of configuration objects.

        Returns:
            dict[str, dict[str, Any]]: Dictionary mapping tool names to their configuration dictionaries.
        """
        return self.config_loader.load_tool_configs(configurations, self.tools)

    def build_mcp_config(self, configurations: list[DigitalEmployeeConfiguration]) -> dict[str, dict[str, Any]]:
        """Build MCP configuration based on identity and configurations.

        This method loads MCP configuration templates from YAML and replaces
        placeholders with values from DigitalEmployeeConfiguration objects. Subclasses can
        override this to provide custom logic and merge with parent configs.

        Example for subclasses:
            def build_mcp_config(self, configurations):
                # Get base configs from parent
                base_configs = super().build_mcp_config(configurations)

                # Load additional configs for this subclass
                additional = self.config_loader.load_mcp_configs(configurations)

                # Merge them
                return self.config_loader.merge_configs(base_configs, additional)

        Args:
            configurations (list[DigitalEmployeeConfiguration]): List of configuration objects.

        Returns:
            dict[str, dict[str, Any]]: Dictionary mapping MCP names to their configuration dictionaries.
        """
        return self.config_loader.load_mcp_configs(configurations, self.mcps)

    def _ensure_valid_agent_config(self, agent_config: dict[str, Any]) -> dict[str, Any]:
        """Ensure agent config is valid.

        Args:
            agent_config (dict[str, Any]): Agent config to validate.

        Returns:
            dict[str, Any]: Validated agent config.
        """
        agent_config = agent_config.copy()
        if AgentConfigKeys.STEP_LIMIT_CONFIG in agent_config:
            step_limit_config = agent_config[AgentConfigKeys.STEP_LIMIT_CONFIG]
            if isinstance(step_limit_config, dict):
                step_limit_object = StepLimitConfig.model_validate(step_limit_config)
            elif isinstance(step_limit_config, StepLimitConfig):
                step_limit_object = step_limit_config
            else:
                raise TypeError(
                    f"step_limit_config must be a StepLimitConfig instance or dict, "
                    f"got {type(step_limit_config).__name__}"
                )
            step_limit_dict = step_limit_object.to_dict()
            if step_limit_dict:
                agent_config[AgentConfigKeys.STEP_LIMIT_CONFIG] = step_limit_dict
            else:
                agent_config.pop(AgentConfigKeys.STEP_LIMIT_CONFIG, None)

        return agent_config

    def _build_runtime_configs(
        self,
        configurations: list[DigitalEmployeeConfiguration],
    ) -> dict[str, dict[str, Any]]:
        """Build runtime configurations for tools and MCPs.

        This method builds all configurations and filters them to only include
        tools and MCPs that are actually used by the agent. Runtime configurations
        are merged with instance configurations (self.configurations), with runtime
        configurations taking precedence.

        Args:
            configurations (list[DigitalEmployeeConfiguration]): List of runtime configuration objects.

        Returns:
            dict[str, dict[str, Any]]: Runtime configuration dictionary.
        """
        # Merge instance configurations with runtime configurations
        # Runtime configurations override instance configurations
        merged_configs = self.configurations + configurations

        tool_configs = self.config_loader.load_tool_configs(merged_configs, self.agent.tools)
        mcp_configs = self.config_loader.load_mcp_configs(merged_configs, self.agent.mcps)

        return {
            "tool_configs": tool_configs,
            "mcp_configs": mcp_configs,
        }

    def _merge_defaults_with_parent(self, child_template_dir: Path) -> Any:
        """Merge parent defaults.yaml with child (subclass) defaults.yaml in memory.

        Args:
            child_template_dir (Path): Path to the subclass's config_templates directory.

        Returns:
            OmegaConf: Merged defaults (parent <- child, child overrides parent).
        """
        parent_defaults_path = self.config_loader.template_dir / "defaults.yaml"
        child_defaults_path = child_template_dir / "defaults.yaml"

        parent_defaults = (
            OmegaConf.load(parent_defaults_path) if parent_defaults_path.exists() else OmegaConf.create({})
        )
        child_defaults = OmegaConf.load(child_defaults_path) if child_defaults_path.exists() else OmegaConf.create({})

        return OmegaConf.merge(parent_defaults, child_defaults)
