"""LLMClient — reusable wrapper around the Anthropic Messages API with tool_use support."""

from __future__ import annotations

import os
import random
import time
from typing import Any, Literal

from anthropic import Anthropic, AnthropicBedrock, AnthropicVertex
from anthropic._exceptions import APIStatusError
from anthropic.types.beta import (
    BetaMessage,
    BetaMessageParam,
    BetaToolResultBlockParam,
)

from .config import config
from .llm_context_manager import LLMContextManager
from .task_logger import get_agent_logger
from .utils import (
    APIProvider,
    ToolResult,
    _inject_prompt_caching,
    _load_aws_bedrock_access,
    _make_api_tool_result,
    _response_to_params,
    get_model_name,
)


class LLMClient:
    """Wraps the Anthropic Messages API for a single agent role.

    Each agent gets its own ``LLMClient`` with independent conversation
    history, system prompt, and tool definitions.
    """

    def __init__(
        self,
        name: str,
        system_prompt: str,
        tools: list[dict[str, Any]],
        model: str,
        provider: APIProvider,
        temperature: float | None = None,
        top_p: float | None = None,
        only_n_most_recent_images: int | None = 10,
        api_key: str | None = None,
        betas: list[str] | None = None,
        effort: Literal["low", "medium", "high", "xhigh", "max"] = "high",
    ) -> None:
        self.betas = [
            "effort-2025-11-24",
            # "computer-use-2025-01-24",
            *(betas if betas else []),
        ]
        self.name = name
        self.effort = effort
        self.system_prompt = system_prompt
        self.tools = tools
        self.provider = provider
        self.model_name = get_model_name(provider, model)
        self.summarization_model = get_model_name(provider, config.summarization_model)
        self.logger = get_agent_logger(self.name)
        self.max_tokens = config.max_tokens

        self.sampling = {}
        if temperature is not None:
            self.sampling["temperature"] = temperature
        if top_p is not None:
            self.sampling["top_p"] = top_p

        self.only_n_most_recent_images = only_n_most_recent_images
        self.api_key = api_key

        self.messages: list[BetaMessageParam] = []
        self.usage = {
            "input_tokens": 0,
            "last_output_tokens": 0,
            "total_input_tokens": 0,
            "total_output_tokens": 0,
        }
        self.client = self._create_client(provider)
        self.context_manager = LLMContextManager(
            name=name,
            model=model,
            summarization_model=self.summarization_model,
            provider=provider,
            client=self.client,
            betas=self.betas,
            only_n_most_recent_images=only_n_most_recent_images,
        )

    def _create_client(self, provider) -> Any:
        if provider == APIProvider.ANTHROPIC:
            return Anthropic(
                api_key=self.api_key or os.environ.get("ANTHROPIC_API_KEY"),
                max_retries=4,
            )
        if provider == APIProvider.BEDROCK:
            bedrock_creds = _load_aws_bedrock_access()
            return AnthropicBedrock(**bedrock_creds, max_retries=4)
        if provider == APIProvider.VERTEX:
            return AnthropicVertex()
        raise ValueError(f"Unknown provider: {provider}")

    def _refresh_client(self) -> None:
        """Re-assume the IAM role and rebuild the client with fresh credentials."""
        self.logger.info("Refreshing Bedrock credentials (assumed role expired)")
        self.client = self._create_client(self.provider)
        self.context_manager.client = self.client

    @staticmethod
    def _is_credential_expiry(e: Exception) -> bool:
        expiry_indicators = (
            "expiredtoken",
            "expiredtokenexception",
            "security token included in the request is expired",
            "request signature we calculated does not match",
            "unrecognizedclientexception",
        )
        if not isinstance(e, APIStatusError):
            return False
        if e.status_code not in (401, 403):
            return False
        msg = (e.body.get("message", "") if isinstance(e.body, dict) else str(e.body) or "").lower()
        return any(ind in msg for ind in expiry_indicators)

    def add_user_message(
        self,
        text: str,
        images: list[str] | None = None,
    ) -> None:
        """Append a user message with optional base64 PNG images."""
        content: list[dict[str, Any]] = []
        for img_b64 in images or []:
            content.append(
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/png",
                        "data": img_b64,
                    },
                }
            )
        content.append({"type": "text", "text": text})
        self.messages.append({"role": "user", "content": content})

    def add_tool_results(
        self,
        results: list[tuple[str, ToolResult | dict]],
    ) -> None:
        """Append tool_result blocks for the given ``(tool_use_id, result)`` pairs."""
        content: list[BetaToolResultBlockParam] = []
        for tool_use_id, result in results:
            if isinstance(result, dict):
                content.append(_make_api_tool_result(result, tool_use_id))
            else:
                content.append(
                    _make_api_tool_result(
                        {
                            "output": result.output,
                            "error": result.error,
                            "base64_data": result.base64_data,
                            "media_type": result.media_type,
                            "system": result.system,
                        }
                        if result
                        else {},
                        tool_use_id,
                    )
                )
        self.messages.append({"role": "user", "content": content})

    # MARK: Generate

    def _prepare_messages(self) -> list[BetaMessageParam]:
        """Create a copy of messages with prompt caching injected (Anthropic only)."""
        messages = list(self.messages)
        if self.provider == APIProvider.ANTHROPIC:
            _inject_prompt_caching(messages)
        return messages

    def generate(self) -> BetaMessage:
        """Call the Anthropic API and return the raw ``BetaMessage``."""
        compacted = False
        if self.count_tokens() > 0.8 * config.max_input_token_threshold:
            self._handle_context_compaction()
            compacted = True

        messages = self._prepare_messages()

        system: list[dict[str, Any]] = [
            {"type": "text", "text": self.system_prompt},
        ]

        extra_body: dict[str, Any] = {"thinking": {"type": "adaptive"}}
        output_config: dict[str, Any] = {"effort": self.effort}
        last_error: Exception | None = None
        for attempt in range(5):
            try:
                response: BetaMessage = self.client.beta.messages.create(
                    max_tokens=self.max_tokens,
                    messages=messages,
                    model=self.model_name,
                    system=system,
                    tools=self.tools,
                    betas=self.betas,
                    extra_body=extra_body,
                    output_config=output_config,
                    **self.sampling,
                )
                self.update_usage(response)
                return response
            except Exception as e:
                last_error = e

                if self._is_credential_expiry(e):
                    self._refresh_client()
                    continue

                err_msg = str(e)
                if "413" in err_msg or "too large" in err_msg.lower() or "too long" in err_msg.lower():
                    if compacted:
                        self.context_manager.aggressive_compaction(self.messages, self.only_n_most_recent_images)
                    else:
                        self._handle_context_compaction()
                        compacted = True
                    messages = self._prepare_messages()

                wait_time = 2 ** (attempt + 1)
                time.sleep(wait_time + 2.0 * random.random())

        raise RuntimeError(f"[{self.name}] All API attempts exhausted. Last error: {last_error}")

    def generate_and_process(
        self,
    ) -> tuple[list[dict[str, Any]], list[str], list[str], str]:
        """Generate a response, store it in history, and parse it.

        Returns:
            ``(tool_calls, text_blocks, thinking_blocks, stop_reason)`` where:
            - *tool_calls* is a list of ``{"id", "name", "input"}`` dicts
            - *text_blocks* is a list of text strings
            - *thinking_blocks* is a list of thinking strings
            - *stop_reason* is the API stop reason
        """
        response = self.generate()
        response_params = _response_to_params(response)
        self.messages.append({"role": "assistant", "content": response_params})

        tool_calls: list[dict[str, Any]] = []
        text_blocks: list[str] = []
        thinking_blocks: list[str] = []

        for block in response_params:
            if block.get("type") == "tool_use":
                tool_calls.append(
                    {
                        "id": block["id"],
                        "name": block["name"],
                        "input": block.get("input", {}),
                    }
                )
            elif block.get("type") == "text" and block.get("text"):
                text_blocks.append(block["text"])
            elif block.get("type") == "thinking" and block.get("thinking"):
                thinking_blocks.append(block["thinking"])

        if not tool_calls:
            self.add_user_message(text="Please use your tools to proceed.")

        return tool_calls, text_blocks, thinking_blocks, response.stop_reason or ""

    # MARK: Token counting & context compaction

    def count_tokens(self) -> int:
        """Return the input token count using the Anthropic count_tokens API."""
        return self.context_manager.count_tokens(self.messages, self.system_prompt, self.tools)

    def _handle_context_compaction(self) -> None:
        self.logger.info("=== COMPACTION TRIGGERED ===")
        self.logger.info(f"Old message length: {len(self.messages)}")
        self.messages = self.context_manager.compact(
            self.messages,
            self.system_prompt,
            self.tools,
            config.max_input_token_threshold,
        )
        self.logger.info(f"New message length: {len(self.messages)}")
        self.logger.info("=== COMPACTION TRIGGERED ===")

    # MARK: Usage tracking

    def update_usage(self, response: BetaMessage) -> None:
        self.usage["input_tokens"] = response.usage.input_tokens
        self.usage["last_output_tokens"] = response.usage.output_tokens
        self.usage["total_input_tokens"] += response.usage.input_tokens
        self.usage["total_output_tokens"] += response.usage.output_tokens

    def print_usage(self) -> str:
        repr = f"[{self.name}] Usage: \n"
        repr += f"Input Tokens: {self.usage['total_input_tokens']}\n"
        repr += f"Output Tokens: {self.usage['total_output_tokens']}\n"
        return repr

    def reset(self) -> None:
        """Clear conversation history."""
        self.messages = []
        self.usage = {
            "input_tokens": 0,
            "last_output_tokens": 0,
            "total_input_tokens": 0,
            "total_output_tokens": 0,
        }
