"""Anthropic integration for AgentCost.

Wraps an Anthropic client to automatically track token usage and cost
from every message creation call. Zero overhead — reads usage from
the response object that's already returned.

Usage:
    from agentcost import AgentCost
    from agentcost.integrations.anthropic import track_anthropic
    import anthropic

    ac = AgentCost(api_key="ac_xxxxx")
    client = track_anthropic(anthropic.Anthropic(), ac, agent_id="my-agent")

    # Every call is now auto-tracked
    response = client.messages.create(
        model="claude-sonnet-4-5-20250929",
        max_tokens=1024,
        messages=[{"role": "user", "content": "Hello"}]
    )
"""

from __future__ import annotations

import time
from typing import Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from agentcost import AgentCost

from ..event import CostEvent
from ..tracker import _active_run


class _TrackedMessages:
    """Wraps client.messages to auto-track create() calls."""

    def __init__(self, original_messages: Any, ac: "AgentCost", agent_id: str):
        self._original = original_messages
        self._ac = ac
        self._agent_id = agent_id

    def create(self, **kwargs: Any) -> Any:
        start = time.time()
        response = self._original.create(**kwargs)
        duration_ms = int((time.time() - start) * 1000)

        model = getattr(response, "model", kwargs.get("model", "unknown"))
        usage = getattr(response, "usage", None)

        tokens_in = getattr(usage, "input_tokens", 0) if usage else 0
        tokens_out = getattr(usage, "output_tokens", 0) if usage else 0
        cost = self._ac.pricing.llm_cost(model, tokens_in, tokens_out)

        active = _active_run.get(None)
        effective_agent = active._agent_id if active else self._agent_id
        effective_run_id = active._run_id if active else None

        event = CostEvent(
            agent_id=effective_agent,
            run_id=effective_run_id,
            component_name=model,
            event_type="llm_call",
            cost_usd=cost,
            duration_ms=duration_ms,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
        )

        if active:
            active._events.append(event)

        self._ac._client.enqueue(event)
        return response

    def __getattr__(self, name: str) -> Any:
        return getattr(self._original, name)


class _TrackedAnthropicClient:
    """Wraps an Anthropic client to auto-track all message creation calls."""

    def __init__(self, client: Any, ac: "AgentCost", agent_id: str):
        self._original = client
        self.messages = _TrackedMessages(client.messages, ac, agent_id)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._original, name)


def track_anthropic(
    client: Any,
    ac: "AgentCost",
    agent_id: str = "anthropic-agent",
) -> Any:
    """Wrap an Anthropic client to automatically track all message calls.

    Args:
        client: An anthropic.Anthropic() or anthropic.AsyncAnthropic() instance.
        ac: Your AgentCost instance.
        agent_id: Agent name for tracking (default: "anthropic-agent").

    Returns:
        A wrapped client that behaves identically but tracks costs.
    """
    return _TrackedAnthropicClient(client, ac, agent_id)
