# -*- coding: utf-8 -*-
"""
Risk Manager Agent - Based on AgentScope ReActAgent
Uses LLM for risk assessment
"""
from typing import Any, Dict, Optional

from agentscope.agent import ReActAgent
from agentscope.memory import InMemoryMemory, LongTermMemoryBase
from agentscope.message import Msg
from agentscope.tool import Toolkit

from ..utils.progress import progress
from .prompt_factory import build_agent_system_prompt, clear_prompt_factory_cache


class RiskAgent(ReActAgent):
    """
    Risk Manager Agent - Uses LLM for risk assessment
    Inherits from AgentScope's ReActAgent
    """

    def __init__(
        self,
        model: Any,
        formatter: Any,
        name: str = "risk_manager",
        config: Optional[Dict[str, Any]] = None,
        long_term_memory: Optional[LongTermMemoryBase] = None,
        toolkit: Optional[Toolkit] = None,
    ):
        """
        Initialize Risk Manager Agent

        Args:
            model: LLM model instance
            formatter: Message formatter instance
            name: Agent name
            config: Configuration dictionary
            long_term_memory: Optional ReMeTaskLongTermMemory instance
        """
        object.__setattr__(self, "config", config or {})
        object.__setattr__(self, "agent_id", name)

        if toolkit is None:
            toolkit = Toolkit()
        object.__setattr__(self, "toolkit", toolkit)

        sys_prompt = self._load_system_prompt()

        kwargs = {
            "name": name,
            "sys_prompt": sys_prompt,
            "model": model,
            "formatter": formatter,
            "toolkit": toolkit,
            "memory": InMemoryMemory(),
            "max_iters": 10,
        }
        if long_term_memory:
            kwargs["long_term_memory"] = long_term_memory
            kwargs["long_term_memory_mode"] = "static_control"

        super().__init__(**kwargs)

    def _load_system_prompt(self) -> str:
        """Load system prompt for risk manager"""
        return build_agent_system_prompt(
            agent_id=self.agent_id,
            config_name=self.config.get("config_name", "default"),
            toolkit=self.toolkit,
        )

    async def reply(self, x: Msg = None) -> Msg:
        """
        Provide risk assessment

        Args:
            x: Input message (content must be str)

        Returns:
            Msg with risk warnings (content is str)
        """
        progress.update_status(self.name, None, "Assessing risk")

        result = await super().reply(x)

        progress.update_status(self.name, None, "Risk assessment completed")

        return result

    def reload_runtime_assets(self, active_skill_dirs: Optional[list] = None) -> None:
        """Reload toolkit and system prompt from current run assets."""
        from .toolkit_factory import create_agent_toolkit

        clear_prompt_factory_cache()
        self.toolkit = create_agent_toolkit(
            self.agent_id,
            self.config.get("config_name", "default"),
            active_skill_dirs=active_skill_dirs,
        )
        self._apply_runtime_sys_prompt(self._load_system_prompt())

    def _apply_runtime_sys_prompt(self, sys_prompt: str) -> None:
        """Update the prompt used by future turns and the cached system msg."""
        self._sys_prompt = sys_prompt
        for msg, _marks in self.memory.content:
            if getattr(msg, "role", None) == "system":
                msg.content = sys_prompt
                break