from collections.abc import Iterator

from openai import AzureOpenAI
from openai.lib._pydantic import to_strict_json_schema

from yera.config2.dataclasses import LLMConfig
from yera.config2.keyring import DevKeyring
from yera.models.data_classes import Message
from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel


class AzureOpenAILLM(BaseLLMInterface):
    def __init__(self, config: LLMConfig, overrides, creds_map):
        kws = {**config.inference.model_dump(), **overrides}
        self.client = None
        self._endpoint_path = f"providers.{creds_map['endpoint']}"
        self._api_key_path = f"providers.{creds_map['api_key']}"

        self.deployment_name = config.model_id

        super().__init__(**kws)

    def start(self):
        endpoint = DevKeyring.get(self._endpoint_path)
        api_key = DevKeyring.get(self._api_key_path)

        self.client = AzureOpenAI(
            api_version="2024-12-01-preview",
            azure_endpoint=endpoint,
            api_key=api_key,
        )

    def stop(self):
        self.client.close()
        self.client = None

    def chat(
        self,
        messages: list[Message],
        **openai_kw,
    ) -> Iterator[str]:
        openai_messages = [
            {"role": msg.role, "content": msg.content} for msg in messages
        ]

        stream = self.client.chat.completions.create(
            model=self.deployment_name,
            messages=openai_messages,
            stream=True,
            **openai_kw,
        )

        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    def make_struct(
        self,
        messages: list[Message],
        cls: type[TBaseModel],
        **openai_kw,
    ) -> Iterator[str]:
        openai_messages = [
            {"role": msg.role, "content": msg.content} for msg in messages
        ]

        try:
            schema = to_strict_json_schema(cls)
            stream = self.client.chat.completions.create(
                model=self.deployment_name,
                messages=openai_messages,
                response_format={
                    "type": "json_schema",
                    "json_schema": {
                        "name": cls.__name__,
                        "schema": schema,
                        "strict": True,
                    },
                },
                stream=True,
                **openai_kw,
            )

            for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content

        except Exception as e:
            error_msg = str(e).lower()

            if any(
                phrase in error_msg
                for phrase in [
                    "does not support",
                    "not supported",
                    "not available",
                    "unsupported",
                ]
            ):
                # Fall back to tool calling
                yield from self._make_struct_via_tools(
                    openai_messages, cls, **openai_kw
                )
            else:
                # Re-raise other errors
                raise

    def _make_struct_via_tools(
        self,
        messages: list[dict],
        cls: type[TBaseModel],
        **openai_kw,
    ) -> Iterator[str]:
        schema = cls.model_json_schema()
        schema.pop("$defs", None)
        schema.pop("definitions", None)

        tool = {
            "type": "function",
            "function": {
                "name": "return_structured_data",
                "description": "Returns structured data matching the required schema",
                "parameters": schema,
            },
        }

        stream = self.client.chat.completions.create(
            model=self.deployment_name,
            messages=messages,
            tools=[tool],
            tool_choice={
                "type": "function",
                "function": {"name": "return_structured_data"},
            },
            stream=True,
            **openai_kw,
        )

        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.tool_calls:
                tool_call = chunk.choices[0].delta.tool_calls[0]
                if tool_call.function and tool_call.function.arguments:
                    yield tool_call.function.arguments
