import functools
import logging
from typing import Any

from opentelemetry.trace import Status, StatusCode

from ...contexts import GenerationContext
from ...sdk import VerseSDK
from .utils import (
    format_messages_for_tracing,
    get_completion_choices,
    get_finish_reason,
    get_request_parameters,
    get_stream_text,
    get_usage_from_response,
)


def configure_google(sdk: VerseSDK):
    """Configure Google GenAI to use the Verse SDK."""
    try:
        from google.genai.models import Models
    except ImportError:
        logging.warning("Google GenAI SDK not installed, skipping integration")
        return

    original_generate_content = Models.generate_content
    original_generate_content_stream = Models.generate_content_stream

    @functools.wraps(original_generate_content)
    def generate_content_with_tracing(self, *args, **kwargs):
        """Wrap synchronous generate_content with tracing."""
        model_name = kwargs.get("model", "unknown")
        contents = kwargs.get("contents", [])

        generation_ctx = sdk.create_generation(
            name="Google: Generate Content",
            model=model_name,
            usage=get_request_parameters(kwargs),
        )

        generation_ctx.provider("gcp.gemini").operation("generate_content").request_params(
            get_request_parameters(kwargs)
        )

        if contents:
            formatted_messages = format_messages_for_tracing(contents)
            generation_ctx.messages(formatted_messages)

        try:
            response = original_generate_content(self, *args, **kwargs)
            _process_response(response, generation_ctx, model_name)
            generation_ctx.span.end()
            return response

        except Exception as e:
            generation_ctx.error(e)
            generation_ctx.span.end()
            raise

    @functools.wraps(original_generate_content_stream)
    def generate_content_stream_with_tracing(self, *args, **kwargs):
        """Wrap streaming generate_content with tracing."""
        model_name = kwargs.get("model", "unknown")
        contents = kwargs.get("contents", [])

        generation_ctx = sdk.create_generation(
            name="Google: Generate Content Stream",
            model=model_name,
            usage=get_request_parameters(kwargs),
        )

        generation_ctx.provider("gcp.gemini").operation("generate_content").request_params(
            get_request_parameters(kwargs)
        )

        if contents:
            formatted_messages = format_messages_for_tracing(contents)
            generation_ctx.messages(formatted_messages)

        try:
            stream = original_generate_content_stream(self, *args, **kwargs)
            return _wrap_stream_response(stream, generation_ctx, sdk)

        except Exception as e:
            generation_ctx.error(e)
            generation_ctx.span.end()
            raise

    Models.generate_content = generate_content_with_tracing
    Models.generate_content_stream = generate_content_stream_with_tracing


def _process_response(response: Any, generation_ctx: GenerationContext, model_name: str):
    """Process non-streaming response and add attributes to span."""
    try:
        generation_ctx.model_used(model_name)

        finish_reason = get_finish_reason(response)
        if finish_reason:
            generation_ctx.reason(finish_reason)

        choices = get_completion_choices(response)
        if choices:
            generation_ctx.completions(choices)

        usage = get_usage_from_response(response)
        if usage:
            generation_ctx.usage(usage)

        generation_ctx.status(Status(StatusCode.OK))

    except Exception as e:
        logging.warning("Error processing Google GenAI response", exc_info=e)


def _wrap_stream_response(stream, generation_ctx: GenerationContext, sdk: VerseSDK):
    """Wrap streaming response to capture events."""
    try:
        for chunk in stream:
            text = get_stream_text(chunk)
            if text:
                sdk.event(
                    "gen_ai.stream.delta",
                    metadata={"content": text[:512]},
                    level="info",
                    span=generation_ctx.span,
                )

            yield chunk

        if hasattr(stream, "_result"):
            _process_response(stream._result, generation_ctx, generation_ctx.span.name)
        generation_ctx.span.end()

    except Exception as e:
        generation_ctx.error(e)
        generation_ctx.span.end()
        raise


async def _wrap_async_stream_response(
    stream, generation_ctx: GenerationContext, sdk: VerseSDK
):
    """Wrap async streaming response to capture events."""
    try:
        async for chunk in stream:
            text = get_stream_text(chunk)
            if text:
                sdk.event(
                    "gen_ai.stream.delta",
                    metadata={"content": text[:512]},
                    level="info",
                    span=generation_ctx.span,
                )

            yield chunk

        if hasattr(stream, "_result"):
            _process_response(stream._result, generation_ctx, generation_ctx.span.name)
        generation_ctx.span.end()

    except Exception as e:
        generation_ctx.error(e)
        generation_ctx.span.end()
        raise
