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

Design Principles (Enterprise Grade):
  - The ENTIRE wrapper method is wrapped in try/except. The original Mistral
    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.
  - String-based fallback guarantees billing capture even when API format changes.
  - Backward compatible: same signature as before + optional customer_name param.
"""

import json
from typing import Any, Optional

try:
    from mistralai import Mistral
except ImportError:
    raise ImportError(
        "mistralai package is a peer-dependency. To use the Paygent wrapper around mistralai "
        "you're assumed to already have mistralai 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. 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 PaygentMistral:
    """Main wrapper class for Mistral that provides automatic usage tracking."""

    def __init__(self, mistral_client: Mistral, paygent_client: Client):
        """
        Create a new PaygentMistral wrapper.

        Args:
            mistral_client: The Mistral client instance
            paygent_client: The Paygent client instance for usage tracking
        """
        self.mistral = mistral_client
        self.paygent_client = paygent_client

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


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

    def __init__(self, mistral_client: Mistral, paygent_client: Client):
        self.mistral = mistral_client
        self.paygent_client = paygent_client

    def complete(
        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 Mistral call always succeeds and the response is always returned.
        Any tracking failure is absorbed silently — 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 Mistral parameters

        Returns:
            The chat completion response from Mistral
        """
        # Make the Mistral API call — this MUST succeed
        response = self.mistral.chat.complete(
            model=model,
            messages=messages,
            **kwargs
        )

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

            usage = getattr(response, 'usage', None)
            inp = getattr(usage, 'prompt_tokens', 0) or 0
            out = getattr(usage, 'completion_tokens', 0) or 0
            tot = getattr(usage, 'total_tokens', inp + out) or inp + out

            if inp > 0 or out > 0:
                usage_data = UsageData(
                    service_provider=ServiceProvider.MISTRAL_AI,
                    model=model,
                    prompt_tokens=inp,
                    completion_tokens=out,
                    total_tokens=tot,
                )
                self.paygent_client.send_usage(
                    external_agent_id, external_customer_id, indicator, usage_data
                )
            else:
                # Fallback: count from raw strings
                prompt_string = ""
                output_string = ""
                try:
                    prompt_string = json.dumps(messages)
                except Exception:
                    try:
                        prompt_string = str(messages)
                    except Exception:
                        prompt_string = ""
                try:
                    choices = getattr(response, 'choices', None) or []
                    if choices:
                        msg = getattr(choices[0], 'message', None)
                        if msg:
                            output_string = getattr(msg, 'content', '') or ''
                except Exception:
                    output_string = ""

                usage_data_ws = UsageDataWithStrings(
                    service_provider=ServiceProvider.MISTRAL_AI,
                    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] MistralWrapper tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response
