import re
import os
import tiktoken
import openai
from openai import AsyncOpenAI

from .base import LLMProvider, ContextLengthExceededError, ProviderAuthError, ProviderConnectionError

class OpenAIProvider(LLMProvider):
    """LLM provider implementation for OpenAI models."""
    _client: AsyncOpenAI | None = None
    _encodings: dict = {}

    def _get_client(self) -> AsyncOpenAI:
        if OpenAIProvider._client is None:
            if not (api_key := os.environ.get('TELLMGRAMBOT_OPENAI_API_KEY')):
                raise ProviderAuthError("OpenAI API key not set. Provide TELLMGRAMBOT_OPENAI_API_KEY.")
            OpenAIProvider._client = AsyncOpenAI(api_key=api_key)
        return OpenAIProvider._client

    async def complete(self, model: str, messages: list[dict], tools: list[dict] | None = None) -> "str | ToolCall":
        from .base import ToolCall
        client = self._get_client()
        kwargs: dict = {'model': model, 'messages': messages}
        if tools:
            kwargs['tools'] = [
                {"type": "function", "function": {"name": t["name"], "description": t["description"], "parameters": t["parameters"]}}
                for t in tools
            ]
        try:
            response = await client.chat.completions.create(**kwargs)
            choice = response.choices[0]
            if choice.finish_reason == 'tool_calls' and choice.message.tool_calls:
                tc = choice.message.tool_calls[0]
                import json
                return ToolCall(
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments),
                    tool_call_id=tc.id,
                )
            return choice.message.content
        except openai.AuthenticationError as e:
            raise ProviderAuthError(f"OpenAI authentication failed: {e}") from e
        except openai.BadRequestError as e:
            if re.search(r'maximum context.+reduce the length', str(e)):
                raise ContextLengthExceededError() from e
            raise
        except openai.APIConnectionError as e:
            raise ProviderConnectionError(f"OpenAI connection error: {e}") from e

    def _get_encoding(self, model: str) -> tiktoken.Encoding:
        """
        Get or load the tiktoken encoding for the specified model.

        Attempts to retrieve the encoding directly from tiktoken. If that fails,
        falls back to a sensible default ('o200k_base' for GPT-4/5, 'cl100k_base' otherwise).

        Args:
            model: The model name (e.g., 'gpt-4o', 'gpt-4-turbo').

        Returns:
            The tiktoken encoding object for this model.
        """
        if model not in OpenAIProvider._encodings:
            try:
                OpenAIProvider._encodings[model] = tiktoken.encoding_for_model(model)
            except Exception:
                fallback = "o200k_base" if re.search("^gpt-[45]", model) else "cl100k_base"
                try:
                    OpenAIProvider._encodings[model] = tiktoken.get_encoding(fallback)
                except Exception:
                    OpenAIProvider._encodings[model] = tiktoken.get_encoding("cl100k_base")
                print(f"Warning: using {OpenAIProvider._encodings[model].name} encoding for model \"{model}\"")
        return OpenAIProvider._encodings[model]

    async def count_tokens(self, model: str, messages: list[dict]) -> int:
        """
        Count tokens in messages for an OpenAI model using tiktoken.

        Encodes each message field using the model's tiktoken encoding and applies
        per-message and per-name overhead as documented in the OpenAI cookbook.

        Args:
            model: The OpenAI model name (e.g., 'gpt-4o', 'gpt-4-turbo').
            messages: List of message dictionaries with 'role' and 'content' keys.

        Returns:
            The total token count for the messages, including overhead.

        Source: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
        """
        encoding = self._get_encoding(model)
        num_tokens = 0
        for message in messages:
            num_tokens += 3  # tokens_per_message
            for key, value in message.items():
                num_tokens += len(encoding.encode(value))
                if key == "name":
                    num_tokens += 1  # tokens_per_name
        num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
        return num_tokens
