"""
Agent Status integration for OpenAI Python client.

Provides automatic monitoring for OpenAI API calls (chat completions, etc).

Usage:

    from openai import OpenAI
    from agent_status.integrations.openai_wrapper import AgentStatusOpenAI

    client = AgentStatusOpenAI(api_key="fb_live_xxx")

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "Hello!"}]
    )

    print(client.rora_metrics)
"""

from __future__ import annotations

import time
import logging
from datetime import datetime, timezone
from functools import wraps
from typing import Any, Dict, Optional, TYPE_CHECKING

OPENAI_AVAILABLE = False

try:
    from openai import OpenAI, AsyncOpenAI
    OPENAI_AVAILABLE = True
except ImportError:
    OpenAI = object
    AsyncOpenAI = object

from agent_status.client import AgentStatusClient

logger = logging.getLogger(__name__)


class TrackedCompletions:
    """Wrapper for chat.completions that tracks metrics."""

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

    def create(self, *args, **kwargs) -> Any:
        """Create a chat completion 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.total_tokens
                self._metrics["prompt_tokens"] += result.usage.prompt_tokens
                self._metrics["completion_tokens"] += result.usage.completion_tokens

            logger.debug(f"OpenAI 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"OpenAI request failed: {e}")
            raise


class TrackedChat:
    """Wrapper for chat that tracks metrics."""

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


class AgentStatusOpenAI:
    """
    OpenAI client wrapper that automatically tracks metrics for Agent Status.

    Drop-in replacement for the OpenAI client that tracks:
    - Request counts and latencies
    - Token usage (prompt, completion, 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 OpenAI client

    Example:
        from agent_status.integrations.openai_wrapper import AgentStatusOpenAI

        client = AgentStatusOpenAI(rora_api_key="fb_live_xxx")

        response = client.chat.completions.create(
            model="gpt-4",
            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 OPENAI_AVAILABLE:
            raise ImportError(
                "OpenAI is not installed. Install with: pip install openai"
            )

        self._openai = OpenAI(**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,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "by_model": {},
            "errors": [],
        }

        self.chat = TrackedChat(self._openai.chat, 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,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "by_model": {},
            "errors": [],
        }
        self.chat = TrackedChat(self._openai.chat, 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._openai, name)


class AsyncTrackedCompletions:
    """Async wrapper for chat.completions that tracks metrics."""

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

    async def create(self, *args, **kwargs) -> Any:
        """Create a chat completion 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.total_tokens
                self._metrics["prompt_tokens"] += result.usage.prompt_tokens
                self._metrics["completion_tokens"] += result.usage.completion_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 AsyncTrackedChat:
    """Async wrapper for chat that tracks metrics."""

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


class AgentStatusAsyncOpenAI:
    """
    Async OpenAI client wrapper that automatically tracks metrics for Agent Status.

    Drop-in replacement for the AsyncOpenAI client.

    Example:
        from agent_status.integrations.openai_wrapper import AgentStatusAsyncOpenAI

        client = AgentStatusAsyncOpenAI(rora_api_key="fb_live_xxx")

        response = await client.chat.completions.create(
            model="gpt-4",
            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 OPENAI_AVAILABLE:
            raise ImportError(
                "OpenAI is not installed. Install with: pip install openai"
            )

        self._openai = AsyncOpenAI(**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,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "by_model": {},
            "errors": [],
        }

        self.chat = AsyncTrackedChat(self._openai.chat, 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,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "by_model": {},
            "errors": [],
        }
        self.chat = AsyncTrackedChat(self._openai.chat, 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._openai, name)


# Backward compatibility aliases
RoraOpenAI = AgentStatusOpenAI
RoraAsyncOpenAI = AgentStatusAsyncOpenAI
