"""Anthropic clients."""

import json
from collections.abc import Generator
from typing import Any

import boto3
from botocore.config import Config as BotoConfig
from dateutil import parser
from openai.types.chat import ChatCompletion, ChatCompletionChunk

from llmax.external_clients.exceptions import ProviderNotFoundError
from llmax.external_clients.universal_client.chat_completion_message import (
    ChatCompletionAssistantMessage,
    ChatCompletionSystemMessage,
    ChatCompletionUserMessage,
)
from llmax.external_clients.universal_client.client import UniversalClient
from llmax.messages import Messages
from llmax.models import Deployment, Model
from llmax.utils import logger

Client = Any


MAPPING_FINISH_REASON = {
    "end_turn": "stop",
    "max_tokens": "stop",
}


_boto3_clients: dict[tuple[str, str], Any] = {}


def client_creation_anthropic(
    aws_key: str,
    aws_secret_key: str,
    region_name: str,
) -> Any:
    """Create the anthropic client, reusing cached clients per (key, region)."""
    cache_key = (aws_key, region_name)
    if cache_key not in _boto3_clients:
        _boto3_clients[cache_key] = boto3.Session(
            aws_access_key_id=aws_key,
            aws_secret_access_key=aws_secret_key,
            region_name=region_name,
        ).client(
            "bedrock-runtime",
            config=BotoConfig(
                max_pool_connections=100,
                retries={"max_attempts": 3, "mode": "adaptive"},
            ),
        )
    return _boto3_clients[cache_key]


def anthropic_parsing(response: dict[str, Any]) -> ChatCompletion | None:
    """The parsing from anthropic to openAi."""
    try:
        data = json.loads(response["body"].read())
        input_tokens = response["ResponseMetadata"]["HTTPHeaders"][
            "x-amzn-bedrock-input-token-count"
        ]
        completion_tokens = response["ResponseMetadata"]["HTTPHeaders"][
            "x-amzn-bedrock-output-token-count"
        ]
        chat_completion = {
            "id": response["ResponseMetadata"]["RequestId"],
            "created": int(
                parser.parse(
                    response["ResponseMetadata"]["HTTPHeaders"]["date"],
                ).timestamp(),
            ),
            "model": data["model"],
            "usage": {
                "completion_tokens": completion_tokens,
                "prompt_tokens": input_tokens,
                "total_tokens": completion_tokens + input_tokens,
            },
            "choices": [
                {
                    "finish_reason": MAPPING_FINISH_REASON[data["stop_reason"]],
                    "index": i,
                    "message": {"content": e["text"], "role": data["role"]},
                }
                for i, e in enumerate(data["content"])
            ],
            "object": "chat.completion",
        }
        return ChatCompletion.model_validate(chat_completion)
    except Exception:
        return None


def anthropic_parsing_stream(  # noqa: C901, PLR0912, PLR0915
    response: dict[str, Any],
) -> Generator[ChatCompletionChunk, None, None] | None:
    """The parsing from anthropic to openAi in streaming mode."""
    try:
        request_id = response["ResponseMetadata"]["RequestId"]
        created = int(
            parser.parse(
                response["ResponseMetadata"]["HTTPHeaders"]["date"],
            ).timestamp(),
        )
    except Exception:
        return None

    _object = "chat.completion.chunk"
    model = None
    role = None
    prompt_tokens = None
    completion_tokens = None
    stop_reason = None

    # Track active tool calls
    current_tool_calls = {}
    active_tool_id = None

    for event in response["body"]:
        chunk = json.loads(event["chunk"]["bytes"])
        chunk_type = chunk.get("type")

        # Default values for this chunk
        content = None
        tool_calls = None

        if chunk_type == "message_start":
            model = chunk["message"]["model"]
            role = chunk["message"]["role"]
            prompt_tokens = chunk["message"].get("usage", {}).get("input_tokens")
            completion_tokens = chunk["message"].get("usage", {}).get("output_tokens")

        elif chunk_type == "content_block_start":
            block_index = chunk.get("index", 0)
            content_block = chunk.get("content_block", {})
            block_type = content_block.get("type")

            if block_type == "text":
                # Regular text content block
                pass

            elif block_type == "tool_use":
                # Start of a tool call
                tool_id = content_block.get("id")
                tool_name = content_block.get("name")
                active_tool_id = tool_id

                # Initialize the tool call
                current_tool_calls[tool_id] = {
                    "index": block_index,
                    "id": tool_id,
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "arguments": "",  # Will be filled with JSON
                    },
                }

                # Create initial tool call to send
                tool_calls = [current_tool_calls[tool_id]]

        elif chunk_type == "content_block_delta":
            delta = chunk.get("delta", {})
            delta_type = delta.get("type")

            if delta_type == "text_delta":
                # Regular text content
                content = delta.get("text", "")

            elif delta_type == "input_json_delta" and active_tool_id:
                # Tool JSON input update
                json_fragment = delta.get("partial_json", "")

                # Update the arguments in the current tool call
                if active_tool_id in current_tool_calls:
                    current_tool_id = active_tool_id
                    current_tool_calls[current_tool_id]["function"]["arguments"] = (
                        json_fragment
                    )

                    # Create tool call update to send
                    tool_calls = [
                        {
                            "index": current_tool_calls[current_tool_id]["index"],
                            "id": current_tool_id,
                            "function": {"arguments": json_fragment},
                        },
                    ]

        elif chunk_type == "content_block_stop":
            if active_tool_id:
                active_tool_id = None

        elif chunk_type == "message_delta":
            if "delta" in chunk and "stop_reason" in chunk["delta"]:
                stop_reason = MAPPING_FINISH_REASON.get(
                    chunk["delta"]["stop_reason"],
                    "stop",
                )

        elif chunk_type == "message_stop":
            prompt_tokens = chunk.get("amazon-bedrock-invocationMetrics", {}).get(
                "inputTokenCount",
            )
            completion_tokens = chunk.get("amazon-bedrock-invocationMetrics", {}).get(
                "outputTokenCount",
            )

        # Only yield if we have content or tool calls
        if content is not None or tool_calls is not None:
            try:
                # Prepare delta based on what we have
                delta = {}
                if content is not None:
                    delta["content"] = content
                if role is not None and chunk_type == "message_start":
                    delta["role"] = role
                if tool_calls is not None:
                    delta["tool_calls"] = tool_calls

                chat_completion_chunk = {
                    "id": request_id,
                    "created": created,
                    "model": model,
                    "usage": {
                        "completion_tokens": completion_tokens,
                        "prompt_tokens": prompt_tokens,
                        "total_tokens": completion_tokens + prompt_tokens
                        if completion_tokens and prompt_tokens
                        else None,
                    },
                    "choices": [
                        {
                            "delta": delta,
                            "index": 0,
                            "finish_reason": stop_reason,
                        },
                    ],
                    "object": _object,
                }
                yield ChatCompletionChunk.model_validate(chat_completion_chunk)
            except Exception as e:
                logger.error(f"Error processing chunk: {e}")
                return None
    return None


def completion_call_anthropic(  # noqa: C901, PLR0912
    client: Any,
    messages: Messages,
    model: Model,
    stream: bool = False,
    *args: Any,  # noqa: ARG001
    **kwargs: Any,
) -> ChatCompletion | None | Generator[ChatCompletionChunk, None, None]:
    """Anthropic call to make a completion."""
    system_message: list[str] = []
    to_remove: list[int] = []

    try:
        counter = 0
        for i, message in enumerate(messages):
            try:
                ChatCompletionAssistantMessage.model_validate(message)
                counter += 1
            except Exception:  # noqa: S110
                pass
            try:
                ChatCompletionSystemMessage.model_validate(message)
                system_message.append(str(message["content"]))
                to_remove.append(i)
                counter += 1
            except Exception:  # noqa: S110
                pass
            try:
                ChatCompletionUserMessage.model_validate(message)
                counter += 1
            except Exception:  # noqa: S110
                pass
        if counter < len(messages):
            logger.error("Incorrect message formatting")
    except Exception:
        return None

    body = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 10_000,
        "messages": [
            message for i, message in enumerate(messages) if i not in to_remove
        ],
    }
    if len(system_message) > 0:
        body.update({"system": " ".join(system_message)})

    if "tools" in kwargs:
        tools = kwargs["tools"]
        if tools:
            # Convert OpenAI-style tools to Anthropic format
            anthropic_tools = []
            for tool in tools:
                if tool["type"] == "function":
                    anthropic_tool = {
                        "name": tool["function"]["name"],
                        "description": tool["function"].get("description", ""),
                        "input_schema": tool["function"].get("parameters", {}),
                    }
                    anthropic_tools.append(anthropic_tool)

            if anthropic_tools:
                body.update({"tools": anthropic_tools})

    if "tool_choice" in kwargs:
        tool_choice = kwargs["tool_choice"]
        if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
            body.update({"tool_choice": tool_choice["function"]["name"]})
        elif tool_choice == "auto":
            body.update({"tool_choice": "auto"})

    if stream:
        response_stream = client.invoke_model_with_response_stream(
            modelId=model,
            body=json.dumps(body),
        )
        return anthropic_parsing_stream(response_stream)
    response = client.invoke_model(modelId=model, body=json.dumps(body))
    return anthropic_parsing(response)


def get_client(deployment: Deployment) -> Client:
    """Get a client for the given deployment."""
    match deployment.provider:
        case "aws-bedrock":
            return UniversalClient(
                client_creation=client_creation_anthropic,
                completion_call=completion_call_anthropic,
                deployment=deployment,
                aws_key=deployment.project_id,
                aws_secret_key=deployment.api_key,
                region_name=deployment.region,
            )
        case _:
            raise ProviderNotFoundError(deployment)
