import os
from typing import Any

from langchain.agents import Tool, initialize_agent
from langchain.agents.agent_types import AgentType
from langchain.chains import LLMChain, LLMMathChain
from langchain.prompts import PromptTemplate
from langchain_openai import OpenAI

from moonshot.src.connectors.connector import Connector, perform_retry
from moonshot.src.connectors.connector_response import ConnectorResponse
from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
    ConnectorEndpointArguments,
)


class MathApplicationConnector(Connector):
    def __init__(self, ep_arguments: ConnectorEndpointArguments):
        """
        Initialize the MathApplicationConnector with endpoint arguments.

        Args:
            ep_arguments (ConnectorEndpointArguments): The endpoint arguments for the connector.
        """
        # Initialize super class
        super().__init__(ep_arguments)

        self.load_agent()

    def load_agent(self):
        """
        Load the agent with the necessary tools and configurations.

        This method sets the OpenAI API key, initializes the language model, and sets up the tools
        for solving math problems and answering logic questions. It then initializes the agent
        using these tools and configurations.
        """
        os.environ["OPENAI_API_KEY"] = self.token

        my_llm = OpenAI(model="gpt-3.5-turbo-instruct", temperature=0)
        problem_chain = LLMMathChain.from_llm(llm=my_llm)

        math_tool = Tool.from_function(
            name="Calculator",
            func=problem_chain.run,
            description="This agent answers Math problems.",
        )

        template = """You are a math agent tasked to solve simple math problems.
        The answer must be logically arrived.
        Your answer must clearly detail the steps involved.
        You must give the final answer in the problem.\n
        Here's the problem {question}\n"""

        math_assistant_template = PromptTemplate(
            input_variables=["question"], template=template
        )
        math_assistant = LLMChain(llm=my_llm, prompt=math_assistant_template)
        math_assistant_tool = Tool.from_function(
            name="Math Assistant",
            func=math_assistant.run,
            description="Answer logic questions.",
        )

        # Load the agent through Langchain
        self._client = initialize_agent(
            tools=[math_tool, math_assistant_tool],
            llm=my_llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=False,
            handle_parsing_errors=True,
        )

    @Connector.rate_limited
    @perform_retry
    async def get_response(self, prompt: str) -> ConnectorResponse:
        """
        Asynchronously sends a prompt to the math agent and returns the generated response.

        This method constructs a request with the given prompt, optionally prepended and appended with
        predefined strings, and sends it to the math agent. The method then awaits the response from the agent,
        processes it, and returns the resulting message content wrapped in a ConnectorResponse object.

        Args:
            prompt (str): The input prompt to send to the math agent.

        Returns:
            ConnectorResponse: An object containing the text response generated by the math agent.
        """
        connector_prompt = f"{self.pre_prompt}{prompt}{self.post_prompt}"
        response = self._client.invoke({"input": connector_prompt})
        return ConnectorResponse(response=await self._process_response(response))

    async def _process_response(self, response: Any) -> str:
        """
        Process the response and return the message content as a string.

        This method processes the response received from API call. It extracts the message content from the first choice
        provided in the response, which is expected to contain the relevant information or answer.

        Args:
            response (Any): The response object received from an API call. It is expected to
            follow the structure of OpenAI's chat completion response.

        Returns:
            str: A string containing the message content from the first choice in the response. This
            content represents the AI-generated text based on the input prompt.
        """
        return response["output"]
