"""
Structured messages to an agent, typically from an LLM, to be handled by
an agent. The messages could represent, for example:
- information or data given to the agent
- request for information or data from the agent
- request to run a method of the agent
"""

import json
import textwrap
from abc import ABC
from random import choice
from typing import Any, Dict, List, Type

from docstring_parser import parse
from pydantic import BaseModel

from langroid.language_models.base import LLMFunctionSpec
from langroid.utils.pydantic_utils import (
    _recursive_purge_dict_key,
    generate_simple_schema,
)


class ToolMessage(ABC, BaseModel):
    """
    Abstract Class for a class that defines the structure of a "Tool" message from an
    LLM. Depending on context, "tools" are also referred to as "plugins",
    or "function calls" (in the context of OpenAI LLMs).
    Essentially, they are a way for the LLM to express its intent to run a special
    function or method. Currently these "tools" are handled by methods of the
    agent.

    Attributes:
        request (str): name of agent method to map to.
        purpose (str): purpose of agent method, expressed in general terms.
            (This is used when auto-generating the tool instruction to the LLM)
        result (str): example of result of agent method.
    """

    request: str
    purpose: str
    result: str = ""

    class Config:
        arbitrary_types_allowed = False
        validate_all = True
        validate_assignment = True
        # do not include these fields in the generated schema
        # since we don't require the LLM to specify them
        schema_extra = {"exclude": {"purpose", "result"}}

    @classmethod
    def instructions(cls) -> str:
        return ""

    @classmethod
    def require_recipient(cls) -> Type["ToolMessage"]:
        class ToolMessageWithRecipient(cls):  # type: ignore
            recipient: str  # no default, so it is required

        return ToolMessageWithRecipient

    @classmethod
    def examples(cls) -> List["ToolMessage"]:
        """
        Examples to use in few-shot demos with JSON formatting instructions.
        Returns:
        """
        return []

    @classmethod
    def usage_example(cls) -> str:
        """
        Instruction to the LLM showing an example of how to use the message.
        Returns:
            str: example of how to use the message
        """
        # pick a random example of the fields
        if len(cls.examples()) == 0:
            return ""
        ex = choice(cls.examples())
        return ex.json_example()

    def to_json(self) -> str:
        return self.json(indent=4, exclude={"result", "purpose"})

    def json_example(self) -> str:
        return self.json(indent=4, exclude={"result", "purpose"})

    def dict_example(self) -> Dict[str, Any]:
        return self.dict(exclude={"result", "purpose"})

    @classmethod
    def default_value(cls, f: str) -> Any:
        """
        Returns the default value of the given field, for the message-class
        Args:
            f (str): field name

        Returns:
            Any: default value of the field, or None if not set or if the
                field does not exist.
        """
        schema = cls.schema()
        properties = schema["properties"]
        return properties.get(f, {}).get("default", None)

    @classmethod
    def json_instructions(cls, tool: bool = False) -> str:
        """
        Default Instructions to the LLM showing how to use the tool/function-call.
        Works for GPT4 but override this for weaker LLMs if needed.

        Args:
            tool: instructions for Langroid-native tool use? (e.g. for non-OpenAI LLM)
                (or else it would be for OpenAI Function calls)
        Returns:
            str: instructions on how to use the message
        """
        # TODO: when we attempt to use a "simpler schema"
        # (i.e. all nested fields explicit without definitions),
        # we seem to get worse results, so we turn it off for now
        param_dict = (
            # cls.simple_schema() if tool else
            cls.llm_function_schema(request=True).parameters
        )
        return textwrap.dedent(
            f"""
            TOOL: {cls.default_value("request")}
            PURPOSE: {cls.default_value("purpose")} 
            JSON FORMAT: {
                json.dumps(param_dict, indent=4)
            }
            {"EXAMPLE: " + cls.usage_example() if cls.examples() else ""}
            """.lstrip()
        )

    @staticmethod
    def json_group_instructions() -> str:
        """Template for instructions for a group of tools.
        Works with GPT4 but override this for weaker LLMs if needed.
        """
        return textwrap.dedent(
            """
            === ALL AVAILABLE TOOLS and THEIR JSON FORMAT INSTRUCTIONS ===
            You have access to the following TOOLS to accomplish your task:

            {json_instructions}
            
            When one of the above TOOLs is applicable, you must express your 
            request as "TOOL:" followed by the request in the above JSON format.
            """
        )

    @classmethod
    def llm_function_schema(
        cls,
        request: bool = False,
        defaults: bool = True,
    ) -> LLMFunctionSpec:
        """
        Clean up the schema of the Pydantic class (which can recursively contain
        other Pydantic classes), to create a version compatible with OpenAI
        Function-call API.

        Adapted from this excellent library:
        https://github.com/jxnl/instructor/blob/main/instructor/function_calls.py

        Args:
            request: whether to include the "request" field in the schema.
                (we set this to True when using Langroid-native TOOLs as opposed to
                OpenAI Function calls)
            defaults: whether to include fields with default values in the schema,
                    in the "properties" section.

        Returns:
            LLMFunctionSpec: the schema as an LLMFunctionSpec

        """
        schema = cls.schema()
        docstring = parse(cls.__doc__ or "")
        parameters = {
            k: v for k, v in schema.items() if k not in ("title", "description")
        }
        for param in docstring.params:
            if (name := param.arg_name) in parameters["properties"] and (
                description := param.description
            ):
                if "description" not in parameters["properties"][name]:
                    parameters["properties"][name]["description"] = description

        excludes = (
            ["result", "purpose"] if request else ["request", "result", "purpose"]
        )
        # exclude 'excludes' from parameters["properties"]:
        parameters["properties"] = {
            field: details
            for field, details in parameters["properties"].items()
            if field not in excludes and (defaults or details.get("default") is None)
        }
        parameters["required"] = sorted(
            k
            for k, v in parameters["properties"].items()
            if ("default" not in v and k not in excludes)
        )
        if request:
            parameters["required"].append("request")

        if "description" not in schema:
            if docstring.short_description:
                schema["description"] = docstring.short_description
            else:
                schema["description"] = (
                    f"Correctly extracted `{cls.__name__}` with all "
                    f"the required parameters with correct types"
                )

        parameters.pop("exclude")
        _recursive_purge_dict_key(parameters, "title")
        _recursive_purge_dict_key(parameters, "additionalProperties")
        return LLMFunctionSpec(
            name=cls.default_value("request"),
            description=cls.default_value("purpose"),
            parameters=parameters,
        )

    @classmethod
    def simple_schema(cls) -> Dict[str, Any]:
        """
        Return a simplified schema for the message, with only the request and
        required fields.
        Returns:
            Dict[str, Any]: simplified schema
        """
        schema = generate_simple_schema(cls, exclude=["result", "purpose"])
        return schema
