"""Google Gemini client wrapper for automatic logging."""

from __future__ import annotations

import time
import uuid
from functools import wraps
from typing import TYPE_CHECKING, Any, Iterator

if TYPE_CHECKING:
    from google.generativeai import GenerativeModel
    from google.generativeai.types import GenerateContentResponse
    from aitracer.client import AITracer


def wrap_gemini_model(model: "GenerativeModel", tracer: "AITracer") -> "GenerativeModel":
    """
    Wrap a Google Gemini GenerativeModel to automatically log all API calls.

    Args:
        model: GenerativeModel instance.
        tracer: AITracer instance.

    Returns:
        Wrapped GenerativeModel (same instance, modified in place).

    Example:
        >>> import google.generativeai as genai
        >>> from aitracer import AITracer
        >>> from aitracer.wrappers import wrap_gemini_model
        >>>
        >>> genai.configure(api_key="your-api-key")
        >>> model = genai.GenerativeModel("gemini-1.5-flash")
        >>> tracer = AITracer(api_key="your-aitracer-key")
        >>> model = wrap_gemini_model(model, tracer)
        >>>
        >>> response = model.generate_content("Hello!")
    """
    model_name = getattr(model, "model_name", "gemini-unknown")

    # Wrap generate_content
    original_generate = model.generate_content

    @wraps(original_generate)
    def wrapped_generate_content(
        contents: Any,
        *args: Any,
        stream: bool = False,
        **kwargs: Any,
    ) -> Any:
        """Wrapped generate_content method."""
        start_time = time.time()
        span_id = str(uuid.uuid4())

        try:
            response = original_generate(contents, *args, stream=stream, **kwargs)

            if stream:
                return _wrap_stream_response(
                    response=response,
                    tracer=tracer,
                    model=model_name,
                    contents=contents,
                    start_time=start_time,
                    span_id=span_id,
                )
            else:
                latency_ms = int((time.time() - start_time) * 1000)
                _log_generation(
                    tracer=tracer,
                    model=model_name,
                    contents=contents,
                    response=response,
                    latency_ms=latency_ms,
                    span_id=span_id,
                )
                return response

        except Exception as e:
            latency_ms = int((time.time() - start_time) * 1000)
            tracer.log(
                model=model_name,
                provider="google",
                input_data={"contents": _serialize_contents(contents)},
                output_data=None,
                latency_ms=latency_ms,
                status="error",
                error_message=str(e),
                span_id=span_id,
            )
            raise

    model.generate_content = wrapped_generate_content  # type: ignore

    # Wrap generate_content_async if available
    if hasattr(model, "generate_content_async"):
        original_generate_async = model.generate_content_async

        @wraps(original_generate_async)
        async def wrapped_generate_content_async(
            contents: Any,
            *args: Any,
            stream: bool = False,
            **kwargs: Any,
        ) -> Any:
            """Wrapped async generate_content method."""
            start_time = time.time()
            span_id = str(uuid.uuid4())

            try:
                response = await original_generate_async(
                    contents, *args, stream=stream, **kwargs
                )

                if stream:
                    return _wrap_async_stream_response(
                        response=response,
                        tracer=tracer,
                        model=model_name,
                        contents=contents,
                        start_time=start_time,
                        span_id=span_id,
                    )
                else:
                    latency_ms = int((time.time() - start_time) * 1000)
                    _log_generation(
                        tracer=tracer,
                        model=model_name,
                        contents=contents,
                        response=response,
                        latency_ms=latency_ms,
                        span_id=span_id,
                    )
                    return response

            except Exception as e:
                latency_ms = int((time.time() - start_time) * 1000)
                tracer.log(
                    model=model_name,
                    provider="google",
                    input_data={"contents": _serialize_contents(contents)},
                    output_data=None,
                    latency_ms=latency_ms,
                    status="error",
                    error_message=str(e),
                    span_id=span_id,
                )
                raise

        model.generate_content_async = wrapped_generate_content_async  # type: ignore

    return model


def _wrap_stream_response(
    response: Iterator["GenerateContentResponse"],
    tracer: "AITracer",
    model: str,
    contents: Any,
    start_time: float,
    span_id: str,
) -> Iterator["GenerateContentResponse"]:
    """Wrap streaming response to log after completion."""
    content_parts: list[str] = []
    input_tokens = 0
    output_tokens = 0
    last_response = None

    try:
        for chunk in response:
            last_response = chunk

            # Accumulate content
            if chunk.text:
                content_parts.append(chunk.text)

            # Get usage if available
            if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
                input_tokens = getattr(chunk.usage_metadata, "prompt_token_count", 0) or 0
                output_tokens = getattr(chunk.usage_metadata, "candidates_token_count", 0) or 0

            yield chunk

        # Log after stream completes
        latency_ms = int((time.time() - start_time) * 1000)
        full_content = "".join(content_parts)

        tracer.log(
            model=model,
            provider="google",
            input_data={"contents": _serialize_contents(contents)},
            output_data={"content": full_content},
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            status="success",
            span_id=span_id,
        )

    except Exception as e:
        latency_ms = int((time.time() - start_time) * 1000)
        tracer.log(
            model=model,
            provider="google",
            input_data={"contents": _serialize_contents(contents)},
            output_data=None,
            latency_ms=latency_ms,
            status="error",
            error_message=str(e),
            span_id=span_id,
        )
        raise


async def _wrap_async_stream_response(
    response: Any,
    tracer: "AITracer",
    model: str,
    contents: Any,
    start_time: float,
    span_id: str,
) -> Any:
    """Wrap async streaming response to log after completion."""
    content_parts: list[str] = []
    input_tokens = 0
    output_tokens = 0

    try:
        async for chunk in response:
            # Accumulate content
            if chunk.text:
                content_parts.append(chunk.text)

            # Get usage if available
            if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
                input_tokens = getattr(chunk.usage_metadata, "prompt_token_count", 0) or 0
                output_tokens = getattr(chunk.usage_metadata, "candidates_token_count", 0) or 0

            yield chunk

        # Log after stream completes
        latency_ms = int((time.time() - start_time) * 1000)
        full_content = "".join(content_parts)

        tracer.log(
            model=model,
            provider="google",
            input_data={"contents": _serialize_contents(contents)},
            output_data={"content": full_content},
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=latency_ms,
            status="success",
            span_id=span_id,
        )

    except Exception as e:
        latency_ms = int((time.time() - start_time) * 1000)
        tracer.log(
            model=model,
            provider="google",
            input_data={"contents": _serialize_contents(contents)},
            output_data=None,
            latency_ms=latency_ms,
            status="error",
            error_message=str(e),
            span_id=span_id,
        )
        raise


def _log_generation(
    tracer: "AITracer",
    model: str,
    contents: Any,
    response: "GenerateContentResponse",
    latency_ms: int,
    span_id: str,
) -> None:
    """Log a non-streaming generation."""
    # Extract response data
    output_content = None
    try:
        output_content = response.text
    except Exception:
        # Response may not have text (e.g., safety blocked)
        if response.candidates:
            parts = []
            for candidate in response.candidates:
                if candidate.content and candidate.content.parts:
                    for part in candidate.content.parts:
                        if hasattr(part, "text"):
                            parts.append(part.text)
            output_content = "".join(parts) if parts else None

    # Extract token usage
    input_tokens = 0
    output_tokens = 0
    if hasattr(response, "usage_metadata") and response.usage_metadata:
        input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0) or 0
        output_tokens = getattr(response.usage_metadata, "candidates_token_count", 0) or 0

    tracer.log(
        model=model,
        provider="google",
        input_data={"contents": _serialize_contents(contents)},
        output_data={"content": output_content},
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        latency_ms=latency_ms,
        status="success",
        span_id=span_id,
    )


def _serialize_contents(contents: Any) -> Any:
    """Serialize contents to JSON-compatible format."""
    if contents is None:
        return None

    if isinstance(contents, str):
        return contents

    if isinstance(contents, list):
        result = []
        for item in contents:
            result.append(_serialize_content_item(item))
        return result

    return _serialize_content_item(contents)


def _serialize_content_item(item: Any) -> Any:
    """Serialize a single content item."""
    if isinstance(item, str):
        return item

    if isinstance(item, dict):
        return item

    if hasattr(item, "to_dict"):
        return item.to_dict()

    if hasattr(item, "model_dump"):
        return item.model_dump()

    if hasattr(item, "__dict__"):
        return {k: v for k, v in item.__dict__.items() if not k.startswith("_")}

    return str(item)
