import re
from collections.abc import Iterator

from anthropic import Anthropic, transform_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 AnthropicLLM(BaseLLMInterface):
    def __init__(self, config: LLMConfig, overrides, creds_map):
        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:
        pattern = r"claude-(?:sonnet|opus|haiku)-(\d+)-(\d+)"
        match = re.search(pattern, model_id)

        if not match:
            return False

        major = int(match.group(1))
        minor = int(match.group(2))

        # Structured outputs supported on 4.5+
        return (major, minor) >= (4, 5)

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

    def stop(self):
        self.client = None

    def chat(
        self, messages: list[Message], *, max_tokens: int = 4096, **anthropic_kw
    ) -> Iterator[str]:
        # todo: claude does not accept system lines, but one a single system param
        #   this means sys lines in the middle of the conversation have to be converted
        #   into user lines. Content: f"<instruction>\n{msg['content']}\n</instruction>"

        system_lines = [msg.content for msg in messages if msg.role == "system"]
        messages = [
            {"role": msg.role, "content": msg.content}
            for msg in messages
            if msg.role != "system"
        ]

        system = "\n\n".join(system_lines) if system_lines else None

        kws = {
            "model": self.model_id,
            "max_tokens": max_tokens,
            "messages": messages,
            **anthropic_kw,
            "thinking": {"type": "enabled", "budget_tokens": 1024},
        }
        if system:
            kws["system"] = system

        with self.client.messages.stream(
            **kws,
        ) as stream:
            for event in stream:
                if event.type == "content_block_start":
                    if event.content_block.type == "thinking":
                        yield "<thinking>"
                    elif event.content_block.type == "text":
                        yield "\n<thinking>\n\n"

                elif event.type == "content_block_delta":
                    if event.delta.type == "thinking_delta":
                        yield event.delta.thinking
                    elif event.delta.type == "text_delta":
                        yield event.delta.text

    def make_struct(
        self,
        messages: list[Message],
        cls: type[TBaseModel],
        *,
        max_tokens: int = 4096,
        **anthropic_kw,
    ) -> Iterator[str]:
        system_lines = [msg.content for msg in messages if msg.role == "system"]
        conversation_messages = [
            {"role": msg.role, "content": msg.content}
            for msg in messages
            if msg.role != "system"
        ]

        system = "\n\n".join(system_lines) if system_lines else None

        if self.struct_gen:
            kws = {
                "model": self.model_id,
                "max_tokens": max_tokens,
                "messages": conversation_messages,
                "betas": ["structured-outputs-2025-11-13"],
                "output_format": {
                    "type": "json_schema",
                    "schema": transform_schema(cls),
                },
                **anthropic_kw,
            }
            if system:
                kws["system"] = system

            with self.client.beta.messages.stream(**kws) as stream:
                yield from stream.text_stream
        else:
            tool_schema = {
                "name": "return_structured_output",
                "description": f"Returns structured data conforming to the {cls.__name__} schema",
                "input_schema": cls.model_json_schema(),
            }

            kws = {
                "model": self.model_id,
                "max_tokens": max_tokens,
                "messages": conversation_messages,
                "tools": [tool_schema],
                "tool_choice": {"type": "tool", "name": "return_structured_output"},
                **anthropic_kw,
            }
            if system:
                kws["system"] = system

            tool_use_started = False

            with self.client.messages.stream(**kws) as stream:
                for event in stream:
                    if event.type == "content_block_start":
                        if (
                            hasattr(event.content_block, "type")
                            and event.content_block.type == "tool_use"
                        ):
                            tool_use_started = True  # We got the tool call!

                    elif event.type == "content_block_delta":
                        delta = event.delta
                        if hasattr(delta, "type") and delta.type == "input_json_delta":
                            yield delta.partial_json

                if not tool_use_started:
                    raise ValueError("Model did not call the structured output tool")
