import json
from collections.abc import Generator, Mapping, Sequence
from typing import Any, Optional, cast

from core.agent.entities import AgentToolEntity
from core.agent.plugin_entities import AgentStrategyParameter
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_manager import ModelInstance, ModelManager
from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
from core.plugin.impl.exc import PluginDaemonClientSideError
from core.plugin.impl.plugin import PluginInstaller
from core.provider_manager import ProviderManager
from core.tools.entities.tool_entities import ToolParameter, ToolProviderType
from core.tools.tool_manager import ToolManager
from core.variables.segments import StringSegment
from core.workflow.entities.node_entities import NodeRunResult
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.enums import SystemVariableKey
from core.workflow.nodes.agent.entities import AgentNodeData, AgentOldVersionModelFeatures, ParamsAutoGenerated
from core.workflow.nodes.base.entities import BaseNodeData
from core.workflow.nodes.enums import NodeType
from core.workflow.nodes.event.event import RunCompletedEvent
from core.workflow.nodes.tool.tool_node import ToolNode
from core.workflow.utils.variable_template_parser import VariableTemplateParser
from extensions.ext_database import db
from factories.agent_factory import get_plugin_agent_strategy
from models.model import Conversation
from models.workflow import WorkflowNodeExecutionStatus


class AgentNode(ToolNode):
    """
    Agent Node
    """

    _node_data_cls = AgentNodeData  # type: ignore
    _node_type = NodeType.AGENT

    def _run(self) -> Generator:
        """
        Run the agent node
        """
        node_data = cast(AgentNodeData, self.node_data)

        try:
            strategy = get_plugin_agent_strategy(
                tenant_id=self.tenant_id,
                agent_strategy_provider_name=node_data.agent_strategy_provider_name,
                agent_strategy_name=node_data.agent_strategy_name,
            )
        except Exception as e:
            yield RunCompletedEvent(
                run_result=NodeRunResult(
                    status=WorkflowNodeExecutionStatus.FAILED,
                    inputs={},
                    error=f"Failed to get agent strategy: {str(e)}",
                )
            )
            return

        agent_parameters = strategy.get_parameters()

        # get parameters
        parameters = self._generate_agent_parameters(
            agent_parameters=agent_parameters,
            variable_pool=self.graph_runtime_state.variable_pool,
            node_data=node_data,
        )
        parameters_for_log = self._generate_agent_parameters(
            agent_parameters=agent_parameters,
            variable_pool=self.graph_runtime_state.variable_pool,
            node_data=node_data,
            for_log=True,
        )

        # get conversation id
        conversation_id = self.graph_runtime_state.variable_pool.get(["sys", SystemVariableKey.CONVERSATION_ID])

        try:
            message_stream = strategy.invoke(
                params=parameters,
                user_id=self.user_id,
                app_id=self.app_id,
                conversation_id=conversation_id.text if conversation_id else None,
            )
        except Exception as e:
            yield RunCompletedEvent(
                run_result=NodeRunResult(
                    status=WorkflowNodeExecutionStatus.FAILED,
                    inputs=parameters_for_log,
                    error=f"Failed to invoke agent: {str(e)}",
                )
            )
            return

        try:
            # convert tool messages

            yield from self._transform_message(
                message_stream,
                {
                    "icon": self.agent_strategy_icon,
                    "agent_strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
                },
                parameters_for_log,
            )
        except PluginDaemonClientSideError as e:
            yield RunCompletedEvent(
                run_result=NodeRunResult(
                    status=WorkflowNodeExecutionStatus.FAILED,
                    inputs=parameters_for_log,
                    error=f"Failed to transform agent message: {str(e)}",
                )
            )

    def _generate_agent_parameters(
        self,
        *,
        agent_parameters: Sequence[AgentStrategyParameter],
        variable_pool: VariablePool,
        node_data: AgentNodeData,
        for_log: bool = False,
    ) -> dict[str, Any]:
        """
        Generate parameters based on the given tool parameters, variable pool, and node data.

        Args:
            agent_parameters (Sequence[AgentParameter]): The list of agent parameters.
            variable_pool (VariablePool): The variable pool containing the variables.
            node_data (AgentNodeData): The data associated with the agent node.

        Returns:
            Mapping[str, Any]: A dictionary containing the generated parameters.

        """
        agent_parameters_dictionary = {parameter.name: parameter for parameter in agent_parameters}

        result: dict[str, Any] = {}
        for parameter_name in node_data.agent_parameters:
            parameter = agent_parameters_dictionary.get(parameter_name)
            if not parameter:
                result[parameter_name] = None
                continue
            agent_input = node_data.agent_parameters[parameter_name]
            if agent_input.type == "variable":
                variable = variable_pool.get(agent_input.value)  # type: ignore
                if variable is None:
                    raise ValueError(f"Variable {agent_input.value} does not exist")
                parameter_value = variable.value
            elif agent_input.type in {"mixed", "constant"}:
                # variable_pool.convert_template expects a string template,
                # but if passing a dict, convert to JSON string first before rendering
                try:
                    parameter_value = json.dumps(agent_input.value, ensure_ascii=False)
                except TypeError:
                    parameter_value = str(agent_input.value)
                segment_group = variable_pool.convert_template(parameter_value)
                parameter_value = segment_group.log if for_log else segment_group.text
                # variable_pool.convert_template returns a string,
                # so we need to convert it back to a dictionary
                try:
                    parameter_value = json.loads(parameter_value)
                except json.JSONDecodeError:
                    parameter_value = parameter_value
            else:
                raise ValueError(f"Unknown agent input type '{agent_input.type}'")
            value = parameter_value
            if parameter.type == "array[tools]":
                value = cast(list[dict[str, Any]], value)
                value = [tool for tool in value if tool.get("enabled", False)]

                for tool in value:
                    if "schemas" in tool:
                        tool.pop("schemas")
                    parameters = tool.get("parameters", {})
                    if all(isinstance(v, dict) for _, v in parameters.items()):
                        params = {}
                        for key, param in parameters.items():
                            if param.get("auto", ParamsAutoGenerated.OPEN.value) == ParamsAutoGenerated.CLOSE.value:
                                value_param = param.get("value", {})
                                params[key] = value_param.get("value", "") if value_param is not None else None
                            else:
                                params[key] = None
                        parameters = params
                    tool["settings"] = {k: v.get("value", None) for k, v in tool.get("settings", {}).items()}
                    tool["parameters"] = parameters

            if not for_log:
                if parameter.type == "array[tools]":
                    value = cast(list[dict[str, Any]], value)
                    tool_value = []
                    for tool in value:
                        provider_type = ToolProviderType(tool.get("type", ToolProviderType.BUILT_IN.value))
                        setting_params = tool.get("settings", {})
                        parameters = tool.get("parameters", {})
                        manual_input_params = [key for key, value in parameters.items() if value is not None]

                        parameters = {**parameters, **setting_params}
                        entity = AgentToolEntity(
                            provider_id=tool.get("provider_name", ""),
                            provider_type=provider_type,
                            tool_name=tool.get("tool_name", ""),
                            tool_parameters=parameters,
                            plugin_unique_identifier=tool.get("plugin_unique_identifier", None),
                        )

                        extra = tool.get("extra", {})

                        tool_runtime = ToolManager.get_agent_tool_runtime(
                            self.tenant_id, self.app_id, entity, self.invoke_from
                        )
                        if tool_runtime.entity.description:
                            tool_runtime.entity.description.llm = (
                                extra.get("descrption", "") or tool_runtime.entity.description.llm
                            )
                        for tool_runtime_params in tool_runtime.entity.parameters:
                            tool_runtime_params.form = (
                                ToolParameter.ToolParameterForm.FORM
                                if tool_runtime_params.name in manual_input_params
                                else tool_runtime_params.form
                            )
                        manual_input_value = {}
                        if tool_runtime.entity.parameters:
                            manual_input_value = {
                                key: value for key, value in parameters.items() if key in manual_input_params
                            }
                        runtime_parameters = {
                            **tool_runtime.runtime.runtime_parameters,
                            **manual_input_value,
                        }
                        tool_value.append(
                            {
                                **tool_runtime.entity.model_dump(mode="json"),
                                "runtime_parameters": runtime_parameters,
                                "provider_type": provider_type.value,
                            }
                        )
                    value = tool_value
                if parameter.type == "model-selector":
                    value = cast(dict[str, Any], value)
                    model_instance, model_schema = self._fetch_model(value)
                    # memory config
                    history_prompt_messages = []
                    if node_data.memory:
                        memory = self._fetch_memory(model_instance)
                        if memory:
                            prompt_messages = memory.get_history_prompt_messages(
                                message_limit=node_data.memory.window.size if node_data.memory.window.size else None
                            )
                            history_prompt_messages = [
                                prompt_message.model_dump(mode="json") for prompt_message in prompt_messages
                            ]
                    value["history_prompt_messages"] = history_prompt_messages
                    if model_schema:
                        # remove structured output feature to support old version agent plugin
                        model_schema = self._remove_unsupported_model_features_for_old_version(model_schema)
                        value["entity"] = model_schema.model_dump(mode="json")
                    else:
                        value["entity"] = None
            result[parameter_name] = value

        return result

    @classmethod
    def _extract_variable_selector_to_variable_mapping(
        cls,
        *,
        graph_config: Mapping[str, Any],
        node_id: str,
        node_data: BaseNodeData,
    ) -> Mapping[str, Sequence[str]]:
        """
        Extract variable selector to variable mapping
        :param graph_config: graph config
        :param node_id: node id
        :param node_data: node data
        :return:
        """
        node_data = cast(AgentNodeData, node_data)
        result: dict[str, Any] = {}
        for parameter_name in node_data.agent_parameters:
            input = node_data.agent_parameters[parameter_name]
            if input.type in ["mixed", "constant"]:
                selectors = VariableTemplateParser(str(input.value)).extract_variable_selectors()
                for selector in selectors:
                    result[selector.variable] = selector.value_selector
            elif input.type == "variable":
                result[parameter_name] = input.value

        result = {node_id + "." + key: value for key, value in result.items()}

        return result

    @property
    def agent_strategy_icon(self) -> str | None:
        """
        Get agent strategy icon
        :return:
        """
        manager = PluginInstaller()
        plugins = manager.list_plugins(self.tenant_id)
        try:
            current_plugin = next(
                plugin
                for plugin in plugins
                if f"{plugin.plugin_id}/{plugin.name}"
                == cast(AgentNodeData, self.node_data).agent_strategy_provider_name
            )
            icon = current_plugin.declaration.icon
        except StopIteration:
            icon = None
        return icon

    def _fetch_memory(self, model_instance: ModelInstance) -> Optional[TokenBufferMemory]:
        # get conversation id
        conversation_id_variable = self.graph_runtime_state.variable_pool.get(
            ["sys", SystemVariableKey.CONVERSATION_ID.value]
        )
        if not isinstance(conversation_id_variable, StringSegment):
            return None
        conversation_id = conversation_id_variable.value

        # get conversation
        conversation = (
            db.session.query(Conversation)
            .filter(Conversation.app_id == self.app_id, Conversation.id == conversation_id)
            .first()
        )

        if not conversation:
            return None

        memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance)

        return memory

    def _fetch_model(self, value: dict[str, Any]) -> tuple[ModelInstance, AIModelEntity | None]:
        provider_manager = ProviderManager()
        provider_model_bundle = provider_manager.get_provider_model_bundle(
            tenant_id=self.tenant_id, provider=value.get("provider", ""), model_type=ModelType.LLM
        )
        model_name = value.get("model", "")
        model_credentials = provider_model_bundle.configuration.get_current_credentials(
            model_type=ModelType.LLM, model=model_name
        )
        provider_name = provider_model_bundle.configuration.provider.provider
        model_type_instance = provider_model_bundle.model_type_instance
        model_instance = ModelManager().get_model_instance(
            tenant_id=self.tenant_id,
            provider=provider_name,
            model_type=ModelType(value.get("model_type", "")),
            model=model_name,
        )
        model_schema = model_type_instance.get_model_schema(model_name, model_credentials)
        return model_instance, model_schema

    def _remove_unsupported_model_features_for_old_version(self, model_schema: AIModelEntity) -> AIModelEntity:
        if model_schema.features:
            for feature in model_schema.features:
                if feature.value not in AgentOldVersionModelFeatures:
                    model_schema.features.remove(feature)
        return model_schema
