from __future__ import annotations

import os
from typing import TYPE_CHECKING, ClassVar, Optional

from openai import AsyncOpenAI

from agents import (
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
)

if TYPE_CHECKING:
    from agents import ModelSettings


class BaseModelProvider(ModelProvider):
    _instance: ClassVar[Optional["BaseModelProvider"]] = None

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None, model: Optional[str] = None):
        self._api_key = api_key or os.getenv("OPENAI_API_KEY") or ""
        self._base_url = base_url or os.getenv("OPENAI_API_BASE_URL") or ""
        self._model = model or os.getenv("OPENAI_API_MODEL") or "gpt-4o"
        self._client: Optional[AsyncOpenAI] = None
        self._init_client()

    def _init_client(self) -> None:
        if self._api_key and self._base_url:
            self._client = AsyncOpenAI(base_url=self._base_url, api_key=self._api_key)

    @property
    def api_key(self) -> str:
        return self._api_key

    @property
    def base_url(self) -> str:
        return self._base_url

    @property
    def model_name(self) -> str:
        return self._model

    def get_model(self, model_name: str | None) -> Model:
        if self._client is None:
            raise ValueError("Model provider not properly initialized. Please check API key and base URL.")
        return OpenAIChatCompletionsModel(model=model_name or self._model, openai_client=self._client)

    def create_sub_model(self, model: Optional[str] = None, model_settings: Optional["ModelSettings"] = None) -> "SubModelProvider":
        return SubModelProvider(parent=self, model=model, model_settings=model_settings)

    @classmethod
    def get_default(cls) -> "BaseModelProvider":
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    @classmethod
    def set_default(cls, provider: "BaseModelProvider") -> None:
        cls._instance = provider


class SubModelProvider(ModelProvider):
    def __init__(
        self,
        parent: BaseModelProvider,
        model: Optional[str] = None,
        model_settings: Optional["ModelSettings"] = None,
    ):
        self._parent = parent
        self._model = model or parent.model_name
        self._model_settings = model_settings

    @property
    def model_name(self) -> str:
        return self._model

    @property
    def parent(self) -> BaseModelProvider:
        return self._parent

    def get_model(self, model_name: str | None) -> Model:
        base_model = self._parent.get_model(model_name or self._model)
        return base_model

    def create_sub_model(self, model: Optional[str] = None, model_settings: Optional["ModelSettings"] = None) -> "SubModelProvider":
        return SubModelProvider(parent=self._parent, model=model, model_settings=model_settings)
