"""
Google Gemini wrapper for automatic usage tracking with Paygent.

Design Principles (Enterprise Grade):
  - The ENTIRE wrapper method is wrapped in try/except. The original Gemini
    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.

Usage is identical to the standard google-genai SDK, with only the addition of
tracking parameters.
"""

import json
from typing import Any, Optional

try:
    from google import genai
except ImportError:
    raise ImportError(
        "google-genai package is a peer-dependency. To use the Paygent wrapper around google-genai "
        "you're assumed to already have google-genai 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


def _extract_gemini_output(response) -> str:
    """Best-effort extraction of text from a Gemini response. Never raises."""
    try:
        return response.text or ""
    except Exception:
        pass
    try:
        cands = getattr(response, 'candidates', []) or []
        parts = []
        for cand in cands:
            for p in getattr(getattr(cand, 'content', None), 'parts', []) or []:
                parts.append(getattr(p, 'text', '') or '')
        return ''.join(parts)
    except Exception:
        return ""


class PaygentGemini:
    """Main wrapper class for Google Gemini that provides automatic usage tracking."""

    def __init__(self, gemini_client: genai.Client, paygent_client: Client):
        """
        Create a new PaygentGemini wrapper.

        Args:
            gemini_client: The GoogleGenAI client instance from google-genai
            paygent_client: The Paygent client instance for usage tracking
        """
        self.gemini = gemini_client
        self.paygent_client = paygent_client

    @property
    def models(self) -> 'ModelsWrapper':
        """Access to models API with automatic usage tracking."""
        return ModelsWrapper(self.gemini, self.paygent_client)


class ModelsWrapper:
    """Wrapper for Gemini models API."""

    def __init__(self, gemini_client: genai.Client, paygent_client: Client):
        self.gemini = gemini_client
        self.paygent_client = paygent_client

    def generate_content(
        self,
        *,
        model: str,
        contents: Any,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
        external_customer_name: str = "",
        **kwargs
    ) -> Any:
        """
        Generate content with automatic usage tracking.

        The Gemini 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
            contents: The content to generate from
            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 Gemini parameters

        Returns:
            The generation response from Gemini
        """
        # Make the Gemini API call — this MUST succeed
        response = self.gemini.models.generate_content(
            model=model,
            contents=contents,
            **kwargs
        )

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

            um = getattr(response, 'usage_metadata', None)
            inp = getattr(um, 'prompt_token_count', 0) or 0
            out = getattr(um, 'candidates_token_count', 0) or 0
            tot = getattr(um, 'total_token_count', inp + out) or inp + out

            if inp > 0 or out > 0:
                usage_data = UsageData(
                    service_provider=ServiceProvider.GOOGLE,
                    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 = ""
                try:
                    prompt_string = json.dumps(contents) if not isinstance(contents, str) else contents
                except Exception:
                    try:
                        prompt_string = str(contents)
                    except Exception:
                        prompt_string = ""

                usage_data_ws = UsageDataWithStrings(
                    service_provider=ServiceProvider.GOOGLE,
                    model=model,
                    prompt_string=prompt_string,
                    output_string=_extract_gemini_output(response),
                )
                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] GeminiWrapper tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response

    def start_chat(
        self,
        *,
        model: str,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
        external_customer_name: str = "",
        **kwargs
    ) -> 'ChatSessionWrapper':
        """
        Start a chat session with automatic usage tracking.

        Args:
            model: The model to use
            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 chat configuration

        Returns:
            A wrapped chat session with automatic tracking
        """
        try:
            _auto_create_customer(self.paygent_client, external_customer_id, external_customer_name)
        except Exception:
            pass

        chat_session = self.gemini.models.start_chat(model=model, **kwargs)

        return ChatSessionWrapper(
            chat_session=chat_session,
            paygent_client=self.paygent_client,
            model_name=model,
            indicator=indicator,
            external_agent_id=external_agent_id,
            external_customer_id=external_customer_id,
        )

    def generate_image(
        self,
        *,
        model: str,
        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 Gemini call always succeeds and the response is always returned.
        Any tracking failure is absorbed silently.

        Args:
            model: The model to use (e.g., "imagen-3.0-generate-001")
            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 image generation parameters

        Returns:
            The image generation response
        """
        # Make the image generation API call — this MUST succeed
        response = self.gemini.models.generate_content(
            model=model,
            contents=prompt,
            **kwargs
        )

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

            um = getattr(response, 'usage_metadata', None)
            inp = getattr(um, 'prompt_token_count', 0) or 0
            out = getattr(um, 'candidates_token_count', 0) or 0
            tot = getattr(um, 'total_token_count', inp + out) or inp + out

            if inp > 0 or out > 0:
                usage_data = UsageData(
                    service_provider=ServiceProvider.GOOGLE,
                    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 via prompt string
                usage_data_ws = UsageDataWithStrings(
                    service_provider=ServiceProvider.GOOGLE,
                    model=model,
                    prompt_string=prompt or "",
                    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] GeminiWrapper generate_image tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response


class ChatSessionWrapper:
    """Wrapper for Gemini ChatSession with automatic per-turn usage tracking."""

    def __init__(
        self,
        chat_session: Any,
        paygent_client: Client,
        model_name: str,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
    ):
        self.chat_session = chat_session
        self.paygent_client = paygent_client
        self.model_name = model_name
        self.indicator = indicator
        self.external_agent_id = external_agent_id
        self.external_customer_id = external_customer_id

    def send_message(self, message: str) -> Any:
        """
        Send a message in the chat with automatic usage tracking.

        The Gemini call always succeeds and the response is always returned.
        Any tracking failure is absorbed silently.

        Args:
            message: The message to send

        Returns:
            The chat response from Gemini
        """
        # Make the Gemini API call — this MUST succeed
        response = self.chat_session.send_message(message)

        # --- Tracking (fire-and-forget, never breaks caller) ---
        try:
            um = getattr(response, 'usage_metadata', None)
            inp = getattr(um, 'prompt_token_count', 0) or 0
            out = getattr(um, 'candidates_token_count', 0) or 0
            tot = getattr(um, 'total_token_count', inp + out) or inp + out

            if inp > 0 or out > 0:
                usage_data = UsageData(
                    service_provider=ServiceProvider.GOOGLE,
                    model=self.model_name,
                    prompt_tokens=inp,
                    completion_tokens=out,
                    total_tokens=tot,
                )
                self.paygent_client.send_usage(
                    self.external_agent_id,
                    self.external_customer_id,
                    self.indicator,
                    usage_data,
                )
            else:
                usage_data_ws = UsageDataWithStrings(
                    service_provider=ServiceProvider.GOOGLE,
                    model=self.model_name,
                    prompt_string=message or "",
                    output_string=_extract_gemini_output(response),
                )
                self.paygent_client.send_usage_with_token_string(
                    self.external_agent_id,
                    self.external_customer_id,
                    self.indicator,
                    usage_data_ws,
                )
        except Exception as e:
            try:
                self.paygent_client.logger.error(
                    f"[Paygent] GeminiChatSession tracking error (response still returned): {e}"
                )
            except Exception:
                pass

        return response

    def get_history(self) -> Any:
        """
        Get the chat history.

        Returns:
            The chat history
        """
        try:
            return self.chat_session.get_history()
        except Exception:
            return []
