from typing import TYPE_CHECKING, Callable, Any, Awaitable, Union, List
from os import PathLike

from .message_types import Msg

if TYPE_CHECKING:
    from .wrappers.prompt_wrapper import PromptWrapper  # noqa: F401
    from .wrappers.llm_response_wrapper import LLMResponse  # noqa: F401

TPrompt = Union[
    dict, Msg, str, "PromptWrapper", List[Union[dict, Msg, str, "PromptWrapper"]]
]
"""Type for prompt argument in LLM requests"""
TplFunctionType = Callable[[Union[PathLike[str], str], Any], str]
"""Function type for rendering prompt templates"""
LLMFunctionType = Callable[[TPrompt, Any], "LLMResponse"]
"""Function type for requesting LLM synchronously"""
LLMAsyncFunctionType = Callable[[TPrompt, Any], Awaitable["LLMResponse"]]
"""Function type for requesting LLM asynchronously"""


class BadAIAnswer(ValueError):
    """Unprocessable response generated by the LLM"""

    def __init__(self, message: str = None, details: str = None):
        self.message = str(message or "Unprocessable response generated by the LLM")
        self.details = details
        super().__init__(self.message + (f": {self.details}" if self.details else ""))

    def __str__(self):
        return self.message + (f": {self.details}" if self.details else "")


class BadAIJsonAnswer(BadAIAnswer):
    def __init__(
        self, message: str = "Invalid JSON generated by the LLM", details=None
    ):
        super().__init__(message, details)


class LLMContextLengthExceededError(BadAIAnswer):
    def __init__(
        self,
        message: str = None,
        actual_tokens: int | None = None,
        max_tokens: int | None = None,
        model: str | None = None,
    ):
        if message is None:
            message = "LLM context length exceeded" + (f" for model '{model}'" if model else "")
            if actual_tokens:
                message += f", actual tokens: {actual_tokens}"
            if max_tokens:
                message += f", max. tokens: {max_tokens}"
        super().__init__(message)
        self.actual_tokens = actual_tokens
        self.max_tokens = max_tokens
        self.model = model
