from typing import Any

import openai
from openai.types.chat import (
    ChatCompletion,
    ChatCompletionMessageParam,
    ChatCompletionToolChoiceOptionParam,
    ChatCompletionToolParam,
)


class LlmClient:
    def __init__(self, base_url: str, api_key: str, timeout: float) -> None:
        self._base_url = base_url
        self._api_key = api_key
        self._timeout = timeout
        self._client: openai.OpenAI | None = None

    def open(self) -> "LlmClient":
        self._client = openai.OpenAI(
            base_url=self._base_url,
            api_key=self._api_key,
            timeout=self._timeout,
        )
        return self

    def close(self) -> None:
        if self._client:
            self._client.close()
            self._client = None

    def __enter__(self) -> "LlmClient":
        return self.open()

    def __exit__(self, *_: object) -> None:
        self.close()

    def create_message(
        self,
        *,
        model: str,
        max_tokens: int,
        system: str,
        messages: list[ChatCompletionMessageParam],
        tools: list[ChatCompletionToolParam] | None = None,
        tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
    ) -> ChatCompletion:
        if self._client is None:
            raise RuntimeError("Client not initialized — call open() or use as context manager")
        msgs: list[ChatCompletionMessageParam] = [{"role": "system", "content": system}, *messages]
        kwargs: dict[str, Any] = {
            "model": model,
            "max_tokens": max_tokens,
            "messages": msgs,
        }
        if tools:
            kwargs["tools"] = tools
        if tool_choice is not None:
            kwargs["tool_choice"] = tool_choice
        return self._client.chat.completions.create(**kwargs)  # type: ignore[no-any-return]
