"""
Google Generative AI SDK instrumentation.

Patches Google GenAI SDK methods to capture:
- GenerativeModel.generate_content (sync and async)
- GenerativeModel.generate_content_stream

Based on ARCHITECTURE_V2.md Section 6.3.
"""

from __future__ import annotations

import logging
import time
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, TypeVar

import wrapt

from risicare.integrations._base import scrub_sensitive

logger = logging.getLogger(__name__)

T = TypeVar("T")

_instrumented = False


def instrument_google(module: Any) -> None:
    """
    Apply instrumentation to Google Generative AI module.

    This patches:
    - google.generativeai.GenerativeModel.generate_content
    - google.generativeai.GenerativeModel.generate_content_async
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Patch sync generate_content
        wrapt.wrap_function_wrapper(
            module,
            "GenerativeModel.generate_content",
            _wrap_generate_content,
        )

        # Patch async generate_content
        wrapt.wrap_function_wrapper(
            module,
            "GenerativeModel.generate_content_async",
            _wrap_async_generate_content,
        )

        _instrumented = True
        logger.debug("Instrumented Google Generative AI SDK")

    except Exception as e:
        logger.warning(f"Failed to instrument Google GenAI SDK: {e}")


def _get_tracer() -> Optional[Any]:
    """Get the Risicare tracer if available."""
    try:
        from risicare.tracer import get_tracer

        return get_tracer()
    except ImportError:
        return None


def _get_config() -> Optional[Any]:
    """Get the Risicare client config if available."""
    try:
        from risicare.client import get_client

        client = get_client()
        return client.config if client else None
    except ImportError:
        return None


def _should_trace_content() -> bool:
    """Check if content should be traced."""
    config = _get_config()
    return config.trace_content if config else False


def _wrap_generate_content(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for GenerativeModel.generate_content (sync)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model_name = getattr(instance, "model_name", "unknown")
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"google.generate_content/{model_name}",
            kind=SpanKind.LLM_CALL,
            attributes={
                "gen_ai.system": "google",
                "gen_ai.request.model": model_name,
                "gen_ai.request.stream": False,
            },
        ) as span:
            # Capture input
            contents = args[0] if args else kwargs.get("contents")
            if trace_content and contents:
                if isinstance(contents, str):
                    span.set_attribute(
                        "gen_ai.prompt.content",
                        contents[:10000] if len(contents) > 10000 else contents,
                    )
                elif isinstance(contents, list):
                    span.set_attribute("gen_ai.prompt.parts", len(contents))

            try:
                start_time = time.perf_counter()
                response = wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000

                span.set_attribute("gen_ai.latency_ms", latency_ms)
                _capture_response(span, response, trace_content)

                return response

            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"google.generate_content/{model_name}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "google",
            "gen_ai.request.model": model_name,
            "gen_ai.request.stream": True,
        },
    )

    # Capture input
    contents = args[0] if args else kwargs.get("contents")
    if trace_content and contents:
        if isinstance(contents, str):
            span.set_attribute(
                "gen_ai.prompt.content",
                contents[:10000] if len(contents) > 10000 else contents,
            )
        elif isinstance(contents, list):
            span.set_attribute("gen_ai.prompt.parts", len(contents))

    try:
        start_time = time.perf_counter()
        result = wrapped(*args, **kwargs)
        return _wrap_stream_sync(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


async def _wrap_async_generate_content(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for GenerativeModel.generate_content_async."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return await wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model_name = getattr(instance, "model_name", "unknown")
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"google.generate_content/{model_name}",
            kind=SpanKind.LLM_CALL,
            attributes={
                "gen_ai.system": "google",
                "gen_ai.request.model": model_name,
                "gen_ai.request.stream": False,
            },
        ) as span:
            # Capture input
            contents = args[0] if args else kwargs.get("contents")
            if trace_content and contents:
                if isinstance(contents, str):
                    span.set_attribute(
                        "gen_ai.prompt.content",
                        contents[:10000] if len(contents) > 10000 else contents,
                    )
                elif isinstance(contents, list):
                    span.set_attribute("gen_ai.prompt.parts", len(contents))

            try:
                start_time = time.perf_counter()
                response = await wrapped(*args, **kwargs)
                latency_ms = (time.perf_counter() - start_time) * 1000

                span.set_attribute("gen_ai.latency_ms", latency_ms)
                _capture_response(span, response, trace_content)

                return response

            except Exception as e:
                span.record_exception(e)
                span.set_attribute("error", True)
                span.set_attribute("error.type", type(e).__name__)
                span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
                raise

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"google.generate_content/{model_name}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "google",
            "gen_ai.request.model": model_name,
            "gen_ai.request.stream": True,
        },
    )

    # Capture input
    contents = args[0] if args else kwargs.get("contents")
    if trace_content and contents:
        if isinstance(contents, str):
            span.set_attribute(
                "gen_ai.prompt.content",
                contents[:10000] if len(contents) > 10000 else contents,
            )
        elif isinstance(contents, list):
            span.set_attribute("gen_ai.prompt.parts", len(contents))

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        return _wrap_stream_async(result, span, start_time, trace_content, tracer)
    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
        span.end()
        tracer.export_span(span)
        raise


def _wrap_stream_sync(
    stream: Iterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> Iterator[Any]:
    """Wrap a sync streaming response from Google GenAI.

    Google GenAI streaming yields GenerateContentResponse chunks.
    Each chunk may have partial text in candidates[0].content.parts[].text
    and usage_metadata on the final chunk.
    """
    chunk_count = 0
    content_buffer = []
    content_length = 0
    usage_data = {}

    try:
        for chunk in stream:
            chunk_count += 1

            # Accumulate content (bounded)
            if trace_content and content_length < 10000:
                candidates = getattr(chunk, "candidates", [])
                if candidates:
                    content = getattr(candidates[0], "content", None)
                    if content:
                        for part in getattr(content, "parts", []):
                            text = getattr(part, "text", "")
                            if text:
                                content_buffer.append(text)
                                content_length += len(text)

            # Capture usage from final chunk
            usage = getattr(chunk, "usage_metadata", None)
            if usage:
                usage_data = {
                    "prompt_tokens": getattr(usage, "prompt_token_count", 0),
                    "completion_tokens": getattr(usage, "candidates_token_count", 0),
                    "total_tokens": getattr(usage, "total_token_count", 0),
                }

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if usage_data:
            span.set_attribute("gen_ai.usage.prompt_tokens", usage_data.get("prompt_tokens", 0))
            span.set_attribute("gen_ai.usage.completion_tokens", usage_data.get("completion_tokens", 0))
            span.set_attribute("gen_ai.usage.total_tokens", usage_data.get("total_tokens", 0))

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus
        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


async def _wrap_stream_async(
    stream: AsyncIterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> AsyncIterator[Any]:
    """Wrap an async streaming response from Google GenAI."""
    chunk_count = 0
    content_buffer = []
    content_length = 0
    usage_data = {}

    try:
        async for chunk in stream:
            chunk_count += 1

            # Accumulate content (bounded)
            if trace_content and content_length < 10000:
                candidates = getattr(chunk, "candidates", [])
                if candidates:
                    content = getattr(candidates[0], "content", None)
                    if content:
                        for part in getattr(content, "parts", []):
                            text = getattr(part, "text", "")
                            if text:
                                content_buffer.append(text)
                                content_length += len(text)

            # Capture usage from final chunk
            usage = getattr(chunk, "usage_metadata", None)
            if usage:
                usage_data = {
                    "prompt_tokens": getattr(usage, "prompt_token_count", 0),
                    "completion_tokens": getattr(usage, "candidates_token_count", 0),
                    "total_tokens": getattr(usage, "total_token_count", 0),
                }

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if usage_data:
            span.set_attribute("gen_ai.usage.prompt_tokens", usage_data.get("prompt_tokens", 0))
            span.set_attribute("gen_ai.usage.completion_tokens", usage_data.get("completion_tokens", 0))
            span.set_attribute("gen_ai.usage.total_tokens", usage_data.get("total_tokens", 0))

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus
        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


def _capture_response(span: Any, response: Any, trace_content: bool) -> None:
    """Capture response data into span."""
    # Candidates
    candidates = getattr(response, "candidates", [])
    span.set_attribute("gen_ai.response.candidates", len(candidates))

    if candidates and trace_content:
        first = candidates[0]
        content = getattr(first, "content", None)
        if content:
            parts = getattr(content, "parts", [])
            if parts:
                text_parts = []
                for part in parts:
                    text = getattr(part, "text", "")
                    if text:
                        text_parts.append(text)
                if text_parts:
                    combined = "\n".join(text_parts)
                    span.set_attribute(
                        "gen_ai.completion.content",
                        combined[:10000] if len(combined) > 10000 else combined,
                    )

        finish_reason = getattr(first, "finish_reason", None)
        if finish_reason:
            span.set_attribute("gen_ai.completion.finish_reason", str(finish_reason))

    # Usage metadata
    usage = getattr(response, "usage_metadata", None)
    if usage:
        span.set_attribute(
            "gen_ai.usage.prompt_tokens",
            getattr(usage, "prompt_token_count", 0),
        )
        span.set_attribute(
            "gen_ai.usage.completion_tokens",
            getattr(usage, "candidates_token_count", 0),
        )
        span.set_attribute(
            "gen_ai.usage.total_tokens",
            getattr(usage, "total_token_count", 0),
        )
