"""
OctoToolsSwarm: A multi-agent system for complex reasoning.
Implements the OctoTools framework using swarms.
"""

import json
import logging
import math  # Import the math module
import re
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from dotenv import load_dotenv
from swarms import Agent
from swarms.structs.conversation import Conversation

# from exa_search import exa_search as web_search_execute


# Load environment variables
load_dotenv()

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ToolType(Enum):
    """Defines the types of tools available."""

    IMAGE_CAPTIONER = "image_captioner"
    OBJECT_DETECTOR = "object_detector"
    WEB_SEARCH = "web_search"
    PYTHON_CALCULATOR = "python_calculator"
    # Add more tool types as needed


@dataclass
class Tool:
    """
    Represents an external tool.

    Attributes:
        name: Unique name of the tool.
        description: Description of the tool's function.
        metadata: Dictionary containing tool metadata.
        execute_func: Callable function that executes the tool's logic.
    """

    name: str
    description: str
    metadata: Dict[str, Any]
    execute_func: Callable

    def execute(self, **kwargs):
        """Executes the tool's logic, handling potential errors."""
        try:
            return self.execute_func(**kwargs)
        except Exception as e:
            logger.error(f"Error executing tool {self.name}: {str(e)}")
            return {"error": str(e)}


class AgentRole(Enum):
    """Defines the roles for agents in the OctoTools system."""

    PLANNER = "planner"
    VERIFIER = "verifier"
    SUMMARIZER = "summarizer"


class OctoToolsSwarm:
    """
    A multi-agent system implementing the OctoTools framework.

    Attributes:
        model_name: Name of the LLM model to use.
        max_iterations: Maximum number of action-execution iterations.
        base_path: Path for saving agent states.
        tools: List of available Tool objects.
    """

    def __init__(
        self,
        tools: List[Tool],
        model_name: str = "gemini/gemini-2.0-flash",
        max_iterations: int = 10,
        base_path: Optional[str] = None,
    ):
        """Initialize the OctoToolsSwarm system."""
        self.model_name = model_name
        self.max_iterations = max_iterations
        self.base_path = Path(base_path) if base_path else Path("./octotools_states")
        self.base_path.mkdir(exist_ok=True)
        self.tools = {tool.name: tool for tool in tools}  # Store tools in a dictionary

        # Initialize agents
        self._init_agents()

        # Create conversation tracker and memory
        self.conversation = Conversation()
        self.memory = []  # Store the trajectory

    def _init_agents(self) -> None:
        """Initialize all agents with their specific roles and prompts."""
        # Planner agent
        self.planner = Agent(
            agent_name="OctoTools-Planner",
            system_prompt=self._get_planner_prompt(),
            model_name=self.model_name,
            max_loops=3,
            saved_state_path=str(self.base_path / "planner.json"),
            verbose=True,
        )

        # Verifier agent
        self.verifier = Agent(
            agent_name="OctoTools-Verifier",
            system_prompt=self._get_verifier_prompt(),
            model_name=self.model_name,
            max_loops=1,
            saved_state_path=str(self.base_path / "verifier.json"),
            verbose=True,
        )

        # Summarizer agent
        self.summarizer = Agent(
            agent_name="OctoTools-Summarizer",
            system_prompt=self._get_summarizer_prompt(),
            model_name=self.model_name,
            max_loops=1,
            saved_state_path=str(self.base_path / "summarizer.json"),
            verbose=True,
        )

    def _get_planner_prompt(self) -> str:
        """Get the prompt for the planner agent (Improved with few-shot examples)."""
        tool_descriptions = "\n".join(
            [
                f"- {tool_name}: {self.tools[tool_name].description}"
                for tool_name in self.tools
            ]
        )
        return f"""You are the Planner in the OctoTools framework. Your role is to analyze the user's query,
        identify required skills, suggest relevant tools, and plan the steps to solve the problem.

        1. **Analyze the user's query:** Understand the requirements and identify the necessary skills and potentially relevant tools.
        2. **Perform high-level planning:**  Create a rough outline of how tools might be used to solve the problem.
        3. **Perform low-level planning (action prediction):**  At each step, select the best tool to use and formulate a specific sub-goal for that tool, considering the current context.

        Available Tools:
        {tool_descriptions}

        Output your response in JSON format.  Here are examples for different stages:

        **Query Analysis (High-Level Planning):**
        Example Input:
        Query: "What is the capital of France?"

        Example Output:
        ```json
        {{
            "summary": "The user is asking for the capital of France.",
            "required_skills": ["knowledge retrieval"],
            "relevant_tools": ["Web_Search_Tool"]
        }}
        ```

        **Action Prediction (Low-Level Planning):**
        Example Input:
        Context: {{ "query": "What is the capital of France?", "available_tools": ["Web_Search_Tool"] }}

        Example Output:
        ```json
        {{
            "justification": "The Web_Search_Tool can be used to directly find the capital of France.",
            "context": {{}},
            "sub_goal": "Search the web for 'capital of France'.",
            "tool_name": "Web_Search_Tool"
        }}
        ```
        Another Example:
        Context: {{"query": "How many objects are in the image?", "available_tools": ["Image_Captioner_Tool", "Object_Detector_Tool"], "image": "objects.png"}}

        Example Output:
        ```json
        {{
            "justification": "First, get a general description of the image to understand the context.",
            "context": {{ "image": "objects.png" }},
            "sub_goal": "Generate a description of the image.",
            "tool_name": "Image_Captioner_Tool"
        }}
        ```

        Example for Finding Square Root:
        Context: {{"query": "What is the square root of the number of objects in the image?", "available_tools": ["Object_Detector_Tool", "Python_Calculator_Tool"], "image": "objects.png", "Object_Detector_Tool_result": ["object1", "object2", "object3", "object4"]}}

        Example Output:
        ```json
        {{
            "justification": "We have detected 4 objects in the image. Now we need to find the square root of 4.",
            "context": {{}},
            "sub_goal": "Calculate the square root of 4",
            "tool_name": "Python_Calculator_Tool"
        }}
        ```

        Your output MUST be a single, valid JSON object with the following keys:
            - justification (string): Your reasoning.
            - context (dict):  A dictionary containing relevant information.
            - sub_goal (string): The specific instruction for the tool.
            - tool_name (string): The EXACT name of the tool to use.

            Do NOT include any text outside of the JSON object.
        """

    def _get_verifier_prompt(self) -> str:
        """Get the prompt for the verifier agent (Improved with few-shot examples)."""
        return """You are the Context Verifier in the OctoTools framework. Your role is to analyze the current context
        and memory to determine if the problem is solved, if there are any inconsistencies, or if further steps are needed.

        Output your response in JSON format:

        Expected output structure:
        ```json
        {
            "completeness": "Indicate whether the query is fully, partially, or not answered.",
            "inconsistencies": "List any inconsistencies found in the context or memory.",
            "verification_needs": "List any information that needs further verification.",
            "ambiguities": "List any ambiguities found in the context or memory.",
            "stop_signal": true/false
        }
        ```

        Example Input:
        Context: { "last_result": { "result": "Caption: The image shows a cat." } }
        Memory: [ { "component": "Action Predictor", "result": { "tool_name": "Image_Captioner_Tool" } } ]

        Example Output:
        ```json
        {
            "completeness": "partial",
            "inconsistencies": [],
            "verification_needs": ["Object detection to confirm the presence of a cat."],
            "ambiguities": [],
            "stop_signal": false
        }
        ```

        Another Example:
        Context: { "last_result": { "result": ["Detected object: cat"] } }
        Memory:  [ { "component": "Action Predictor", "result": { "tool_name": "Object_Detector_Tool" } } ]

        Example Output:
        ```json
        {
            "completeness": "yes",
            "inconsistencies": [],
            "verification_needs": [],
            "ambiguities": [],
            "stop_signal": true
        }
        ```

        Square Root Example:
        Context: {
            "query": "What is the square root of the number of objects in the image?",
            "image": "example.png",
            "Object_Detector_Tool_result": ["object1", "object2", "object3", "object4"],
            "Python_Calculator_Tool_result": "Result of 4**0.5 is 2.0"
        }
        Memory: [
            { "component": "Action Predictor", "result": { "tool_name": "Object_Detector_Tool" } },
            { "component": "Action Predictor", "result": { "tool_name": "Python_Calculator_Tool" } }
        ]

        Example Output:
        ```json
        {
            "completeness": "yes",
            "inconsistencies": [],
            "verification_needs": [],
            "ambiguities": [],
            "stop_signal": true
        }
        ```
        """

    def _get_summarizer_prompt(self) -> str:
        """Get the prompt for the summarizer agent (Improved with few-shot examples)."""
        return """You are the Solution Summarizer in the OctoTools framework.  Your role is to synthesize the final
        answer to the user's query based on the complete trajectory of actions and results.

        Output your response in JSON format:

        Expected output structure:
         ```json
        {
            "final_answer": "Provide a clear and concise answer to the original query."
        }
        ```
        Example Input:
        Memory: [
            {"component": "Query Analyzer", "result": {"summary": "Find the capital of France."}},
            {"component": "Action Predictor", "result": {"tool_name": "Web_Search_Tool"}},
            {"component": "Tool Execution", "result": {"result": "The capital of France is Paris."}}
        ]

        Example Output:
        ```json
        {
            "final_answer": "The capital of France is Paris."
        }
        ```

        Square Root Example:
        Memory: [
            {"component": "Query Analyzer", "result": {"summary": "Find the square root of the number of objects in the image."}},
            {"component": "Action Predictor", "result": {"tool_name": "Object_Detector_Tool", "sub_goal": "Detect objects in the image"}},
            {"component": "Tool Execution", "result": {"result": ["object1", "object2", "object3", "object4"]}},
            {"component": "Action Predictor", "result": {"tool_name": "Python_Calculator_Tool", "sub_goal": "Calculate the square root of 4"}},
            {"component": "Tool Execution", "result": {"result": "Result of 4**0.5 is 2.0"}}
        ]

        Example Output:
        ```json
        {
            "final_answer": "The square root of the number of objects in the image is 2.0. There are 4 objects in the image, and the square root of 4 is 2.0."
        }
        ```
        """

    def _safely_parse_json(self, json_str: str) -> Dict[str, Any]:
        """Safely parse JSON, handling errors and using recursive descent."""
        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            logger.warning(
                f"JSONDecodeError: Attempting to extract JSON from: {json_str}"
            )
            try:
                # More robust JSON extraction with recursive descent
                def extract_json(s):
                    stack = []
                    start = -1
                    for i, c in enumerate(s):
                        if c == "{":
                            if not stack:
                                start = i
                            stack.append(c)
                        elif c == "}":
                            if stack:
                                stack.pop()
                                if not stack and start != -1:
                                    return s[start : i + 1]
                    return None

                extracted_json = extract_json(json_str)
                if extracted_json:
                    logger.info(f"Extracted JSON: {extracted_json}")
                    return json.loads(extracted_json)
                else:
                    logger.error("Failed to extract JSON using recursive descent.")
                    return {
                        "error": "Failed to parse JSON",
                        "content": json_str,
                    }
            except Exception as e:
                logger.exception(f"Error during JSON extraction: {e}")
                return {
                    "error": "Failed to parse JSON",
                    "content": json_str,
                }

    def _execute_tool(self, tool_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Executes a tool based on its name and provided context."""
        if tool_name not in self.tools:
            return {"error": f"Tool '{tool_name}' not found."}

        tool = self.tools[tool_name]
        try:
            # For Python Calculator tool, handle object counts from Object Detector
            if tool_name == "Python_Calculator_Tool":
                # Check for object detector results
                object_detector_result = context.get("Object_Detector_Tool_result")
                if object_detector_result and isinstance(object_detector_result, list):
                    # Calculate the number of objects
                    num_objects = len(object_detector_result)
                    # If sub_goal doesn't already contain an expression, create one
                    if (
                        "sub_goal" in context
                        and "Calculate the square root" in context["sub_goal"]
                    ):
                        context["expression"] = f"{num_objects}**0.5"
                    elif "expression" not in context:
                        # Default to square root if no expression is specified
                        context["expression"] = f"{num_objects}**0.5"

            # Filter context: only pass expected inputs to the tool
            valid_inputs = {
                k: v
                for k, v in context.items()
                if k in tool.metadata.get("input_types", {})
            }
            result = tool.execute(**valid_inputs)
            return {"result": result}
        except Exception as e:
            logger.exception(f"Error executing tool {tool_name}: {e}")
            return {"error": str(e)}

    def _run_agent(self, agent: Agent, input_prompt: str) -> Dict[str, Any]:
        """Runs a swarms agent, handling output and JSON parsing."""
        try:
            # Construct the full input, including the system prompt
            full_input = f"{agent.system_prompt}\n\n{input_prompt}"

            # Run the agent and capture the output
            agent_response = agent.run(full_input)

            logger.info(f"DEBUG: Raw agent response: {agent_response}")

            # Extract the LLM's response (remove conversation history, etc.)
            response_text = agent_response  # Assuming direct return

            # Try to parse the response as JSON
            parsed_response = self._safely_parse_json(response_text)

            return parsed_response

        except Exception as e:
            logger.exception(f"Error running agent {agent.agent_name}: {e}")
            return {"error": f"Agent {agent.agent_name} failed: {str(e)}"}

    def run(self, query: str, image: Optional[str] = None) -> Dict[str, Any]:
        """Execute the task through the multi-agent workflow."""
        logger.info(f"Starting task: {query}")

        try:
            # Step 1: Query Analysis (High-Level Planning)
            planner_input = (
                f"Analyze the following query and determine the necessary skills and"
                f" relevant tools: {query}"
            )
            query_analysis = self._run_agent(self.planner, planner_input)

            if "error" in query_analysis:
                return {
                    "error": f"Planner query analysis failed: {query_analysis['error']}",
                    "trajectory": self.memory,
                    "conversation": self.conversation.return_history_as_string(),
                }

            self.memory.append(
                {
                    "step": 0,
                    "component": "Query Analyzer",
                    "result": query_analysis,
                }
            )
            self.conversation.add(
                role=self.planner.agent_name,
                content=json.dumps(query_analysis),
            )

            # Initialize context with the query and image (if provided)
            context = {"query": query}
            if image:
                context["image"] = image

            # Add available tools to context
            if "relevant_tools" in query_analysis:
                context["available_tools"] = query_analysis["relevant_tools"]
            else:
                # If no relevant tools specified, make all tools available
                context["available_tools"] = list(self.tools.keys())

            step_count = 1

            # Step 2: Iterative Action-Execution Loop
            while step_count <= self.max_iterations:
                logger.info(f"Starting iteration {step_count} of {self.max_iterations}")

                # Step 2a: Action Prediction (Low-Level Planning)
                action_planner_input = (
                    f"Current Context: {json.dumps(context)}\nAvailable Tools:"
                    f" {', '.join(context.get('available_tools', list(self.tools.keys())))}\nPlan the"
                    " next step."
                )
                action = self._run_agent(self.planner, action_planner_input)
                if "error" in action:
                    logger.error(f"Error in action prediction: {action['error']}")
                    return {
                        "error": f"Planner action prediction failed: {action['error']}",
                        "trajectory": self.memory,
                        "conversation": self.conversation.return_history_as_string(),
                    }
                self.memory.append(
                    {
                        "step": step_count,
                        "component": "Action Predictor",
                        "result": action,
                    }
                )
                self.conversation.add(
                    role=self.planner.agent_name,
                    content=json.dumps(action),
                )

                # Input Validation for Action (Relaxed)
                if (
                    not isinstance(action, dict)
                    or "tool_name" not in action
                    or "sub_goal" not in action
                ):
                    error_msg = (
                        "Action prediction did not return required fields (tool_name,"
                        " sub_goal) or was not a dictionary."
                    )
                    logger.error(error_msg)
                    self.memory.append(
                        {
                            "step": step_count,
                            "component": "Error",
                            "result": error_msg,
                        }
                    )
                    break

                # Step 2b: Execute Tool
                tool_execution_context = {
                    **context,
                    **action.get("context", {}),  # Add any additional context
                    "sub_goal": action["sub_goal"],  # Pass sub_goal to tool
                }

                tool_result = self._execute_tool(
                    action["tool_name"], tool_execution_context
                )

                self.memory.append(
                    {
                        "step": step_count,
                        "component": "Tool Execution",
                        "result": tool_result,
                    }
                )

                # Step 2c: Context Update - Store result with a descriptive key
                if "result" in tool_result:
                    context[f"{action['tool_name']}_result"] = tool_result["result"]
                if "error" in tool_result:
                    context[f"{action['tool_name']}_error"] = tool_result["error"]

                # Step 2d: Context Verification
                verifier_input = (
                    f"Current Context: {json.dumps(context)}\nMemory:"
                    f" {json.dumps(self.memory)}\nQuery: {query}"
                )
                verification = self._run_agent(self.verifier, verifier_input)
                if "error" in verification:
                    return {
                        "error": f"Verifier failed: {verification['error']}",
                        "trajectory": self.memory,
                        "conversation": self.conversation.return_history_as_string(),
                    }

                self.memory.append(
                    {
                        "step": step_count,
                        "component": "Context Verifier",
                        "result": verification,
                    }
                )
                self.conversation.add(
                    role=self.verifier.agent_name,
                    content=json.dumps(verification),
                )

                # Check for stop signal from Verifier
                if verification.get("stop_signal") is True:
                    logger.info(
                        "Received stop signal from verifier. Stopping iterations."
                    )
                    break

                # Safety mechanism - if we've executed the same tool multiple times
                same_tool_count = sum(
                    1
                    for m in self.memory
                    if m.get("component") == "Action Predictor"
                    and m.get("result", {}).get("tool_name") == action.get("tool_name")
                )

                if same_tool_count > 3:
                    logger.warning(
                        f"Tool {action.get('tool_name')} used more than 3 times. Forcing stop."
                    )
                    break

                step_count += 1

            # Step 3: Solution Summarization
            summarizer_input = f"Complete Trajectory: {json.dumps(self.memory)}\nOriginal Query: {query}"

            summarization = self._run_agent(self.summarizer, summarizer_input)
            if "error" in summarization:
                return {
                    "error": f"Summarizer failed: {summarization['error']}",
                    "trajectory": self.memory,
                    "conversation": self.conversation.return_history_as_string(),
                }
            self.conversation.add(
                role=self.summarizer.agent_name,
                content=json.dumps(summarization),
            )

            return {
                "final_answer": summarization.get("final_answer", "No answer found."),
                "trajectory": self.memory,
                "conversation": self.conversation.return_history_as_string(),
            }

        except Exception as e:
            logger.exception(f"Unexpected error in run method: {e}")  # More detailed
            return {
                "error": str(e),
                "trajectory": self.memory,
                "conversation": self.conversation.return_history_as_string(),
            }

    def save_state(self) -> None:
        """Save the current state of all agents."""
        for agent in [self.planner, self.verifier, self.summarizer]:
            try:
                agent.save_state()
            except Exception as e:
                logger.error(f"Error saving state for {agent.agent_name}: {str(e)}")

    def load_state(self) -> None:
        """Load the saved state of all agents."""
        for agent in [self.planner, self.verifier, self.summarizer]:
            try:
                agent.load_state()
            except Exception as e:
                logger.error(f"Error loading state for {agent.agent_name}: {str(e)}")


# --- Example Usage ---


# Define dummy tool functions (replace with actual implementations)
def image_captioner_execute(
    image: str, prompt: str = "Describe the image", **kwargs
) -> str:
    """Dummy image captioner."""
    print(f"image_captioner_execute called with image: {image}, prompt: {prompt}")
    return f"Caption for {image}: A descriptive caption (dummy)."  # Simplified


def object_detector_execute(image: str, labels: List[str] = [], **kwargs) -> List[str]:
    """Dummy object detector, handles missing labels gracefully."""
    print(f"object_detector_execute called with image: {image}, labels: {labels}")
    if not labels:
        return [
            "object1",
            "object2",
            "object3",
            "object4",
        ]  # Return default objects if no labels
    return [f"Detected {label}" for label in labels]  # Simplified


def web_search_execute(query: str, **kwargs) -> str:
    """Dummy web search."""
    print(f"web_search_execute called with query: {query}")
    return f"Search results for '{query}'..."  # Simplified


def python_calculator_execute(expression: str, **kwargs) -> str:
    """Python calculator (using math module)."""
    print(f"python_calculator_execute called with: {expression}")
    try:
        # Safely evaluate only simple expressions involving numbers and basic operations
        if re.match(r"^[0-9+\-*/().\s]+$", expression):
            result = eval(expression, {"__builtins__": {}, "math": math})
            return f"Result of {expression} is {result}"
        else:
            return "Error: Invalid expression for calculator."
    except Exception as e:
        return f"Error: {e}"


# Create utility function to get default tools
def get_default_tools() -> List[Tool]:
    """Returns a list of default tools that can be used with OctoToolsSwarm."""
    image_captioner = Tool(
        name="Image_Captioner_Tool",
        description="Generates a caption for an image.",
        metadata={
            "input_types": {"image": "str", "prompt": "str"},
            "output_type": "str",
            "limitations": "May struggle with complex scenes or ambiguous objects.",
            "best_practices": "Use with clear, well-lit images. Provide specific prompts for better results.",
        },
        execute_func=image_captioner_execute,
    )

    object_detector = Tool(
        name="Object_Detector_Tool",
        description="Detects objects in an image.",
        metadata={
            "input_types": {"image": "str", "labels": "list"},
            "output_type": "list",
            "limitations": "Accuracy depends on the quality of the image and the clarity of the objects.",
            "best_practices": "Provide a list of specific object labels to detect. Use high-resolution images.",
        },
        execute_func=object_detector_execute,
    )

    web_search = Tool(
        name="Web_Search_Tool",
        description="Performs a web search.",
        metadata={
            "input_types": {"query": "str"},
            "output_type": "str",
            "limitations": "May not find specific or niche information.",
            "best_practices": "Use specific and descriptive keywords for better results.",
        },
        execute_func=web_search_execute,
    )

    calculator = Tool(
        name="Python_Calculator_Tool",
        description="Evaluates a Python expression.",
        metadata={
            "input_types": {"expression": "str"},
            "output_type": "str",
            "limitations": "Cannot handle complex mathematical functions or libraries.",
            "best_practices": "Use for basic arithmetic and simple calculations.",
        },
        execute_func=python_calculator_execute,
    )

    return [image_captioner, object_detector, web_search, calculator]


# Only execute the example when this script is run directly
# if __name__ == "__main__":
#     print("Running OctoToolsSwarm example...")

#     # Create an OctoToolsSwarm agent with default tools
#     tools = get_default_tools()
#     agent = OctoToolsSwarm(tools=tools)

#     # Example query
#     query = "What is the square root of the number of objects in this image?"

#     # Create a dummy image file for testing if it doesn't exist
#     image_path = "example.png"
#     if not os.path.exists(image_path):
#         with open(image_path, "w") as f:
#             f.write("Dummy image content")
#         print(f"Created dummy image file: {image_path}")

#     # Run the agent
#     result = agent.run(query, image=image_path)

#     # Display results
#     print("\n=== FINAL ANSWER ===")
#     print(result["final_answer"])

#     print("\n=== TRAJECTORY SUMMARY ===")
#     for step in result["trajectory"]:
#         print(f"Step {step.get('step', 'N/A')}: {step.get('component', 'Unknown')}")

#     print("\nOctoToolsSwarm example completed.")
