import re
import logging
from datetime import datetime
from typing import Any

from .utils import run_parallel, RETURN_EXCEPTION
from .wrappers.llm_response_wrapper import LLMResponse, DictFromLLMResponse
from .types import TPrompt, LLMContextLengthExceededError
from ._env import env


def convert_exception(e: Exception, model: str = None) -> Exception | None:
    """
    Convert LLM exceptions microcore-specific exceptions if possible.
    Args:
        e (Exception): Original exception
        model (str): LLM model name, used for better error messages
    Returns:
        Converted exception or None if no conversion is possible
    """
    if not isinstance(e, Exception):
        return None
    t, msg = f"{type(e).__module__}.{type(e).__name__}", str(e)
    max_tokens, actual_tokens = None, None
    if t == "openai.BadRequestError" and "context_length_exceeded" in msg:
        match = re.search(
            r"maximum context length is (\d+) tokens.*?resulted in (\d+) tokens",
            msg
        )
        if match:
            max_tokens = int(match.group(1))
            actual_tokens = int(match.group(2))
        return LLMContextLengthExceededError(
            actual_tokens=actual_tokens,
            max_tokens=max_tokens,
            model=model
        )
    if t == "anthropic.BadRequestError" and "prompt is too long:" in msg:
        if match := re.search(r"(\d+)\s+tokens\s+>\s+(\d+)\s+maximum", msg):
            max_tokens = int(match.group(2))
            actual_tokens = int(match.group(1))
        return LLMContextLengthExceededError(
            actual_tokens=actual_tokens,
            max_tokens=max_tokens,
            model=model
        )
    if (
        t == "google.api_core.exceptions.InvalidArgument"
        and "The input token count exceeds the maximum number of tokens allowed" in msg
    ):
        if match := re.search(
            r"The input token count exceeds the maximum number of tokens allowed (\d+)",
            msg
        ):
            max_tokens = int(match.group(1))
        return LLMContextLengthExceededError(
            actual_tokens=actual_tokens,
            max_tokens=max_tokens,
            model=model
        )
    return None


def llm(
    prompt: TPrompt,
    retries: int = 0,
    parse_json: bool | dict = False,
    **kwargs
) -> str | LLMResponse:
    """
    Request Large Language Model synchronously

    Args:
        prompt (str | Msg | dict | list[str | Msg | dict]): Text to send to LLM
        retries (int): Number of retries in case of error
        parse_json (bool|dict):
            If True, parses response as JSON,
            alternatively non-empty dict can be used as parse_json arguments
        **kwargs (dict): Parameters supported by the LLM API

            See parameters supported by the OpenAI:

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

            **Additional parameters:**

                - callback: callable - callback function
                to be called on each chunk of text,
                enables response streaming if supported by the LLM API
                - callbacks: list[callable] - collection of callbacks
                to be called on each chunk of text,
                enables response streaming if supported by the LLM API

    Returns:

        Text generated by the LLM as string
        with all fields returned by 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
    """
    [h(prompt, **kwargs) for h in env().llm_before_handlers]
    start = datetime.now()
    tries = retries + 1
    while tries > 0:
        try:
            tries -= 1
            response = env().llm_function(prompt, **kwargs)
            break
        except Exception as e:  # pylint: disable=W0718
            converted_exception = convert_exception(e)
            # If context length exceeded, or no tries left --> do not retry
            if tries == 0 or isinstance(converted_exception, LLMContextLengthExceededError):
                if converted_exception:
                    raise converted_exception from e
                raise e
            logging.error(f"LLM error: {e}")
            logging.info(f"Retrying... {tries} retries left")
            continue
    try:
        response.gen_duration = (datetime.now() - start).total_seconds()
        if not env().config.SAVE_MEMORY:
            response.prompt = prompt
    except AttributeError:
        ...
    [h(response) for h in env().llm_after_handlers]
    if tries > 0:
        retry_params = dict(**kwargs)
        retry_params["retries"] = tries - 1
        setattr(
            response,
            "_retry_callback",
            lambda: llm(prompt, **retry_params)
        )
    if parse_json:
        parsing_params = parse_json if isinstance(parse_json, dict) else {}
        return response.parse_json(**parsing_params)
    return response


async def allm(
    prompt: TPrompt,
    retries: int = 0,
    parse_json: bool | dict = False,
    **kwargs
) -> str | LLMResponse | DictFromLLMResponse:
    """
    Request Large Language Model asynchronously

    Args:
        prompt (str | Msg | dict | list[str | Msg | dict]): Text to send to LLM
        retries (int): Number of retries in case of error
        parse_json (bool|dict):
            If True, parses response as JSON,
            alternatively non-empty dict can be used as parse_json arguments
        **kwargs (dict): Parameters supported by the LLM API

            See parameters supported by the OpenAI:

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

            **Additional parameters:**

            - callback: callable - callback function
            to be called on each chunk of text,
            enables response streaming if supported by the LLM API
            - callbacks: list[callable] - collection of callbacks
            to be called on each chunk of text,
            enables response streaming if supported by the LLM API

            Note: async callbacks are supported only for async LLM API calls

    Returns:

        Text generated by the LLM as string
        with all fields returned by 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
    """
    [h(prompt, **kwargs) for h in env().llm_before_handlers]
    start = datetime.now()
    tries = retries + 1
    while tries > 0:
        try:
            tries -= 1
            response = await env().llm_async_function(prompt, **kwargs)
            break
        except Exception as e:  # pylint: disable=W0718
            converted_exception = convert_exception(e)
            # If context length exceeded, or no tries left --> do not retry
            if tries == 0 or isinstance(converted_exception, LLMContextLengthExceededError):
                if converted_exception:
                    raise converted_exception from e
                raise e
            logging.error(f"LLM error: {e}")
            logging.info(f"Retrying... {tries} retries left")
            continue
    try:
        response.gen_duration = (datetime.now() - start).total_seconds()
        if not env().config.SAVE_MEMORY:
            response.prompt = prompt
    except AttributeError:
        ...
    [h(response) for h in env().llm_after_handlers]
    if parse_json:
        try:
            parsing_params = parse_json if isinstance(parse_json, dict) else {}
            return response.parse_json(**parsing_params)
        except Exception as e:  # pylint: disable=W0718
            if tries > 0:
                logging.error(f"LLM error: {e}")
                logging.info(f"Retrying... {tries} retries left")
                return await allm(prompt, retries=tries - 1, parse_json=parse_json, **kwargs)
    return response


async def llm_parallel(
    prompts: list[TPrompt],
    max_concurrent_tasks: int = None,
    allow_failures: bool = False,
    return_on_failure: Any = RETURN_EXCEPTION,
    log_errors: bool = True,
    **kwargs
) -> list[str | LLMResponse]:
    """
    Execute multiple LLM requests in parallel

    Returns (list[LLMResponse | str]): a list of responses in the same order as the prompts
    """
    tasks = [allm(prompt, **kwargs) for prompt in prompts]

    if max_concurrent_tasks is None:
        max_concurrent_tasks = int(env().config.MAX_CONCURRENT_TASKS or 0)
    if not max_concurrent_tasks:
        max_concurrent_tasks = len(tasks)

    return await run_parallel(
        tasks,
        max_concurrent_tasks=max_concurrent_tasks,
        allow_failures=allow_failures,
        return_on_failure=return_on_failure,
        log_errors=log_errors,
    )
