import os

import anthropic
from agentlys.base import AgentlysBase
from agentlys.model import Message, MessagePart
from agentlys.providers.base_provider import BaseProvider
from agentlys.providers.utils import add_empty_function_result

AGENTLYS_HOST = os.getenv("AGENTLYS_HOST")


def part_to_anthropic_dict(part: MessagePart) -> dict:
    if part.type == "text":
        return {
            "type": "text",
            "text": part.content,
        }
    elif part.type == "image":
        if part.image is None:
            raise ValueError("Image part must have an image")
        return {
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": part.image.format,
                "data": part.image.to_base64(),
            },
        }
    elif part.type == "function_call":
        return {
            "type": "tool_use",
            "id": part.function_call_id,
            "name": part.function_call["name"],
            "input": part.function_call["arguments"],
        }
    elif part.type == "function_result_image":
        if part.image is None:
            raise ValueError("Function result image part must have an image")
        # Build content blocks: include text content if present, then image
        content_blocks = []
        if part.content:
            content_blocks.append(
                {
                    "type": "text",
                    "text": part.content,
                }
            )
        content_blocks.append(
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": part.image.format,
                    "data": part.image.to_base64(),
                },
            }
        )
        return {
            "type": "tool_result",
            "content": content_blocks,
            "tool_use_id": part.function_call_id,
        }
    elif part.type == "function_result":
        return {
            "type": "tool_result",
            "tool_use_id": part.function_call_id,
            "content": part.content,
        }
    elif part.type == "thinking":
        if part.is_redacted:
            return {
                "type": "redacted_thinking",
                "data": part.thinking_signature,
            }
        return {
            "type": "thinking",
            "thinking": part.thinking,
            "signature": part.thinking_signature,
        }
    elif part.type == "compaction":
        return {
            "type": "text",
            "text": f"[Previous conversation summary]\n{part.content}",
        }
    raise ValueError(f"Unknown part type: {part.type}")


def message_to_anthropic_dict(message: Message) -> dict:
    res = {
        "role": message.role if message.role in ["user", "assistant"] else "user",
        "content": [],
    }

    for part in message.parts:
        if part.type == "text" and (not part.content or not part.content.strip()):
            continue
        res["content"].append(part_to_anthropic_dict(part))

    return res


DEFAULT_MAX_TOKENS = int(os.getenv("ANTHROPIC_MAX_TOKENS", "10000"))


class AnthropicProvider(BaseProvider):
    def __init__(self, chat: AgentlysBase, model: str, max_tokens: int | None = None):
        self.model = model
        self.client = anthropic.AsyncAnthropic(
            base_url=AGENTLYS_HOST if AGENTLYS_HOST else "https://api.anthropic.com",
        )
        self.chat = chat
        self.max_tokens = DEFAULT_MAX_TOKENS if max_tokens is None else max_tokens

    @staticmethod
    def _strip_thinking_from_prior_turns(messages: list[dict]) -> list[dict]:
        """Strip thinking/redacted_thinking blocks from assistant messages.

        Anthropic's thinking signatures are cryptographic proofs tied to the
        model version that generated them.  When the model is upgraded, old
        signatures become invalid and the API returns 400.  The docs confirm
        it is safe to omit thinking blocks from prior assistant turns.

        The last assistant message's thinking is only preserved when followed
        by a tool_result (active tool-use loop), since in that case the
        thinking was just generated by the current model.  Otherwise it is
        also stripped (e.g. conversation resumed from DB after a model upgrade).
        """

        def _has_tool_result_after(idx: int) -> bool:
            """Check if the message after idx contains a tool_result."""
            next_idx = idx + 1
            if next_idx >= len(messages):
                return False
            next_content = messages[next_idx].get("content")
            if not isinstance(next_content, list):
                return False
            return any(
                block.get("type") == "tool_result" for block in next_content
            )

        # Find the index of the last assistant message
        last_assistant_idx = None
        for i in range(len(messages) - 1, -1, -1):
            if messages[i].get("role") == "assistant":
                last_assistant_idx = i
                break

        if last_assistant_idx is None:
            return messages

        result = []
        for i, msg in enumerate(messages):
            if (
                msg.get("role") == "assistant"
                and isinstance(msg.get("content"), list)
            ):
                # Preserve thinking only on the last assistant when a tool
                # loop is pending (next message is a tool_result).
                if i == last_assistant_idx and _has_tool_result_after(i):
                    result.append(msg)
                else:
                    filtered_content = [
                        block
                        for block in msg["content"]
                        if block.get("type") not in ("thinking", "redacted_thinking")
                    ]
                    result.append({**msg, "content": filtered_content})
            else:
                result.append(msg)
        return result

    def _prepare_request_params(self, **kwargs):
        """Prepare messages, tools, and kwargs for Anthropic API request."""
        messages = self.prepare_messages(
            transform_function=lambda m: message_to_anthropic_dict(m),
            transform_list_function=add_empty_function_result,
        )

        if self.chat.instruction:
            system = self.chat.instruction
        else:
            system = None

        def merge_messages(messages):
            """
            When two messages are in the same role, we merge the following message into the previous.
            {
                "role": "user",
                "content": [
                    {
                    "type": "tool_result",
                    "tool_use_id": "example_19",
                    "content": ""
                    }
                ]
            },
            {
                "role": "user",
                "content": "Plot distribution of stations per city"
            }
            """
            merged_messages = []
            for message in messages:
                if merged_messages and merged_messages[-1]["role"] == message["role"]:
                    # Convert string content to list format for merging
                    if isinstance(merged_messages[-1]["content"], str):
                        merged_messages[-1]["content"] = [
                            {
                                "type": "text",
                                "text": merged_messages[-1]["content"],
                            }
                        ]

                    if not isinstance(merged_messages[-1]["content"], list):
                        raise ValueError(
                            f"Invalid content type: {type(merged_messages[-1]['content'])}"
                        )

                    # Track existing tool_use_ids to prevent duplicates
                    existing_tool_use_ids = {
                        c.get("tool_use_id")
                        for c in merged_messages[-1]["content"]
                        if isinstance(c, dict) and c.get("type") == "tool_result"
                    }
                    # Only add content blocks that aren't duplicate tool_results
                    for content_block in message["content"]:
                        if (
                            isinstance(content_block, dict)
                            and content_block.get("type") == "tool_result"
                        ):
                            tool_use_id = content_block.get("tool_use_id")
                            if tool_use_id in existing_tool_use_ids:
                                continue
                            existing_tool_use_ids.add(tool_use_id)
                        merged_messages[-1]["content"].append(content_block)
                else:
                    merged_messages.append(message)
            return merged_messages

        messages = merge_messages(messages)
        messages = self._strip_thinking_from_prior_turns(messages)

        # Need to map field "parameters" to "input_schema"
        tools = [
            {
                "name": s["name"],
                "description": s["description"],
                "input_schema": s["parameters"],
            }
            for s in self.chat.functions_schema
        ]
        # Add description to the function is their description is empty
        for tool in tools:
            if not tool["description"]:
                tool["description"] = "No description provided"

        # === Add cache_controls ===
        # Anthropic's prompt caching uses up to 4 breakpoints per request.
        # We use all 4: system[-1], tools[-1], messages[-3], messages[-1].
        #
        # Why messages[-3]?  Each tool-loop iteration appends exactly 2
        # messages (1 assistant + 1 tool_result).  So messages[-3] in the
        # current call corresponds to messages[-1] from the *previous* call,
        # whose prefix was already cached.  By keeping a breakpoint there,
        # Anthropic can serve that prefix from cache_read instead of
        # re-caching the entire message history on every iteration.

        def _set_cache_control(msg_index):
            """Add cache_control to the last content block of messages[msg_index]."""
            msg = messages[msg_index]
            if (
                isinstance(msg["content"], list)
                and len(msg["content"]) > 0
                and isinstance(msg["content"][-1], dict)
            ):
                msg["content"][-1]["cache_control"] = {"type": "ephemeral"}
            elif isinstance(msg["content"], str):
                msg["content"] = [
                    {
                        "type": "text",
                        "text": msg["content"],
                        "cache_control": {"type": "ephemeral"},
                    }
                ]

        if messages:
            # Breakpoint on messages[-1]: caches the full conversation
            _set_cache_control(-1)

            # Breakpoint on messages[-3]: retains the previous iteration's
            # cache so Anthropic can read from it (cache_read) instead of
            # re-caching the entire prefix (cache_creation).
            if len(messages) >= 3:
                _set_cache_control(-3)
        # Tools: Add cache_control to the last tool function
        if tools:
            tools[-1]["cache_control"] = {"type": "ephemeral"}

        # System: add cache_control to the LAST system block so the entire
        # system section (instruction + tool states) is cached as a unit.
        # Anthropic uses cumulative hashes — placing cache_control on
        # system[0] while system[1] varies invalidates every downstream block.
        system_messages = []
        if system is not None:
            system_messages.append({"type": "text", "text": system})

        if self.chat.initial_tools_states:
            system_messages.append(
                {"type": "text", "text": self.chat.initial_tools_states}
            )

        if system_messages:
            system_messages[-1]["cache_control"] = {"type": "ephemeral"}

        # === End of cache_control ===

        if system_messages:
            kwargs["system"] = system_messages

        if self.chat.use_tools_only and "tool_choice" not in kwargs:
            kwargs["tool_choice"] = {"type": "any"}

        # Add thinking config if set at class level and not already in kwargs
        if getattr(self.chat, "thinking", None) and "thinking" not in kwargs:
            kwargs["thinking"] = self.chat.thinking

        return messages, tools, kwargs

    async def fetch_async(self, **kwargs):
        messages, tools, kwargs = self._prepare_request_params(**kwargs)

        res = await self.client.messages.create(
            model=self.model,
            messages=messages,
            tools=tools,
            max_tokens=self.max_tokens,
            **kwargs,
        )
        res_dict = res.to_dict()
        msg = Message.from_anthropic_dict(
            role=res_dict["role"],
            content=res_dict["content"],
        )
        msg.usage = res_dict.get("usage")
        return msg

    async def fetch_stream_async(self, **kwargs):
        """Stream response tokens from Anthropic.

        Yields text chunks as they arrive. Returns the final Message
        (with potential tool calls) after streaming completes.
        """
        messages, tools, kwargs = self._prepare_request_params(**kwargs)

        async with self.client.messages.stream(
            model=self.model,
            messages=messages,
            tools=tools if tools else anthropic.NOT_GIVEN,
            max_tokens=self.max_tokens,
            **kwargs,
        ) as stream:
            async for event in stream:
                if event.type == "content_block_delta":
                    if event.delta.type == "thinking_delta":
                        yield {"type": "thinking", "content": event.delta.thinking}
                    elif event.delta.type == "text_delta":
                        yield {"type": "text", "content": event.delta.text}

            # Get final message for tool handling
            response = await stream.get_final_message()
            res_dict = response.to_dict()
            final_message = Message.from_anthropic_dict(
                role=res_dict["role"],
                content=res_dict["content"],
            )
            final_message.usage = res_dict.get("usage")
            yield {"type": "message", "message": final_message}
