import logging
import typing
from abc import ABC, ABCMeta, abstractmethod
from dataclasses import asdict, dataclass, field
from typing import Any, Generic, List, Optional, TypeAlias

from anthropic.types import Message as AnthropicMessage
from anthropic.types import TextBlock, ToolUseBlock

from src.agents.tools.base_tool import AnthropicTool

log = logging.getLogger(__name__)

if typing.TYPE_CHECKING:
    from src.agents.basic_agent import LLMAgent


@dataclass(kw_only=True)
class AbstractLLMMessage(ABC):
    creator: 'LLMAgent'

    # generated after object initialization
    # multiple messages can be generated by a single API call responses

    generation_batch_metadata: "LLMMessage_Generation_Metadata"

    @classmethod
    @abstractmethod
    def from_anthropic_message(cls, anthropic_message: AnthropicMessage,
                               creator: 'LLMAgent'):
        """An anthropic message can take two forms:\n
        if no tooluse -> message.content = [TextBlock]\n
        if tooluse    -> message.content = [*TextBlock, *ToolUseBlock]
        """
        pass


@dataclass(kw_only=True)
class LLMMessage(AbstractLLMMessage):
    content: str

    @classmethod
    def from_anthropic_message(cls, anthropic_message: AnthropicMessage,
                               creator: 'LLMAgent'):
        metadata = LLMMessage_Generation_Metadata.from_anthropic_message(
            anthropic_message)

        result: list["LLMMessage"] = []

        for i in anthropic_message.content:
            if isinstance(i, ToolUseBlock):
                # skip tool use responses

                continue
            result.append(
                cls(creator=creator,
                    content=i.text,
                    generation_batch_metadata=metadata))
        return result


@dataclass
class LLMToolUseDecision(AbstractLLMMessage):
    """An agent's choice of which tool to use when it is provided with a list of tools and a task"""

    context: (
        str  # the agent's explanation and rationale of why it chose to use the tool
    )
    tool_name: str
    tool_kwargs: dict  # the keyword arguments to the tool

    @classmethod
    def decided_to_use_tool(cls, anthropic_message: AnthropicMessage):
        return any(
            isinstance(i, ToolUseBlock) for i in anthropic_message.content)

    @classmethod
    def from_anthropic_message(cls, anthropic_message: AnthropicMessage,
                               creator: 'LLMAgent'):
        """Takes an LLM's response (an anthropic message) and converts it into a list of tools the LLM decided to use.
        
        The parsing is done like so:
        
        All text responses before a tool use block is folded into the rationale behind a tool use.
        
        Multiple tool uses can result from a single LLM response.
        """
        metadata = LLMMessage_Generation_Metadata.from_anthropic_message(
            anthropic_message)

        result: list[LLMToolUseDecision] = []

        context = ""
        for i in anthropic_message.content:
            if isinstance(i, TextBlock):
                context += i.text
            elif isinstance(i, ToolUseBlock):

                if not isinstance(i.input, dict):
                    log.warning(f"ToolUseBlock input is not a dict {i.input}")
                    continue
                result.append(
                    cls(
                        creator=creator,
                        generation_batch_metadata=metadata,
                        context=context,
                        tool_name=i.name,
                        tool_kwargs=i.input,
                    ))
                context = ""
        return result

    def format_as_message(self):
        """Formats the decision as a string that details what tool was used, and why."""
        s = '\n'
        s += 'Tool use thought process:\n'
        s += f'Tool name: {self.tool_name}\n'
        s += f'Input to the tool: {str(self.tool_kwargs)}\n'
        s += f'Context: {self.context}'
        return s


@dataclass
class ToolUsedOrFailed(ABC):
    tool_use_decision: LLMToolUseDecision | None
    tool: "AnthropicTool"


@dataclass
class ToolNotInvoked(ToolUsedOrFailed):
    """When a agent was unable to use a tool for any reason."""

    reason: str | None = None
    tool: Optional["AnthropicTool"]

    def __str__(self) -> str:
        return f"Tool Not Invoked. {self.tool}, Reason: {self.reason}"


@dataclass
class ToolInvoked(ToolUsedOrFailed):
    """The result of a successful tool use by an agent."""

    result: Any

    def __str__(self) -> str:
        return f"Tool Invoked. {self.tool}, Result:\n```\n{self.result}\n```\n"


class NoToolsAvailable:
    """When an Agent was asked to choose a tool to use, but did not have any tools at 
    its disposal"""

    @classmethod
    def to_dict(cls):
        return {"class_name": cls.__name__}

    @classmethod
    def from_dict(cls, d):
        if d["class_name"] != cls.__name__:
            return


@dataclass(kw_only=True)
class LLMMessage_Generation_Metadata:
    message_id: str
    model: str
    input_tokens: int
    output_tokens: int
    stop_reason: Optional[str] = None
    price_in_microcents: Optional[int] = None

    @classmethod
    def from_anthropic_message(cls, message: AnthropicMessage):
        price = cls.calculate_price_of_generation(
            model=message.model,
            input_tokens=message.usage.input_tokens,
            output_tokens=message.usage.output_tokens)

        return cls(message_id=message.id,
                   model=message.model,
                   stop_reason=message.stop_reason,
                   input_tokens=message.usage.input_tokens,
                   output_tokens=message.usage.output_tokens,
                   price_in_microcents=price)

    @classmethod
    def calculate_price_of_generation(cls, model, input_tokens, output_tokens):
        # token caching not implemented yet

        model_name_for_pricing = next(
            (model for model in anthropic_pricing_cents_per_million_tokens
             if model in model),
            None,
        )

        if not model_name_for_pricing:
            log.warning(
                f"Could not find model name in pricing list: {model =}")
            return
        price_dict = anthropic_pricing_cents_per_million_tokens[
            model_name_for_pricing]
        price_in_microcents = (input_tokens * price_dict["input"] +
                               output_tokens * price_dict["output"])
        return price_in_microcents

    def get_price_in_dollars(self):
        if self.price_in_microcents:
            return self.price_in_microcents / 1_000_000


anthropic_pricing_cents_per_million_tokens = {
    # hard coded because not exposed on API
    # last updated: 2024-9-29
    "claude-3-5-sonnet": {
        "input": 300,
        "cached_input_read": 30,
        "cached_input_write": 375,
        "output": 1500,
    },
    "claude-3-opus": {
        "input": 1500,
        "cached_input_read": 150,
        "cached_input_write": 1875,
        "output": 7500,
    },
    "claude-3-haiku": {
        "input": 25,
        "cached_input_read": 3,
        "cached_input_write": 30,
        "output": 125,
    },
}
