from collections.abc import Iterator

import boto3

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


def _supports_struct_gen(model_id: str) -> bool:
    model_lower = model_id.lower()

    # Anthropic Claude 4.5 family
    claude_45_patterns = [
        "claude-opus-4-5",
        "claude-sonnet-4-5",
        "claude-haiku-4-5",
        "claude-opus-4.5",
        "claude-sonnet-4.5",
        "claude-haiku-4.5",
    ]
    if any(p in model_lower for p in claude_45_patterns):
        return True

    # Qwen3 family
    if "qwen3" in model_lower or "qwen.qwen3" in model_lower:
        return True

    # OpenAI GPT-OSS
    if "gpt-oss" in model_lower:
        return True

    # DeepSeek V3.1
    if "deepseek" in model_lower and ("v3" in model_lower or "v3.1" in model_lower):
        return True

    # Google Gemma 3
    if "gemma-3" in model_lower or "gemma3" in model_lower:
        return True

    # MiniMax M2
    if "minimax" in model_lower and "m2" in model_lower:
        return True

    # Mistral AI (newer models with structured output support)
    mistral_structured_patterns = [
        "ministral-3",
        "magistral",
        "voxtral",
        "mistral-large-3",
    ]
    if any(p in model_lower for p in mistral_structured_patterns):
        return True

    # Moonshot Kimi K2
    if "kimi" in model_lower and "k2" in model_lower:
        return True

    # NVIDIA Nemotron Nano v2
    return bool("nemotron-nano" in model_lower and "v2" in model_lower)


class AwsBedrockLLM(BaseLLMInterface):
    def __init__(self, config: LLMConfig, overrides: dict, creds_map: dict):
        kws = {**config.inference.model_dump(), **overrides}
        self.client = None
        self.config = config
        self.creds_map = creds_map
        self.native_struct = _supports_struct_gen(config.model_id)
        super().__init__(**kws)

    def start(self):
        access_key = DevKeyring.get("providers." + self.creds_map["access_key"])
        secret_key = DevKeyring.get("providers." + self.creds_map["secret_key"])
        region = DevKeyring.get("providers." + self.creds_map["region"])
        sess = boto3.Session(
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            region_name=region,
        )
        self.client = sess.client("bedrock-runtime")

    def stop(self):
        self.client = None

    @staticmethod
    def _map_messages(messages: list[Message]) -> tuple[list[dict], list[dict]]:
        bedrock_msgs = []
        system_prompts = []

        for msg in messages:
            if msg.role == "system":
                system_prompts.append({"text": msg.content})
            else:
                bedrock_msgs.append(
                    {"role": msg.role, "content": [{"text": msg.content}]}
                )

        return bedrock_msgs, system_prompts

    def chat(self, messages: list[Message], **bedrock_kw) -> Iterator[str]:
        bedrock_msgs, system_prompts = self._map_messages(messages)

        request_body = {
            "modelId": self.config.model_id,
            "messages": bedrock_msgs,
        }

        if system_prompts:
            request_body["system"] = system_prompts

        inference_config = bedrock_kw.pop("inferenceConfig", {})
        if inference_config:
            request_body["inferenceConfig"] = inference_config

        response = self.client.converse_stream(**request_body)

        for event in response["stream"]:
            if "contentBlockDelta" in event:
                delta = event["contentBlockDelta"]["delta"]
                if "text" in delta:
                    yield delta["text"]

    def make_struct(
        self, messages: list[Message], cls: type[TBaseModel], **bedrock_kw
    ) -> Iterator[str]:
        bedrock_msgs, system_prompts = self._map_messages(messages)

        schema = cls.model_json_schema()

        tool_schema = {
            "toolSpec": {
                "name": "return_structured_output",
                "description": f"Returns structured data conforming to the {cls.__name__} schema",
                "strict": True,
                "inputSchema": {"json": schema},
            }
        }

        request_body = {
            "modelId": self.config.model_id,
            "messages": bedrock_msgs,
            "toolConfig": {
                "tools": [tool_schema],
                "toolChoice": {"tool": {"name": "return_structured_output"}},
            },
        }

        if system_prompts:
            request_body["system"] = system_prompts

        inference_config = bedrock_kw.pop(
            "inferenceConfig", {"maxTokens": 4096 * 4, "temperature": 0.0}
        )
        if inference_config:
            request_body["inferenceConfig"] = inference_config

        response = self.client.converse_stream(**request_body)

        for event in response["stream"]:
            if "contentBlockDelta" in event:
                delta = event["contentBlockDelta"]["delta"]
                if "toolUse" in delta and "input" in delta["toolUse"]:
                    yield delta["toolUse"]["input"]
