import re
import time

from loguru import logger
from openai import OpenAI
from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel

__all__ = ['ModelResponse', 'ModelClient']


class ModelResponse(BaseModel):
    """Base class for the model response."""

    content: str

    time_to_first_token: float | None
    time_to_last_token: float | None

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._thinking = self._action = None

    def _parse_content(self) -> tuple[str, str]:
        """Parse the content to extract thinking and action parts.

        Returns:
            tuple[str, str]: A tuple containing the thinking part and the action part of the content.
        """
        logger.debug('Parsing model response content to extract thinking and action...')

        thinking, action = '', self.content

        # For example, the Qwen2.5-VL-7B-Instruct model returns this format.
        if '<answer>' in self.content:
            match = re.search(
                r'<think>(?P<thinking>.*?)</think>.*?<answer>(?P<action>.*?)</answer>', self.content, re.DOTALL
            )
            if match:
                thinking, action = match.group('thinking').strip(), match.group('action').strip()

        elif 'finish(message=' in self.content:
            parts = self.content.split('finish(message=', 1)
            thinking, action = parts[0].strip(), 'finish(message=' + parts[1]

        elif 'do(action=' in self.content:
            parts = self.content.split('do(action=', 1)
            thinking, action = parts[0].strip(), 'do(action=' + parts[1]

        return thinking, action

    @property
    def thinking(self) -> str:
        """Get the thinking part of the model response.

        Returns:
            str: The thinking part of the model response.
        """
        if self._thinking is None:
            self._thinking, self._action = self._parse_content()
        return self._thinking

    @property
    def action(self) -> str:
        """Get the action part of the model response.

        Returns:
            str: The action part of the model response.
        """
        if self._action is None:
            self._thinking, self._action = self._parse_content()
        return self._action


class ModelClient:
    """Model client for interacting with the OpenAI API."""

    def __init__(self, *, base_url: str, model: str, verbose: bool = True, api_key: str = '') -> None:
        """Initialize the client.

        Args:
            base_url (str): The base URL for the OpenAI API.
            model (str): The model to use for generating responses.
            verbose (bool, optional): Whether to print verbose stream output. Defaults to True.
            api_key (str, optional): The API key for authenticating with the OpenAI API. Defaults to ''.
        """
        logger.debug(f'Initialized ModelClient with model: {model}, base_url: {base_url}, verbose: {verbose}')

        self.model, self.verbose = model, verbose
        self.openai = OpenAI(base_url=base_url, api_key=api_key)

    def verbose_print(self, msg: str) -> None:
        """Print message if verbose is enabled.

        Args:
            msg (str): Message to print.
        """
        if self.verbose:
            print(msg, end='', flush=True)

    def chat(
        self,
        messages: list[ChatCompletionMessageParam],
        *,
        max_tokens: int = 3000,
        temperature: float = 0.0,
        top_p: float = 0.85,
        frequency_penalty: float = 0.2,
        extra_body: dict = None,
    ) -> ModelResponse:
        """Send a chat request to the model.

        Args:
            messages (list[ChatCompletionMessageParam]): The list of messages for the chat.

            max_tokens (int, optional): The maximum number of tokens to generate. Defaults to 3000.
            temperature (float, optional): The sampling temperature. Defaults to 0.0.
            top_p (float, optional): The nucleus sampling probability. Defaults to 0.85.
            frequency_penalty (float, optional): The frequency penalty. Defaults to 0.2.
            extra_body (dict, optional): Additional parameters to include in the request body. Defaults to {}.

        Returns:
            ModelResponse: The response from the model.
        """
        start_time = time.perf_counter()
        logger.info('Requesting chat streaming response from model...')

        stream = self.openai.chat.completions.create(
            messages=messages,
            model=self.model,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            frequency_penalty=frequency_penalty,
            extra_body=extra_body,
            stream=True,
        )

        full_content, time_to_first_token = '', None

        for chunk in stream:
            if len(chunk.choices) == 0:
                continue
            elif not (content := chunk.choices[0].delta.content):
                continue

            if time_to_first_token is None:
                time_to_first_token = round(time.perf_counter() - start_time, 4)

            full_content += content
            self.verbose_print(msg=content)

        self.verbose_print(msg='\n')
        time_to_last_token = round(time.perf_counter() - start_time, 4)

        logger.info(
            'Completed processing streaming response.'
            f'Time to first token: {time_to_first_token}s, time to last token: {time_to_last_token}s.'
        )

        return ModelResponse(
            content=full_content,
            time_to_first_token=time_to_first_token,
            time_to_last_token=time_to_last_token,
        )
