"""
An Explore and Exploit agent implementation.
This module contains the Explore and Exploit agent class and its configuration.
The agent is designed to first explore available tools to understand their capabilities
and effects, then exploit this knowledge to effectively solve tasks.
"""
# pylint: disable=broad-exception-caught
import os
import json
from typing import Optional, Union, Dict, List
from dataclasses import dataclass
from mcp.types import TextContent

from mcpuniverse.mcp.manager import MCPManager
from mcpuniverse.llm.base import BaseLLM
from mcpuniverse.common.logger import get_logger
from mcpuniverse.tracer import Tracer
from mcpuniverse.callbacks.base import CallbackMessage, send_message_async, MessageType

from .base import BaseAgentConfig, BaseAgent
from .utils import build_system_prompt
from .types import AgentResponse

DEFAULT_CONFIG_FOLDER = os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs")


@dataclass
class ExploreAndExploitConfig(BaseAgentConfig):
    """
    Configuration class for Explore and Exploit agents.
    Attributes:
        system_prompt (str): The system prompt template file or string.
        context_examples (str): Additional context examples for the agent.
        max_iterations (int): Maximum number of reasoning iterations.
        exploration_phase_iterations (int): Number of iterations dedicated to tool exploration.
    """
    system_prompt: str = os.path.join(DEFAULT_CONFIG_FOLDER, "explore_and_exploit_prompt.j2")
    context_examples: str = ""
    max_iterations: int = 15
    summarize_tool_response: bool = False
    exploration_phase_iterations: int = 6


class ExploreAndExploit(BaseAgent):
    """
    Explore and Exploit agent implementation.
    This class implements an agent that first explores available tools to understand
    their capabilities and effects, then exploits this knowledge to effectively solve tasks.
    The agent maintains a knowledge base of tool experiences to inform future decisions.
    Attributes:
        config_class (Type[ExploreAndExploitConfig]): The configuration class for this agent.
        alias (List[str]): Alternative names for this agent type.
    """
    config_class = ExploreAndExploitConfig
    alias = ["explore_and_exploit", "explore_exploit"]

    def __init__(
            self,
            mcp_manager: MCPManager,
            llm: BaseLLM,
            config: Optional[Union[Dict, str]] = None
    ):
        """
        Initialize an Explore and Exploit agent.
        Args:
            mcp_manager (MCPManager): An MCP server manager for handling tool interactions.
            llm (BaseLLM): A language model for generating responses.
            config (Optional[Union[Dict, str]]): Agent configuration as a dictionary or file path.
        """
        super().__init__(mcp_manager=mcp_manager, llm=llm, config=config)
        self._logger = get_logger(f"{self.__class__.__name__}:{self._name}")
        self._history: List[str] = []
        self._tool_knowledge: Dict[str, Dict] = {}
        self._exploration_complete: bool = False
        self._current_iteration: int = 0

    def _clean_json_response(self, response: str) -> str:
        """Clean and extract JSON content from LLM response using base class logic."""
        # Use the same logic as base class call_tool method
        response = response.strip().strip('`').strip()
        if response.startswith("json"):
            response = response[4:].strip()
        # Handle thinking tags
        if "<think>" in response:
            response = response.split("</think>")[1].strip()
        return response

    def _build_prompt(self, question: str):
        """
        Construct the prompt for the language model.
        Args:
            question (str): The user's question or task.
        Returns:
            str: The constructed prompt including system instructions, context, and history.
        """
        params = {"INSTRUCTION": self._config.instruction, "QUESTION": question}
        if self._config.context_examples:
            params.update({"CONTEXT_EXAMPLES": self._config.context_examples})
        params.update(self._config.template_vars)
        if self._history:
            params.update({"HISTORY": "\n".join(self._history)})

        # Add tool knowledge to the prompt
        if self._tool_knowledge:
            tool_knowledge_str = self._format_tool_knowledge()
            params.update({"TOOL_KNOWLEDGE": tool_knowledge_str})

        # Add phase information
        phase = "exploration" if not self._exploration_complete else "exploitation"
        params.update({
            "CURRENT_PHASE": phase,
            "EXPLORATION_ITERATIONS_LEFT": max(0, self._config.exploration_phase_iterations - self._current_iteration)
        })

        return build_system_prompt(
            system_prompt_template=self._config.system_prompt,
            tool_prompt_template=self._config.tools_prompt,
            tools=self._tools,
            **params
        )

    def _format_tool_knowledge(self) -> str:
        """
        Format the accumulated tool knowledge for inclusion in the prompt.
        Returns:
            str: Formatted tool knowledge string.
        """
        if not self._tool_knowledge:
            return "No tool knowledge accumulated yet."

        knowledge_lines = []
        for tool_key, knowledge in self._tool_knowledge.items():
            server, tool_name = tool_key.split(".", 1)
            knowledge_lines.extend([
                f"Tool: {server}.{tool_name}",
                f"  Purpose: {knowledge.get('purpose', 'Unknown')}",
                f"  Caveats: {knowledge.get('caveats', 'To be determined')}",
                ""
            ])
        return "\n".join(knowledge_lines)

    def _add_history(self, history_type: str, message: str):
        """
        Add a record to the agent's conversation history.
        Args:
            history_type (str): The type of the history entry.
            message (str): The content of the history entry.
        """
        self._history.append(f"{history_type.title()}: {message}")

    def _update_tool_knowledge(self, server: str, tool_name: str, purpose: str, caveats: str):
        """
        Update the tool knowledge base with new insights.
        Args:
            server (str): The server name.
            tool_name (str): The tool name.
            purpose (str): The tool's purpose.
            caveats (str): Caveats of the tool.
        """
        tool_key = f"{server}.{tool_name}"
        self._tool_knowledge[tool_key] = {"purpose": purpose, "caveats": caveats}

    async def _send_callback_message(self, callbacks: list, data: str):
        """Send a callback message with consistent formatting."""
        message = CallbackMessage(
            source=f"{self._name}:agent:{self.__class__.__name__}",
            type=MessageType.LOG,
            data=data,
            metadata={"event": "plain_text", "data": data}
        )
        await send_message_async(callbacks, message)

    def _parse_llm_response(self, response: str) -> dict:
        """Parse and validate LLM response using base class logic."""
        # Use the same cleaning logic as base class
        cleaned_response = self._clean_json_response(response)
        parsed_response = json.loads(cleaned_response)

        if "thought" not in parsed_response:
            raise ValueError("Invalid response format: missing 'thought' field")
        return parsed_response

    async def _analyze_and_update_tool_knowledge(self, action: dict, observation: str, tracer: Tracer, callbacks: list):
        """
        Analyze actual tool execution results and update knowledge base.
        Args:
            action (dict): The tool action that was executed
            observation (str): The actual observation/result from tool execution
            tracer (Tracer): Tracer for LLM calls
            callbacks (list): Callbacks for LLM calls
        """
        # Only analyze during exploration phase
        if self._exploration_complete:
            return

        tool_key = f"{action['server']}.{action['tool']}"
        existing_knowledge = self._tool_knowledge.get(tool_key, {})

        # Build context for existing knowledge
        knowledge_context = (
            f"Previous Knowledge About This Tool:\n"
            f"- Purpose: {existing_knowledge.get('purpose', 'Unknown')}\n"
            f"- Caveats: {existing_knowledge.get('caveats', 'To be determined')}\n\n"
            if existing_knowledge else "Previous Knowledge: None (first time using this tool)\n\n"
        )

        analysis_prompt = f"""
Analyze the following tool usage and its actual results to update knowledge about the tool's capabilities and behavior.
{knowledge_context}Current Tool Usage:
Tool Used: {action['server']}.{action['tool']}
Arguments: {action.get('arguments', 'None')}
Actual Result: {observation}
Based on this actual execution and any previous knowledge, provide a JSON response with updated/refined knowledge:
```json
{{
    "purpose": "What is this tool's main purpose? (refine based on previous knowledge and current observation)",
    "caveats": "What caveats for failure cases, limitations, or constraints were observed? \
          (add to or refine previous caveats)"
}}
```
Instructions:
- If this is the first use, provide new knowledge based on the observation
- If there's previous knowledge, BUILD UPON and REFINE it with the new observation
- Don't discard previous insights unless they're clearly contradicted by the new observation
- Enhance and improve the knowledge rather than replacing it
- Make your response concise and to the point
"""

        try:
            analysis_response = self._llm.generate(
                messages=[{"role": "user", "content": analysis_prompt}],
                tracer=tracer,
                callbacks=callbacks
            )

            # Use base class logic for JSON cleaning and parsing
            cleaned_response = self._clean_json_response(analysis_response)
            update_type = "Updated" if existing_knowledge else "Created"

            await self._send_callback_message(callbacks, f"{update_type} tool knowledge: {cleaned_response}")

            analysis = json.loads(cleaned_response)
            self._update_tool_knowledge(
                action['server'],
                action['tool'],
                analysis.get('purpose', 'Unknown'),
                analysis.get('caveats', 'To be determined')
            )

            self._add_history(
                history_type="knowledge_update",
                message=f"{update_type} knowledge for {action['server']}.{action['tool']} based on observed behavior"
            )

        except Exception as e:
            self._logger.warning("Failed to analyze tool knowledge: %s", str(e))
            # Fallback: record basic usage
            if tool_key not in self._tool_knowledge:
                self._update_tool_knowledge(
                    action['server'],
                    action['tool'],
                    "Purpose to be determined",
                    "Caveats to be determined"
                )

    async def _handle_final_answer(self, parsed_response: dict, tracer: Tracer,
                                   callbacks: list) -> Optional[AgentResponse]:
        """Handle final answer response."""
        if "answer" not in parsed_response:
            return None

        # Check if we're in exploration phase - final answers are not allowed
        if not self._exploration_complete:
            self._logger.warning("Final answer attempted during exploration phase - ignoring")
            error_msg = "🚫 Final answers are not allowed during exploration phase. " + \
                        "You must continue exploring tools to build knowledge first."
            self._add_history(history_type="error", message=error_msg)
            await self._send_callback_message(callbacks, error_msg)
            return None

        self._add_history(history_type="answer", message=parsed_response["answer"])
        return AgentResponse(
            name=self._name,
            class_name=self.__class__.__name__,
            response=parsed_response["answer"],
            trace_id=tracer.trace_id
        )

    async def _execute_tool_action(self, action: dict, parsed_response: dict,
                                   tracer: Tracer, callbacks: list, iteration: int) -> tuple[bool, str]:
        """Execute the tool action and handle results."""
        self._add_history(
            history_type="action",
            message=f"Using tool `{action['tool']}` in server `{action['server']}`"
        )

        try:
            # Use base class call_tool method directly
            result = await self.call_tool(action, tracer=tracer, callbacks=callbacks)
            tool_content = result.content[0]

            if not isinstance(tool_content, TextContent):
                raise ValueError("Tool output is not a text")

            self._add_history(history_type="action input", message=str(action.get("arguments", "none")))

            # Use base class summarize_tool_response method if enabled
            tool_feedback = (
                self.summarize_tool_response(tool_content.text, context=json.dumps(action, indent=2), tracer=tracer)
                if self._config.summarize_tool_response
                else tool_content.text
            )

            self._add_history(history_type="observation", message=tool_feedback)
            await self._display_verbose_output(parsed_response, tool_feedback, iteration, callbacks)
            return True, tool_feedback

        except Exception as e:
            error_msg = str(e)
            self._add_history(history_type="observation", message=error_msg)
            return False, error_msg

    async def _handle_tool_action(self, parsed_response: dict, tracer: Tracer,
                                  callbacks: list, iteration: int) -> tuple[bool, str]:
        """Handle tool action execution."""
        if "action" not in parsed_response:
            raise ValueError("Invalid response format: no action or answer provided")

        self._add_history(history_type="thought", message=parsed_response["thought"])
        action = parsed_response["action"]

        if not self._is_valid_action(action):
            self._add_history(history_type="action", message=str(action))
            error_msg = "Invalid action"
            self._add_history(history_type="observation", message=error_msg)
            return False, error_msg

        return await self._execute_tool_action(action, parsed_response, tracer, callbacks, iteration)

    def _is_valid_action(self, action) -> bool:
        """Validate action format."""
        return (isinstance(action, dict) and "server" in action and "tool" in action)

    async def _display_verbose_output(self, parsed_response: dict, observation: str, iteration: int, callbacks: list):
        """Display verbose output through callbacks."""
        phase = "Exploration" if not self._exploration_complete else "Exploitation"

        output_lines = [
            f"\n{'=' * 66}",
            f"{phase} Phase - Iteration: {iteration + 1}",
            f"{'-' * 66}",
            f"\033[32mThought: {parsed_response['thought']}\n\033[0m",
            f"\033[31mAction: {parsed_response['action']}\n\033[0m",
            f"\033[33mObservation: {observation}\033[0m"
        ]

        if "tool_knowledge" in parsed_response:
            output_lines.append(f"\033[36mTool Knowledge Updated: {parsed_response['tool_knowledge']}\033[0m")

        await self._send_callback_message(callbacks, "\n".join(output_lines) + "\n")

    def _handle_execution_error(self, error: Exception, error_type: str):
        """Handle execution errors with appropriate logging and history updates."""
        self._logger.error("Failed to %s: %s", error_type, str(error))

        message = (
            "I encountered an error in parsing LLM response. Let me try again."
            if isinstance(error, json.JSONDecodeError)
            else "I encountered an unexpected error. Let me try a different approach."
        )
        self._add_history(history_type="error", message=message)

    async def _execute(
            self,
            message: Union[str, List[str]],
            output_format: Optional[Union[str, Dict]] = None,
            **kwargs
    ) -> AgentResponse:
        """
        Execute the Explore and Exploit agent's reasoning and action loop.
        This method first explores available tools to understand their capabilities,
        then exploits this knowledge to solve the task effectively.
        Args:
            message (Union[str, List[str]]): The user's message or a list of messages.
            output_format (Optional[Union[str, Dict]]): Desired format for the output.
            **kwargs: Additional keyword arguments.
        Returns:
            AgentResponse: The agent's final response, including the answer and trace information.
        """
        # Prepare message and initialize
        if isinstance(message, (list, tuple)):
            message = "\n".join(message)
        if output_format is not None:
            # Use base class method for output format prompt
            message = message + "\n\n" + self._get_output_format_prompt(output_format)

        tracer = kwargs.get("tracer", Tracer())
        callbacks = kwargs.get("callbacks", [])
        self._current_iteration = 0
        self._exploration_complete = False

        # Main execution loop
        for iteration in range(self._config.max_iterations):
            self._current_iteration = iteration

            # Check phase transition
            if iteration >= self._config.exploration_phase_iterations:
                self._exploration_complete = True

            # Generate and process response
            prompt = self._build_prompt(message)
            response = self._llm.generate(
                messages=[{"role": "user", "content": prompt}],
                tracer=tracer,
                callbacks=callbacks
            )

            try:
                # Parse response
                parsed_response = self._parse_llm_response(response)

                # Check for final answer
                final_response = await self._handle_final_answer(parsed_response, tracer, callbacks)
                if final_response:
                    return final_response

                # Handle tool actions
                tool_executed, observation = await self._handle_tool_action(
                    parsed_response, tracer, callbacks, iteration)

                # Update tool knowledge based on actual observation
                if tool_executed:
                    await self._analyze_and_update_tool_knowledge(
                        parsed_response["action"], observation, tracer, callbacks
                    )

            except json.JSONDecodeError as e:
                self._handle_execution_error(e, "parse response")
            except Exception as e:
                self._handle_execution_error(e, "process response")

        # Return default response if max iterations reached
        return AgentResponse(
            name=self._name,
            class_name=self.__class__.__name__,
            response="I'm sorry, but I couldn't find a satisfactory answer within the allowed number of iterations.",
            trace_id=tracer.trace_id
        )

    def get_history(self) -> str:
        """Retrieve the agent's conversation history."""
        return "\n".join(self._history)

    def get_tool_knowledge(self) -> Dict[str, Dict]:
        """Retrieve the agent's accumulated tool knowledge."""
        return self._tool_knowledge.copy()

    def clear_history(self):
        """Clear the agent's conversation history."""
        self._history = []

    def reset(self):
        """Reset the agent."""
        self.clear_history()
        self._tool_knowledge = {}
        self._exploration_complete = False
        self._current_iteration = 0
