import re
from collections.abc import Iterator

from openai import OpenAI
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 OpenAILLM(BaseLLMInterface):
    def __init__(self, config: LLMConfig, overrides: dict, creds_map: dict[str, str]):
        kws = {**config.inference.model_dump(), **overrides}
        self.client = None
        self.model_id = config.model_id
        self.struct_gen = self._supports_structured_outputs(self.model_id)
        self.creds_map = creds_map
        super().__init__(**kws)

    @staticmethod
    def _supports_structured_outputs(model_id: str) -> bool:
        """Check if a model supports structured outputs via response_format.

        Structured outputs via response_format are supported on:
        - gpt-4o models from 2024-08-06 onwards
        - gpt-4o-mini models from 2024-07-18 onwards
        - gpt-5+ and o3/o4 models

        Older models can use tool calling with strict: true instead.

        Args:
            model_id: Model identifier (e.g., "gpt-4o-2024-08-06")

        Returns:
            True if model supports response_format structured outputs
        """
        # GPT-5+ models support structured outputs
        if (
            model_id.startswith("gpt-5")
            or model_id.startswith("o3")
            or model_id.startswith("o4")
        ):
            return True

        # Check gpt-4o and gpt-4o-mini with date snapshots
        gpt4o_pattern = r"gpt-4o(?:-mini)?-(\d{4})-(\d{2})-(\d{2})"
        match = re.search(gpt4o_pattern, model_id)

        if match:
            year = int(match.group(1))
            month = int(match.group(2))
            day = int(match.group(3))
            date_tuple = (year, month, day)

            # gpt-4o-mini: supported from 2024-07-18+
            if "mini" in model_id:
                return date_tuple >= (2024, 7, 18)
            # gpt-4o: supported from 2024-08-06+
            return date_tuple >= (2024, 8, 6)

        return False

    def start(self):
        api_key = DevKeyring.get("providers." + self.creds_map["api_key"])
        self.client = OpenAI(api_key=api_key)

    def stop(self):
        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.model_id,
            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
        ]

        schema = to_strict_json_schema(cls)

        if self.struct_gen:
            stream = self.client.chat.completions.create(
                model=self.model_id,
                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

        else:
            tool_schema = {
                "type": "function",
                "function": {
                    "name": "return_structured_output",
                    "description": f"Returns structured data conforming to the {cls.__name__} schema",
                    "parameters": schema,
                    "strict": True,
                },
            }

            stream = self.client.chat.completions.create(
                model=self.model_id,
                messages=openai_messages,
                tools=[tool_schema],
                tool_choice={
                    "type": "function",
                    "function": {"name": "return_structured_output"},
                },
                stream=True,
                **openai_kw,
            )

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