import json
from typing import Any

from ..types import BadAIJsonAnswer
from ..utils import ExtendedString
from ..message_types import Role


def remove_json_wrapper(input_string: str) -> str:
    input_string = str(input_string).strip()
    if input_string.startswith("```json") and input_string.endswith("```"):
        json_content = input_string[7:-3].strip()
        return json_content

    return input_string


class LLMResponse(ExtendedString):
    """
    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
    """

    def __new__(cls, string: str, attrs: dict = None):
        obj = ExtendedString.__new__(cls, string, attrs)
        # Same fields like in Msg
        setattr(obj, "role", Role.ASSISTANT)
        setattr(obj, "content", str(string))
        return obj

    def parse_json(
        self, raise_errors: bool = True, required_fields: list[str] = None
    ) -> list | dict | Any:
        try:
            res = json.loads(remove_json_wrapper(self.content))
            if required_fields:
                if not isinstance(res, dict):
                    raise BadAIJsonAnswer("Not an object")
                for field in required_fields:
                    if field not in res:
                        raise BadAIJsonAnswer(f'Missing field "{field}"')
            return res
        except json.decoder.JSONDecodeError as e:
            if raise_errors:
                raise BadAIJsonAnswer() from e
            return False
