"""Gemini integration with auto-instrumentation via monkey patching."""

import functools
import logging
from typing import Any

from fenra._context import get_context
from fenra._core import enqueue_transaction

logger = logging.getLogger(__name__)

_patched_sync = False
_patched_async = False
_patched_stream_sync = False
_patched_stream_async = False

# Store original methods for unpatching
_original_generate_content: Any = None
_original_generate_content_async: Any = None
_original_generate_content_stream: Any = None
_original_generate_content_stream_async: Any = None


def _track_generate_content_response(
    response: Any,
    model: str,
    context: dict[str, Any],
    config: Any = None,
) -> None:
    """
    Extract usage from GenerateContentResponse and queue transaction.

    Handles:
    - Token usage (including thinking tokens for reasoning models)
    - Google Search tool usage (web_search requests)
    - Image generation output
    """
    if not hasattr(response, "usage_metadata") or response.usage_metadata is None:
        return

    usage = response.usage_metadata

    # Build token metrics
    # thoughts_token_count is used by thinking models like gemini-2.0-flash-thinking
    output_tokens = (getattr(usage, "candidates_token_count", 0) or 0) + (
        getattr(usage, "thoughts_token_count", None) or 0
    )

    metrics: dict[str, Any] = {
        "input_tokens": getattr(usage, "prompt_token_count", 0) or 0,
        "output_tokens": output_tokens,
        "total_tokens": getattr(usage, "total_token_count", 0) or 0,
    }

    # Include cached tokens if present
    cached = getattr(usage, "cached_content_token_count", None)
    if cached:
        metrics["cached_tokens"] = cached

    usage_entries: list[dict[str, Any]] = [
        {
            "type": "tokens",
            "metrics": metrics,
        }
    ]

    # Detect Google Search tool usage
    search_usage = _detect_tool_usage(response, config)
    if search_usage:
        usage_entries.append(search_usage)

    # Detect image generation output
    image_usage = _detect_image_output(response, config)
    if image_usage:
        usage_entries.append(image_usage)

    transaction: dict[str, Any] = {
        "provider": "gemini",
        "model": model,
        "usage": usage_entries,
        "context": context,
    }

    # Include raw usage for debugging if available
    try:
        transaction["provider_usage_raw"] = {
            "prompt_token_count": getattr(usage, "prompt_token_count", None),
            "candidates_token_count": getattr(usage, "candidates_token_count", None),
            "total_token_count": getattr(usage, "total_token_count", None),
            "thoughts_token_count": getattr(usage, "thoughts_token_count", None),
            "cached_content_token_count": getattr(
                usage, "cached_content_token_count", None
            ),
        }
    except Exception:
        pass

    enqueue_transaction(transaction)


def _detect_tool_usage(response: Any, config: Any) -> dict[str, Any] | None:
    """
    Detect Google Search tool usage from grounding metadata.

    Returns a requests usage entry if web search was used, None otherwise.
    """
    if config is None:
        return None

    # Check if GoogleSearch tool was configured
    tools = getattr(config, "tools", None)
    if not tools:
        return None

    has_google_search = False
    for tool in tools:
        if hasattr(tool, "google_search") and tool.google_search is not None:
            has_google_search = True
            break

    if not has_google_search:
        return None

    # Extract search count from grounding metadata
    search_count = 0
    candidates = getattr(response, "candidates", None) or []
    for candidate in candidates:
        grounding_metadata = getattr(candidate, "grounding_metadata", None)
        if grounding_metadata:
            web_search_queries = getattr(grounding_metadata, "web_search_queries", None)
            if web_search_queries:
                search_count = len(web_search_queries)
                break

    # Default to 1 if tool was configured but we couldn't detect count
    if search_count == 0:
        search_count = 1

    return {
        "type": "requests",
        "metrics": {"count": search_count, "request_type": "web_search"},
    }


def _detect_image_output(response: Any, config: Any) -> dict[str, Any] | None:
    """
    Detect image generation from response parts.

    Returns an images usage entry if images were generated, None otherwise.
    """
    if config is None:
        return None

    # Check if IMAGE was in response modalities
    response_modalities = getattr(config, "response_modalities", None)
    if not response_modalities or "IMAGE" not in response_modalities:
        return None

    # Count images from response parts
    image_count = 0
    candidates = getattr(response, "candidates", None) or []
    for candidate in candidates:
        content = getattr(candidate, "content", None)
        if content:
            parts = getattr(content, "parts", None) or []
            for part in parts:
                if hasattr(part, "inline_data") and part.inline_data:
                    image_count += 1

    if image_count == 0:
        return None

    return {
        "type": "images",
        "metrics": {
            "generated": image_count,
            "size_px": 1024,  # Default Gemini image size
            "quality": "standard",
        },
    }


def _track_stream(
    stream: Any,
    model: str,
    context: dict[str, Any],
    config: Any = None,
) -> Any:
    """
    Wrap a streaming response to track usage from the final chunk.

    Returns a wrapped stream that tracks usage when the stream is exhausted.
    """
    last_chunk = None

    def wrapped_stream() -> Any:
        nonlocal last_chunk
        try:
            for chunk in stream:
                last_chunk = chunk
                yield chunk
        finally:
            # After stream is exhausted, track usage from final chunk
            if last_chunk:
                try:
                    _track_generate_content_response(last_chunk, model, context, config)
                except Exception as e:
                    logger.error(
                        f"Error tracking Gemini streaming usage: {e}", exc_info=True
                    )

    return wrapped_stream()


async def _track_stream_async(
    stream: Any,
    model: str,
    context: dict[str, Any],
    config: Any = None,
) -> Any:
    """
    Wrap an async streaming response to track usage from the final chunk.

    Returns a wrapped async stream that tracks usage when the stream is exhausted.
    """
    last_chunk = None

    try:
        async for chunk in stream:
            last_chunk = chunk
            yield chunk
    finally:
        # After stream is exhausted, track usage from final chunk
        if last_chunk:
            try:
                _track_generate_content_response(last_chunk, model, context, config)
            except Exception as e:
                logger.error(
                    f"Error tracking Gemini streaming usage: {e}", exc_info=True
                )


# ============================================================================
# Sync Patching
# ============================================================================


def patch_gemini() -> None:
    """
    Patch Gemini client to auto-track usage.

    This patches the synchronous `Models.generate_content` method.
    """
    global _patched_sync, _original_generate_content

    if _patched_sync:
        logger.debug("Gemini sync already patched, skipping")
        return

    try:
        from google.genai.models import Models
    except ImportError:
        logger.warning(
            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
        )
        return

    original_generate_content = Models.generate_content
    _original_generate_content = original_generate_content

    @functools.wraps(original_generate_content)
    def patched_generate_content(self: Models, *args: Any, **kwargs: Any) -> Any:
        response = original_generate_content(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or (args[0] if args else "unknown")

        config = kwargs.get("config")

        try:
            _track_generate_content_response(response, model, context, config)
        except Exception as e:
            logger.error(f"Error tracking Gemini usage: {e}", exc_info=True)

        return response

    Models.generate_content = patched_generate_content  # type: ignore[assignment]
    _patched_sync = True
    logger.info("Gemini SDK patched for auto-instrumentation")


def patch_gemini_stream() -> None:
    """
    Patch Gemini client to auto-track streaming usage.

    This patches the synchronous `Models.generate_content_stream` method.
    """
    global _patched_stream_sync, _original_generate_content_stream

    if _patched_stream_sync:
        logger.debug("Gemini sync stream already patched, skipping")
        return

    try:
        from google.genai.models import Models
    except ImportError:
        logger.warning(
            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
        )
        return

    if not hasattr(Models, "generate_content_stream"):
        logger.debug("Gemini Models has no generate_content_stream method, skipping")
        return

    original_generate_content_stream = Models.generate_content_stream
    _original_generate_content_stream = original_generate_content_stream

    @functools.wraps(original_generate_content_stream)
    def patched_generate_content_stream(
        self: Models, *args: Any, **kwargs: Any
    ) -> Any:
        stream = original_generate_content_stream(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or (args[0] if args else "unknown")

        config = kwargs.get("config")

        try:
            return _track_stream(stream, model, context, config)
        except Exception as e:
            logger.error(f"Error wrapping Gemini stream: {e}", exc_info=True)
            return stream

    Models.generate_content_stream = patched_generate_content_stream  # type: ignore[assignment]
    _patched_stream_sync = True
    logger.info("Gemini streaming SDK patched for auto-instrumentation")


# ============================================================================
# Async Patching
# ============================================================================


def patch_gemini_async() -> None:
    """
    Patch Gemini async client to auto-track usage.

    This patches the asynchronous `AsyncModels.generate_content` method.
    """
    global _patched_async, _original_generate_content_async

    if _patched_async:
        logger.debug("Gemini async already patched, skipping")
        return

    try:
        from google.genai.models import AsyncModels
    except ImportError:
        logger.warning(
            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
        )
        return

    original_generate_content = AsyncModels.generate_content
    _original_generate_content_async = original_generate_content

    @functools.wraps(original_generate_content)
    async def patched_generate_content(
        self: AsyncModels, *args: Any, **kwargs: Any
    ) -> Any:
        response = await original_generate_content(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or (args[0] if args else "unknown")

        config = kwargs.get("config")

        try:
            _track_generate_content_response(response, model, context, config)
        except Exception as e:
            logger.error(f"Error tracking Gemini async usage: {e}", exc_info=True)

        return response

    AsyncModels.generate_content = patched_generate_content  # type: ignore[assignment]
    _patched_async = True
    logger.info("Gemini async SDK patched for auto-instrumentation")


def patch_gemini_stream_async() -> None:
    """
    Patch Gemini async client to auto-track streaming usage.

    This patches the asynchronous `AsyncModels.generate_content_stream` method.
    """
    global _patched_stream_async, _original_generate_content_stream_async

    if _patched_stream_async:
        logger.debug("Gemini async stream already patched, skipping")
        return

    try:
        from google.genai.models import AsyncModels
    except ImportError:
        logger.warning(
            "Gemini SDK not installed. Install with: pip install 'fenra[gemini]'"
        )
        return

    if not hasattr(AsyncModels, "generate_content_stream"):
        logger.debug(
            "Gemini AsyncModels has no generate_content_stream method, skipping"
        )
        return

    original_generate_content_stream = AsyncModels.generate_content_stream
    _original_generate_content_stream_async = original_generate_content_stream

    @functools.wraps(original_generate_content_stream)
    async def patched_generate_content_stream(
        self: AsyncModels, *args: Any, **kwargs: Any
    ) -> Any:
        stream = await original_generate_content_stream(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or (args[0] if args else "unknown")

        config = kwargs.get("config")

        try:
            return _track_stream_async(stream, model, context, config)
        except Exception as e:
            logger.error(f"Error wrapping Gemini async stream: {e}", exc_info=True)
            return stream

    AsyncModels.generate_content_stream = patched_generate_content_stream  # type: ignore[assignment]
    _patched_stream_async = True
    logger.info("Gemini async streaming SDK patched for auto-instrumentation")


# ============================================================================
# Unpatch Functions
# ============================================================================


def unpatch_gemini() -> None:
    """Restore original Gemini Models.generate_content method."""
    global _patched_sync, _original_generate_content

    if not _patched_sync or _original_generate_content is None:
        return

    try:
        from google.genai.models import Models

        Models.generate_content = _original_generate_content  # type: ignore[assignment]
        _patched_sync = False
        _original_generate_content = None
        logger.info("Gemini SDK unpatched")
    except ImportError:
        pass


def unpatch_gemini_stream() -> None:
    """Restore original Gemini Models.generate_content_stream method."""
    global _patched_stream_sync, _original_generate_content_stream

    if not _patched_stream_sync or _original_generate_content_stream is None:
        return

    try:
        from google.genai.models import Models

        Models.generate_content_stream = _original_generate_content_stream  # type: ignore[assignment]
        _patched_stream_sync = False
        _original_generate_content_stream = None
        logger.info("Gemini streaming SDK unpatched")
    except ImportError:
        pass


def unpatch_gemini_async() -> None:
    """Restore original Gemini AsyncModels.generate_content method."""
    global _patched_async, _original_generate_content_async

    if not _patched_async or _original_generate_content_async is None:
        return

    try:
        from google.genai.models import AsyncModels

        AsyncModels.generate_content = _original_generate_content_async  # type: ignore[assignment]
        _patched_async = False
        _original_generate_content_async = None
        logger.info("Gemini async SDK unpatched")
    except ImportError:
        pass


def unpatch_gemini_stream_async() -> None:
    """Restore original Gemini AsyncModels.generate_content_stream method."""
    global _patched_stream_async, _original_generate_content_stream_async

    if not _patched_stream_async or _original_generate_content_stream_async is None:
        return

    try:
        from google.genai.models import AsyncModels

        AsyncModels.generate_content_stream = _original_generate_content_stream_async  # type: ignore[assignment]
        _patched_stream_async = False
        _original_generate_content_stream_async = None
        logger.info("Gemini async streaming SDK unpatched")
    except ImportError:
        pass


def unpatch_gemini_all() -> None:
    """Restore all original Gemini SDK methods."""
    unpatch_gemini()
    unpatch_gemini_stream()
    unpatch_gemini_async()
    unpatch_gemini_stream_async()
