"""
OpenAI wrapper for automatic usage tracking with Paygent.

Design Principles (Enterprise Grade):
  - The ENTIRE wrapper method is wrapped in try/except. Even if response
    parsing or tracking explodes, the original OpenAI response is ALWAYS
    returned to the caller. The SDK NEVER breaks client flow.
  - Automatic create_or_get_customer is fired in the background on every call.
  - Cached token extraction follows the patching.py pattern exactly.
  - String-based fallback guarantees billing capture even when the API
    response format changes.
  - send_token_string is used when tokens are 0 so billing always works.
"""

import json
from typing import Any, Optional

try:
    from openai import OpenAI
except ImportError:
    raise ImportError(
        "openai package is a peer-dependency. To use the Paygent wrapper around openai "
        "you're assumed to already have openai package installed."
    )

from ..client import Client
from ..models import UsageData, UsageDataWithStrings, CustomerCreateOrGetRequest
from ..constants import ServiceProvider


def _auto_create_customer(paygent_client: Client, customer_id: str, customer_name: str = "") -> None:
    """Fire-and-forget customer creation in the background. Never raises."""
    try:
        if customer_id:
            paygent_client.create_or_get_customer(
                CustomerCreateOrGetRequest(
                    external_id=customer_id,
                    name=customer_name or customer_id,
                )
            )
    except Exception:
        pass


class PaygentOpenAI:
    """Main wrapper class for OpenAI that provides automatic usage tracking."""

    def __init__(self, openai_client: OpenAI, paygent_client: Client):
        """
        Create a new PaygentOpenAI wrapper.

        Args:
            openai_client: The OpenAI client instance
            paygent_client: The Paygent client instance for usage tracking
        """
        self.openai = openai_client
        self.paygent_client = paygent_client

    @property
    def chat(self) -> 'ChatWrapper':
        """Access to chat API with automatic usage tracking."""
        return ChatWrapper(self.openai, self.paygent_client)

    @property
    def embeddings(self) -> 'EmbeddingsWrapper':
        """Access to embeddings API with automatic usage tracking."""
        return EmbeddingsWrapper(self.openai, self.paygent_client)

    @property
    def images(self) -> 'ImagesWrapper':
        """Access to images API with automatic usage tracking."""
        return ImagesWrapper(self.openai, self.paygent_client)


class ChatWrapper:
    """Wrapper for OpenAI chat API."""

    def __init__(self, openai_client: OpenAI, paygent_client: Client):
        self.openai = openai_client
        self.paygent_client = paygent_client

    @property
    def completions(self) -> 'ChatCompletionsWrapper':
        """Access to chat completions API with automatic usage tracking."""
        return ChatCompletionsWrapper(self.openai, self.paygent_client)


class ChatCompletionsWrapper:
    """Wrapper for OpenAI chat completions API."""

    def __init__(self, openai_client: OpenAI, paygent_client: Client):
        self.openai = openai_client
        self.paygent_client = paygent_client

    def create(
        self,
        *,
        model: str,
        messages: list,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
        external_customer_name: str = "",
        **kwargs
    ) -> Any:
        """
        Create a chat completion with automatic usage tracking.

        The OpenAI call always succeeds and the response is always returned.
        Any tracking failure is absorbed silently — it NEVER crashes the caller.

        Args:
            model: The model to use
            messages: The messages to send
            indicator: Indicator for the usage event
            external_agent_id: External agent identifier
            external_customer_id: External customer identifier
            external_customer_name: Optional customer display name
            **kwargs: Additional OpenAI parameters

        Returns:
            The chat completion response from OpenAI
        """
        # Ensure streaming is disabled for automatic tracking
        kwargs['stream'] = False

        # Make the OpenAI API call (non-streaming) — this MUST succeed
        response = self.openai.chat.completions.create(
            model=model,
            messages=messages,
            **kwargs
        )

        # --- Tracking (fire-and-forget, never breaks caller) ---
        try:
            # Auto-create customer in background
            _auto_create_customer(self.paygent_client, external_customer_id, external_customer_name)

            has_valid_usage = (
                hasattr(response, 'usage') and
                response.usage is not None and
                (getattr(response.usage, 'prompt_tokens', 0) or 0) > 0 and
                (getattr(response.usage, 'completion_tokens', 0) or 0) > 0
            )

            if has_valid_usage:
                # Primary path: Use token counts from API response
                cached_tokens = 0
                try:
                    ptd = getattr(response.usage, 'prompt_tokens_details', None)
                    if ptd:
                        cached_tokens = getattr(ptd, 'cached_tokens', 0) or 0
                except Exception:
                    cached_tokens = 0

                usage_data = UsageData(
                    service_provider=ServiceProvider.OPENAI,
                    model=model,
                    prompt_tokens=response.usage.prompt_tokens,
                    completion_tokens=response.usage.completion_tokens,
                    total_tokens=getattr(response.usage, 'total_tokens', 0) or 0,
                    cached_tokens=cached_tokens,
                )
                self.paygent_client.send_usage(
                    external_agent_id, external_customer_id, indicator, usage_data
                )
            else:
                # Fallback: count from raw strings — never lose a billing event
                prompt_string  = ""
                output_string  = ""
                try:
                    prompt_string = json.dumps(messages)
                except Exception:
                    try:
                        prompt_string = str(messages)
                    except Exception:
                        prompt_string = ""
                try:
                    if hasattr(response, 'choices') and response.choices:
                        msg = getattr(response.choices[0], 'message', None)
                        if msg:
                            output_string = getattr(msg, 'content', '') or ''
                except Exception:
                    output_string = ""

                usage_data_ws = UsageDataWithStrings(
                    service_provider=ServiceProvider.OPENAI,
                    model=model,
                    prompt_string=prompt_string,
                    output_string=output_string,
                )
                self.paygent_client.send_usage_with_token_string(
                    external_agent_id, external_customer_id, indicator, usage_data_ws
                )
        except Exception as e:
            try:
                self.paygent_client.logger.error(
                    f"[Paygent] ChatCompletionsWrapper tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response


class EmbeddingsWrapper:
    """Wrapper for OpenAI embeddings API."""

    def __init__(self, openai_client: OpenAI, paygent_client: Client):
        self.openai = openai_client
        self.paygent_client = paygent_client

    def create(
        self,
        *,
        model: str,
        input: Any,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
        external_customer_name: str = "",
        **kwargs
    ) -> Any:
        """
        Create embeddings with automatic usage tracking.

        The OpenAI call always succeeds and the response is always returned.
        Tracking failure is absorbed silently.

        Args:
            model: The model to use
            input: The input text(s) to embed
            indicator: Indicator for the usage event
            external_agent_id: External agent identifier
            external_customer_id: External customer identifier
            external_customer_name: Optional customer display name
            **kwargs: Additional OpenAI parameters

        Returns:
            The embeddings response from OpenAI
        """
        # Make the OpenAI API call — this MUST succeed
        response = self.openai.embeddings.create(
            model=model,
            input=input,
            **kwargs
        )

        # --- Tracking (fire-and-forget, never breaks caller) ---
        try:
            _auto_create_customer(self.paygent_client, external_customer_id, external_customer_name)

            has_valid_usage = (
                hasattr(response, 'usage') and
                response.usage is not None and
                (getattr(response.usage, 'prompt_tokens', 0) or 0) > 0
            )

            if has_valid_usage:
                usage_data = UsageData(
                    service_provider=ServiceProvider.OPENAI,
                    model=model,
                    prompt_tokens=response.usage.prompt_tokens,
                    completion_tokens=0,  # Embeddings don't have completion tokens
                    total_tokens=getattr(response.usage, 'total_tokens', 0) or 0,
                )
                self.paygent_client.send_usage(
                    external_agent_id, external_customer_id, indicator, usage_data
                )
            else:
                # Fallback: count from input text
                input_text = ""
                try:
                    if isinstance(input, list):
                        input_text = " ".join(str(i) for i in input)
                    else:
                        input_text = str(input)
                except Exception:
                    input_text = ""

                usage_data_ws = UsageDataWithStrings(
                    service_provider=ServiceProvider.OPENAI,
                    model=model,
                    prompt_string=input_text,
                    output_string="",
                )
                self.paygent_client.send_usage_with_token_string(
                    external_agent_id, external_customer_id, indicator, usage_data_ws
                )
        except Exception as e:
            try:
                self.paygent_client.logger.error(
                    f"[Paygent] EmbeddingsWrapper tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response


class ImagesWrapper:
    """Wrapper for OpenAI images API."""

    def __init__(self, openai_client: OpenAI, paygent_client: Client):
        self.openai = openai_client
        self.paygent_client = paygent_client

    def generate(
        self,
        *,
        model: Optional[str] = None,
        prompt: str,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
        external_customer_name: str = "",
        **kwargs
    ) -> Any:
        """
        Generate images with automatic usage tracking.

        The OpenAI call always succeeds and the response is always returned.
        Tracking failure is absorbed silently.

        Note: OpenAI's image generation API doesn't return token usage — we
        use the string-based fallback path to capture the prompt cost.

        Args:
            model: The model to use (optional, defaults to dall-e-2)
            prompt: The prompt for image generation
            indicator: Indicator for the usage event
            external_agent_id: External agent identifier
            external_customer_id: External customer identifier
            external_customer_name: Optional customer display name
            **kwargs: Additional OpenAI parameters

        Returns:
            The images response from OpenAI
        """
        effective_model = model or "dall-e-2"

        # Make the OpenAI API call — this MUST succeed
        if model:
            response = self.openai.images.generate(model=model, prompt=prompt, **kwargs)
        else:
            response = self.openai.images.generate(prompt=prompt, **kwargs)

        # --- Tracking (fire-and-forget, never breaks caller) ---
        try:
            _auto_create_customer(self.paygent_client, external_customer_id, external_customer_name)

            # Images don't return token counts — use string tracking so the
            # backend's pricing table handles cost calculation via the prompt.
            usage_data_ws = UsageDataWithStrings(
                service_provider=ServiceProvider.OPENAI,
                model=effective_model,
                prompt_string=prompt or "",
                output_string="",  # No text output for image generation
            )
            self.paygent_client.send_usage_with_token_string(
                external_agent_id, external_customer_id, indicator, usage_data_ws
            )
        except Exception as e:
            try:
                self.paygent_client.logger.error(
                    f"[Paygent] ImagesWrapper tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response
