"""
AICompliant Discover SDK v2.0.0 — Python

Enterprise-grade AI system auto-discovery via explicit wrapping.
No monkey-patching, no global state mutation.

Usage:
    from aicompliant_discover import AIDiscovery

    discovery = AIDiscovery(api_key="dk_...", endpoint="https://...")

    # Wrap your Anthropic client — automatic usage tracking
    client = discovery.wrap_anthropic(anthropic.Anthropic())
    msg = client.messages.create(model="claude-sonnet-4-20250514", ...)

    # Wrap Google GenAI model
    model = discovery.wrap_google(genai.GenerativeModel("gemini-2.5-flash"))
    response = model.generate_content("Hello")

    # Generic wrapper for any provider
    result = discovery.trace("My System", "OpenAI", lambda: openai_call())

    # Manual reporting
    discovery.capture(system_name="...", provider="...", model="...", input_tokens=100, output_tokens=200)
"""

from __future__ import annotations

__version__ = "2.0.0"

import atexit
import json
import os
import threading
import time
from datetime import datetime, timezone
from typing import Any, Callable, Dict, List, Optional, TypeVar

try:
    import urllib.request
    import urllib.error
except ImportError:
    pass

T = TypeVar("T")

# ─── Event Buffer ─────────────────────────────────────────────

_DEFAULT_FLUSH_SIZE = 50
_DEFAULT_FLUSH_INTERVAL = 30  # seconds
_MAX_RETRIES = 3
_RETRY_BASE_DELAY = 1.0


class _EventBuffer:
    """Thread-safe event buffer with periodic flush."""

    def __init__(self, endpoint: str, api_key: str, flush_size: int, flush_interval: float):
        self._endpoint = endpoint
        self._api_key = api_key
        self._flush_size = flush_size
        self._flush_interval = flush_interval
        self._events: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self._timer: Optional[threading.Timer] = None
        self._closed = False
        self._start_timer()

    def _start_timer(self) -> None:
        if self._closed:
            return
        self._timer = threading.Timer(self._flush_interval, self._timer_flush)
        self._timer.daemon = True
        self._timer.start()

    def _timer_flush(self) -> None:
        self.flush()
        self._start_timer()

    def add(self, event: Dict[str, Any]) -> None:
        flush_now = False
        events: List[Dict[str, Any]] = []
        with self._lock:
            self._events.append(event)
            if len(self._events) >= self._flush_size:
                events = self._events[:]
                self._events = []
                flush_now = True
        if flush_now:
            self._send(events)

    def flush(self) -> None:
        with self._lock:
            if not self._events:
                return
            events = self._events[:]
            self._events = []
        self._send(events)

    def _send(self, events: List[Dict[str, Any]]) -> None:
        if not events:
            return
        for event in events:
            self._send_single(event)

    def _send_single(self, event: Dict[str, Any]) -> None:
        payload = json.dumps(event).encode("utf-8")
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self._api_key}",
        }
        for attempt in range(_MAX_RETRIES):
            try:
                req = urllib.request.Request(
                    self._endpoint,
                    data=payload,
                    headers=headers,
                    method="POST",
                )
                with urllib.request.urlopen(req, timeout=10) as resp:
                    if resp.status < 300:
                        return
                    if resp.status in (429, 500, 502, 503, 504):
                        time.sleep(_RETRY_BASE_DELAY * (2 ** attempt))
                        continue
                    return  # non-retryable error
            except (urllib.error.URLError, OSError):
                time.sleep(_RETRY_BASE_DELAY * (2 ** attempt))
            except Exception:
                return  # unexpected error, don't retry

    def close(self) -> None:
        self._closed = True
        if self._timer:
            self._timer.cancel()
        self.flush()


# ─── Anthropic Wrapper ────────────────────────────────────────

class _WrappedAnthropicMessages:
    """Wraps anthropic.resources.messages.Messages to intercept create() calls."""

    def __init__(self, original_messages: Any, discovery: "AIDiscovery"):
        self._original = original_messages
        self._discovery = discovery

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

        # Extract usage from response
        try:
            model = getattr(result, "model", kwargs.get("model", "unknown"))
            usage = getattr(result, "usage", None)
            input_tokens = getattr(usage, "input_tokens", 0) if usage else 0
            output_tokens = getattr(usage, "output_tokens", 0) if usage else 0
            self._discovery._record(
                provider="Anthropic",
                model=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                latency_ms=latency_ms,
            )
        except Exception:
            pass

        return result

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


class _WrappedAnthropicAsyncMessages:
    """Wraps async anthropic.resources.messages.AsyncMessages."""

    def __init__(self, original_messages: Any, discovery: "AIDiscovery"):
        self._original = original_messages
        self._discovery = discovery

    async def create(self, **kwargs: Any) -> Any:
        start = time.time()
        result = await self._original.create(**kwargs)
        latency_ms = int((time.time() - start) * 1000)

        try:
            model = getattr(result, "model", kwargs.get("model", "unknown"))
            usage = getattr(result, "usage", None)
            input_tokens = getattr(usage, "input_tokens", 0) if usage else 0
            output_tokens = getattr(usage, "output_tokens", 0) if usage else 0
            self._discovery._record(
                provider="Anthropic",
                model=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                latency_ms=latency_ms,
            )
        except Exception:
            pass

        return result

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


class _WrappedAnthropicClient:
    """Wraps an Anthropic client instance. Access .messages.create() transparently."""

    def __init__(self, client: Any, discovery: "AIDiscovery"):
        self._client = client
        self._wrapped_messages = _WrappedAnthropicMessages(client.messages, discovery)

    @property
    def messages(self) -> _WrappedAnthropicMessages:
        return self._wrapped_messages

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


class _WrappedAsyncAnthropicClient:
    """Wraps an AsyncAnthropic client instance."""

    def __init__(self, client: Any, discovery: "AIDiscovery"):
        self._client = client
        self._wrapped_messages = _WrappedAnthropicAsyncMessages(client.messages, discovery)

    @property
    def messages(self) -> _WrappedAnthropicAsyncMessages:
        return self._wrapped_messages

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


# ─── Google GenAI Wrapper ─────────────────────────────────────

class _WrappedGoogleModel:
    """Wraps a Google GenerativeModel to intercept generate_content() calls."""

    def __init__(self, model: Any, discovery: "AIDiscovery"):
        self._model = model
        self._discovery = discovery

    def generate_content(self, *args: Any, **kwargs: Any) -> Any:
        start = time.time()
        result = self._model.generate_content(*args, **kwargs)
        latency_ms = int((time.time() - start) * 1000)

        try:
            model_name = getattr(self._model, "model_name", "gemini-unknown")
            usage = getattr(result, "usage_metadata", None)
            input_tokens = getattr(usage, "prompt_token_count", 0) if usage else 0
            output_tokens = getattr(usage, "candidates_token_count", 0) if usage else 0
            self._discovery._record(
                provider="Google AI",
                model=model_name,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                latency_ms=latency_ms,
            )
        except Exception:
            pass

        return result

    async def generate_content_async(self, *args: Any, **kwargs: Any) -> Any:
        start = time.time()
        result = await self._model.generate_content_async(*args, **kwargs)
        latency_ms = int((time.time() - start) * 1000)

        try:
            model_name = getattr(self._model, "model_name", "gemini-unknown")
            usage = getattr(result, "usage_metadata", None)
            input_tokens = getattr(usage, "prompt_token_count", 0) if usage else 0
            output_tokens = getattr(usage, "candidates_token_count", 0) if usage else 0
            self._discovery._record(
                provider="Google AI",
                model=model_name,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                latency_ms=latency_ms,
            )
        except Exception:
            pass

        return result

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


# ─── Main Discovery Class ────────────────────────────────────

class AIDiscovery:
    """
    AICompliant Discovery SDK v2.0.0

    Enterprise-grade AI system auto-discovery using explicit wrapping.
    No monkey-patching, no prototype mutation.

    Args:
        api_key: AICompliant API key (dk_...). Falls back to AIC_API_KEY env var.
        endpoint: Discovery webhook endpoint. Falls back to AIC_ENDPOINT env var.
        system_name: Default system name. Falls back to AIC_SYSTEM_NAME / SERVICE_NAME / APP_NAME env vars.
        flush_size: Number of events before auto-flush (default: 50).
        flush_interval: Seconds between periodic flushes (default: 30).
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        endpoint: Optional[str] = None,
        system_name: Optional[str] = None,
        flush_size: int = _DEFAULT_FLUSH_SIZE,
        flush_interval: float = _DEFAULT_FLUSH_INTERVAL,
    ):
        self._api_key = api_key or os.environ.get("AIC_API_KEY", "")
        self._endpoint = endpoint or os.environ.get("AIC_ENDPOINT", "")
        self._system_name = (
            system_name
            or os.environ.get("AIC_SYSTEM_NAME")
            or os.environ.get("SERVICE_NAME")
            or os.environ.get("APP_NAME")
            or "unknown"
        )

        if not self._api_key:
            raise ValueError("api_key is required (or set AIC_API_KEY env var)")
        if not self._endpoint:
            raise ValueError("endpoint is required (or set AIC_ENDPOINT env var)")

        self._buffer = _EventBuffer(self._endpoint, self._api_key, flush_size, flush_interval)
        atexit.register(self._buffer.close)

    def _record(
        self,
        provider: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: int = 0,
        system_name: Optional[str] = None,
    ) -> None:
        """Record a usage event to the buffer."""
        event = {
            "system_name": system_name or self._system_name,
            "provider": provider,
            "model": model,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
            "metadata": {"source": "sdk-python-v2", "latency_ms": latency_ms},
        }
        self._buffer.add(event)

    # ── Explicit Wrapping Methods ──

    def wrap_anthropic(self, client: Any) -> Any:
        """
        Wrap an Anthropic client for automatic usage tracking.

        Args:
            client: An anthropic.Anthropic() or anthropic.AsyncAnthropic() instance.

        Returns:
            A wrapped client with the same API. All .messages.create() calls
            are automatically tracked.

        Example:
            client = discovery.wrap_anthropic(anthropic.Anthropic())
            msg = client.messages.create(model="claude-sonnet-4-20250514", ...)
        """
        # Detect sync vs async
        client_type = type(client).__name__
        if "Async" in client_type:
            return _WrappedAsyncAnthropicClient(client, self)
        return _WrappedAnthropicClient(client, self)

    def wrap_google(self, model: Any) -> _WrappedGoogleModel:
        """
        Wrap a Google GenerativeModel for automatic usage tracking.

        Args:
            model: A google.generativeai.GenerativeModel instance.

        Returns:
            A wrapped model with the same API. All .generate_content() calls
            are automatically tracked.

        Example:
            model = discovery.wrap_google(genai.GenerativeModel("gemini-2.5-flash"))
            response = model.generate_content("Hello")
        """
        return _WrappedGoogleModel(model, self)

    def trace(
        self,
        system_name: str,
        provider: str = "unknown",
        fn: Optional[Callable[[], T]] = None,
        *,
        model: str = "unknown",
        extract_usage: Optional[Callable[[Any], Dict[str, int]]] = None,
    ) -> T:
        """
        Generic wrapper for any AI provider call.

        Args:
            system_name: Name for this AI system.
            provider: Provider name (e.g. "OpenAI").
            fn: The function to call (your AI API call).
            model: Model name (default: "unknown").
            extract_usage: Optional function to extract usage from the result.
                Should return {"input_tokens": int, "output_tokens": int}.

        Returns:
            The result of fn().

        Example:
            result = discovery.trace("My System", "OpenAI",
                lambda: openai.chat.completions.create(...),
                model="gpt-4",
                extract_usage=lambda r: {
                    "input_tokens": r.usage.prompt_tokens,
                    "output_tokens": r.usage.completion_tokens
                })
        """
        if fn is None:
            raise ValueError("fn (callable) is required")

        start = time.time()
        result = fn()
        latency_ms = int((time.time() - start) * 1000)

        input_tokens = 0
        output_tokens = 0

        if extract_usage:
            try:
                usage = extract_usage(result)
                input_tokens = usage.get("input_tokens", 0)
                output_tokens = usage.get("output_tokens", 0)
            except Exception:
                pass

        self._record(
            provider=provider,
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            system_name=system_name,
        )

        return result

    def capture(
        self,
        system_name: Optional[str] = None,
        provider: str = "unknown",
        model: str = "unknown",
        input_tokens: int = 0,
        output_tokens: int = 0,
    ) -> None:
        """
        Manually report an AI usage event.

        Args:
            system_name: Name for this AI system.
            provider: Provider name.
            model: Model name.
            input_tokens: Number of input tokens.
            output_tokens: Number of output tokens.

        Example:
            discovery.capture(
                system_name="Resume Screener",
                provider="Anthropic",
                model="claude-sonnet-4-20250514",
                input_tokens=500,
                output_tokens=1500
            )
        """
        self._record(
            provider=provider,
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            system_name=system_name,
        )

    def flush(self) -> None:
        """Flush all buffered events immediately."""
        self._buffer.flush()

    def close(self) -> None:
        """Flush and close the discovery client."""
        self._buffer.close()
