from typing import Any
from typing import TYPE_CHECKING

from ..types import BadAIAnswer, TPrompt
from ..utils import ExtendedString, ConvertableToMessage, extract_number
from ..message_types import Role, AssistantMsg

if TYPE_CHECKING:
    from ..mcp import MCPConnection, MCPAnswer


class DictFromLLMResponse(dict):
    llm_response: "LLMResponse"

    def from_llm_response(self, llm_response: "LLMResponse"):
        self.llm_response = llm_response
        return self

    async def to_mcp(self, mcp: "MCPConnection"):
        return await mcp.exec(self)


class LLMResponse(ExtendedString, ConvertableToMessage):
    """
    Response from the Large Language Model.

    If treated as a string, it returns the text generated by the LLM.

    Also, it contains all fields returned by the API accessible as an attributes.

    See fields returned by the OpenAI:

    - https://platform.openai.com/docs/api-reference/completions/object
    - https://platform.openai.com/docs/api-reference/chat/object
    """

    role: Role
    content: str
    prompt: TPrompt
    gen_duration: float

    def __new__(cls, string: str, attrs: dict = None):
        attrs = {
            **(attrs or {}),
            "role": Role.ASSISTANT,
            "content": str(string),
            "prompt": None,
            # generation duration in seconds (float), used in metrics
            "gen_duration": None,
        }
        obj = ExtendedString.__new__(cls, string, attrs)
        return obj

    def parse_json(
        self,
        raise_errors: bool = True,
        required_fields: list[str] = None,
        validator: callable = None,
    ) -> list | dict | float | int | str | DictFromLLMResponse:
        try:
            res = super().parse_json(raise_errors=True, required_fields=required_fields)
            if validator:
                try:
                    validator(res)
                except Exception as e:
                    raise BadAIAnswer(f"Language model response validation failed: {e}") from None
        except Exception as e:  # pylint: disable=W0718
            if hasattr(self, "_retry_callback"):
                res = self._retry_callback()
                if isinstance(res, DictFromLLMResponse):
                    return res
                return res.parse_json(raise_errors, required_fields, validator)
            if raise_errors:
                raise e
            res = False
        if isinstance(res, dict):
            res = DictFromLLMResponse(res)
            res.llm_response = self
        return res

    def parse_number(
        self,
        default=BadAIAnswer,
        position="last",
        dtype: type | str = float,
        rounding: bool = False,
    ) -> int | float | Any:
        return extract_number(self.content, default, position, dtype, rounding)

    def as_message(self) -> AssistantMsg:
        return self.as_assistant

    async def to_mcp(self, mcp: "MCPConnection") -> "MCPAnswer":
        return await mcp.exec(self)

    def is_tool_call(self):
        from .._env import env
        return self.parse_json(
            raise_errors=False,
            required_fields=[env().config.AI_SYNTAX_FUNCTION_NAME_FIELD],
        ) is not False
