"""
Agent Status integration for Anthropic Python client.

Provides automatic monitoring for Anthropic API calls (messages, etc).

Usage:

    from agent_status.integrations.anthropic_wrapper import AgentStatusAnthropic

    client = AgentStatusAnthropic(rora_api_key="fb_live_xxx")

    message = client.messages.create(
        model="claude-3-opus-20240229",
        max_tokens=1024,
        messages=[{"role": "user", "content": "Hello!"}]
    )

    print(client.rora_metrics)
"""

from __future__ import annotations

import time
import logging
from datetime import datetime, timezone
from typing import Any, Dict, Optional

ANTHROPIC_AVAILABLE = False

try:
    from anthropic import Anthropic, AsyncAnthropic
    ANTHROPIC_AVAILABLE = True
except ImportError:
    Anthropic = object
    AsyncAnthropic = object

from agent_status.client import AgentStatusClient

logger = logging.getLogger(__name__)


class TrackedMessages:
    """Wrapper for messages that tracks metrics."""

    def __init__(self, original_messages: Any, metrics: Dict[str, Any]):
        self._original = original_messages
        self._metrics = metrics

    def create(self, *args, **kwargs) -> Any:
        """Create a message with tracking."""
        start = time.time()
        model = kwargs.get("model", "unknown")

        try:
            result = self._original.create(*args, **kwargs)

            latency_ms = (time.time() - start) * 1000
            self._metrics["total_requests"] += 1
            self._metrics["total_latency_ms"] += latency_ms
            self._metrics["successful_requests"] += 1
            self._metrics["by_model"].setdefault(model, {"requests": 0, "latency_ms": 0})
            self._metrics["by_model"][model]["requests"] += 1
            self._metrics["by_model"][model]["latency_ms"] += latency_ms

            if hasattr(result, "usage") and result.usage:
                self._metrics["total_tokens"] += (
                    result.usage.input_tokens + result.usage.output_tokens
                )
                self._metrics["input_tokens"] += result.usage.input_tokens
                self._metrics["output_tokens"] += result.usage.output_tokens

            logger.debug(f"Anthropic request completed: model={model}, latency={latency_ms:.0f}ms")
            return result

        except Exception as e:
            self._metrics["total_requests"] += 1
            self._metrics["failed_requests"] += 1
            self._metrics["errors"].append({
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "model": model,
                "error": str(e),
            })
            logger.error(f"Anthropic request failed: {e}")
            raise


class AgentStatusAnthropic:
    """
    Anthropic client wrapper that automatically tracks metrics for Agent Status.

    Drop-in replacement for the Anthropic client that tracks:
    - Request counts and latencies
    - Token usage (input, output, total)
    - Errors and failures
    - Per-model breakdown

    Args:
        rora_api_key: Your API key
        rora_base_url: API base URL
        **kwargs: All other arguments passed to Anthropic client

    Example:
        from agent_status.integrations.anthropic_wrapper import AgentStatusAnthropic

        client = AgentStatusAnthropic(rora_api_key="fb_live_xxx")

        message = client.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hello!"}]
        )

        print(client.rora_metrics)
    """

    def __init__(
        self,
        rora_api_key: str,
        rora_base_url: str = "https://api.fabric.carmel.so",
        **kwargs,
    ):
        if not ANTHROPIC_AVAILABLE:
            raise ImportError(
                "Anthropic is not installed. Install with: pip install anthropic"
            )

        self._anthropic = Anthropic(**kwargs)
        self._rora_client = AgentStatusClient(api_key=rora_api_key, base_url=rora_base_url)

        self._metrics: Dict[str, Any] = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_latency_ms": 0,
            "total_tokens": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "by_model": {},
            "errors": [],
        }

        self.messages = TrackedMessages(self._anthropic.messages, self._metrics)

    @property
    def rora_metrics(self) -> Dict[str, Any]:
        """Get current metrics summary."""
        total = self._metrics["total_requests"]
        return {
            **self._metrics,
            "avg_latency_ms": (
                self._metrics["total_latency_ms"] / total if total > 0 else 0
            ),
            "success_rate": (
                self._metrics["successful_requests"] / total if total > 0 else 0
            ),
        }

    def reset_metrics(self) -> None:
        """Reset all tracked metrics."""
        self._metrics = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_latency_ms": 0,
            "total_tokens": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "by_model": {},
            "errors": [],
        }
        self.messages = TrackedMessages(self._anthropic.messages, self._metrics)

    def report(self) -> Dict[str, Any]:
        """Generate a summary report."""
        return {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "metrics": self.rora_metrics,
        }

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


class AsyncTrackedMessages:
    """Async wrapper for messages that tracks metrics."""

    def __init__(self, original_messages: Any, metrics: Dict[str, Any]):
        self._original = original_messages
        self._metrics = metrics

    async def create(self, *args, **kwargs) -> Any:
        """Create a message with tracking."""
        start = time.time()
        model = kwargs.get("model", "unknown")

        try:
            result = await self._original.create(*args, **kwargs)

            latency_ms = (time.time() - start) * 1000
            self._metrics["total_requests"] += 1
            self._metrics["total_latency_ms"] += latency_ms
            self._metrics["successful_requests"] += 1
            self._metrics["by_model"].setdefault(model, {"requests": 0, "latency_ms": 0})
            self._metrics["by_model"][model]["requests"] += 1
            self._metrics["by_model"][model]["latency_ms"] += latency_ms

            if hasattr(result, "usage") and result.usage:
                self._metrics["total_tokens"] += (
                    result.usage.input_tokens + result.usage.output_tokens
                )
                self._metrics["input_tokens"] += result.usage.input_tokens
                self._metrics["output_tokens"] += result.usage.output_tokens

            return result

        except Exception as e:
            self._metrics["total_requests"] += 1
            self._metrics["failed_requests"] += 1
            self._metrics["errors"].append({
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "model": model,
                "error": str(e),
            })
            raise


class AgentStatusAsyncAnthropic:
    """
    Async Anthropic client wrapper that automatically tracks metrics for Agent Status.

    Drop-in replacement for the AsyncAnthropic client.

    Example:
        from agent_status.integrations.anthropic_wrapper import AgentStatusAsyncAnthropic

        client = AgentStatusAsyncAnthropic(rora_api_key="fb_live_xxx")

        message = await client.messages.create(
            model="claude-3-opus-20240229",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hello!"}]
        )

        print(client.rora_metrics)
    """

    def __init__(
        self,
        rora_api_key: str,
        rora_base_url: str = "https://api.fabric.carmel.so",
        **kwargs,
    ):
        if not ANTHROPIC_AVAILABLE:
            raise ImportError(
                "Anthropic is not installed. Install with: pip install anthropic"
            )

        self._anthropic = AsyncAnthropic(**kwargs)
        self._rora_client = AgentStatusClient(api_key=rora_api_key, base_url=rora_base_url)

        self._metrics: Dict[str, Any] = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_latency_ms": 0,
            "total_tokens": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "by_model": {},
            "errors": [],
        }

        self.messages = AsyncTrackedMessages(self._anthropic.messages, self._metrics)

    @property
    def rora_metrics(self) -> Dict[str, Any]:
        """Get current metrics summary."""
        total = self._metrics["total_requests"]
        return {
            **self._metrics,
            "avg_latency_ms": (
                self._metrics["total_latency_ms"] / total if total > 0 else 0
            ),
            "success_rate": (
                self._metrics["successful_requests"] / total if total > 0 else 0
            ),
        }

    def reset_metrics(self) -> None:
        """Reset all tracked metrics."""
        self._metrics = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "total_latency_ms": 0,
            "total_tokens": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "by_model": {},
            "errors": [],
        }
        self.messages = AsyncTrackedMessages(self._anthropic.messages, self._metrics)

    def report(self) -> Dict[str, Any]:
        """Generate a summary report."""
        return {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "metrics": self.rora_metrics,
        }

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


# Backward compatibility aliases
RoraAnthropic = AgentStatusAnthropic
RoraAsyncAnthropic = AgentStatusAsyncAnthropic
