"""OpenTelemetry Google Generative AI API instrumentation"""

import json
import logging
import os
from collections import defaultdict
from typing import AsyncGenerator, Callable, Collection, Generator

from google.genai import types
from opentelemetry import context as context_api
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
from opentelemetry.trace import Span, Status, StatusCode, Tracer, get_tracer
from wrapt import wrap_function_wrapper

from lmnr.opentelemetry_lib.opentelemetry.instrumentation.shared.utils import (
    safe_start_span,
)
from lmnr.opentelemetry_lib.tracing.context import (
    get_event_attributes_from_context,
)
from lmnr.sdk.laminar import Laminar
from lmnr.sdk.utils import json_dumps

from .config import (
    Config,
)
from .schema_utils import SchemaJSONEncoder, process_schema
from .utils import (
    content_union_to_dict,
    dont_throw,
    get_content,
    merge_text_parts,
    process_content_union,
    process_stream_chunk,
    set_span_attribute,
    to_dict,
    with_tracer_wrapper,
)

logger = logging.getLogger(__name__)

_instruments = ("google-genai >= 1.0.0",)

WRAPPED_METHODS = [
    {
        "package": "google.genai.models",
        "object": "Models",
        "method": "generate_content",
        "span_name": "gemini.generate_content",
        "is_streaming": False,
        "is_async": False,
    },
    {
        "package": "google.genai.models",
        "object": "AsyncModels",
        "method": "generate_content",
        "span_name": "gemini.generate_content",
        "is_streaming": False,
        "is_async": True,
    },
    {
        "package": "google.genai.models",
        "object": "Models",
        "method": "generate_content_stream",
        "span_name": "gemini.generate_content_stream",
        "is_streaming": True,
        "is_async": False,
    },
    {
        "package": "google.genai.models",
        "object": "AsyncModels",
        "method": "generate_content_stream",
        "span_name": "gemini.generate_content_stream",
        "is_streaming": True,
        "is_async": True,
    },
]


def should_send_prompts():
    return (
        os.getenv("LAMINAR_TRACE_CONTENT") or "true"
    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")


@dont_throw
def _set_request_attributes(span, args, kwargs):
    config_dict = to_dict(kwargs.get("config", {}))
    set_span_attribute(
        span, gen_ai_attributes.GEN_AI_REQUEST_MODEL, kwargs.get("model")
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_REQUEST_TEMPERATURE,
        config_dict.get("temperature"),
    )
    set_span_attribute(
        span, gen_ai_attributes.GEN_AI_REQUEST_TOP_P, config_dict.get("top_p")
    )
    set_span_attribute(
        span, gen_ai_attributes.GEN_AI_REQUEST_TOP_K, config_dict.get("top_k")
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_REQUEST_CHOICE_COUNT,
        config_dict.get("candidate_count"),
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_REQUEST_MAX_TOKENS,
        config_dict.get("max_output_tokens"),
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_REQUEST_STOP_SEQUENCES,
        config_dict.get("stop_sequences"),
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_REQUEST_FREQUENCY_PENALTY,
        config_dict.get("frequency_penalty"),
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_REQUEST_PRESENCE_PENALTY,
        config_dict.get("presence_penalty"),
    )
    set_span_attribute(
        span, gen_ai_attributes.GEN_AI_REQUEST_SEED, config_dict.get("seed")
    )

    if schema := config_dict.get("response_schema"):
        try:
            set_span_attribute(
                span,
                "gen_ai.request.structured_output_schema",
                json.dumps(process_schema(schema), cls=SchemaJSONEncoder),
            )
        except Exception:
            pass
    elif json_schema := config_dict.get("response_json_schema"):
        try:
            set_span_attribute(
                span,
                "gen_ai.request.structured_output_schema",
                json_dumps(json_schema),
            )
        except Exception:
            pass

    tools: list[types.FunctionDeclaration] = []
    arg_tools = config_dict.get("tools", kwargs.get("tools"))
    if arg_tools:
        for tool in arg_tools:
            if isinstance(tool, types.Tool):
                tools.extend(tool.function_declarations or [])
            elif isinstance(tool, dict) and isinstance(
                tool.get("function_declarations"), list
            ):
                tools.extend(tool.get("function_declarations", []))
            elif isinstance(tool, Callable):
                tools.append(types.FunctionDeclaration.from_callable(tool))

    if should_send_prompts():
        messages = []
        system_instruction = config_dict.get("system_instruction")
        if system_instruction:
            msg = content_union_to_dict(system_instruction, default_role="system")
            msg["role"] = "system"
            messages.append(msg)

        contents = kwargs.get("contents", [])
        if not isinstance(contents, list):
            contents = [contents]
        for content in contents:
            messages.append(content_union_to_dict(content))

        set_span_attribute(span, "gen_ai.input.messages", json_dumps(messages))
    if tools:
        span.set_attribute(
            "gen_ai.tool.definitions",
            json_dumps([to_dict(tool) for tool in tools]),
        )


@dont_throw
def _set_response_attributes(span, response: types.GenerateContentResponse):
    candidates = response.candidates or []
    set_span_attribute(
        span, gen_ai_attributes.GEN_AI_RESPONSE_ID, to_dict(response).get("response_id")
    )
    set_span_attribute(
        span,
        gen_ai_attributes.GEN_AI_RESPONSE_MODEL,
        to_dict(response).get("model_version"),
    )

    if response.usage_metadata:
        usage_dict = to_dict(response.usage_metadata)
        candidates_token_count = usage_dict.get("candidates_token_count")
        # unlike OpenAI, and unlike input cached tokens, thinking tokens are
        # not counted as part of candidates token count, so we need to add them
        # separately for consistency with other instrumentations
        thoughts_token_count = usage_dict.get("thoughts_token_count")
        output_token_count = (
            (candidates_token_count or 0) + (thoughts_token_count or 0)
            if candidates_token_count is not None or thoughts_token_count is not None
            else None
        )
        set_span_attribute(
            span,
            gen_ai_attributes.GEN_AI_USAGE_INPUT_TOKENS,
            usage_dict.get("prompt_token_count"),
        )
        set_span_attribute(
            span,
            gen_ai_attributes.GEN_AI_USAGE_OUTPUT_TOKENS,
            output_token_count,
        )
        set_span_attribute(
            span,
            "llm.usage.total_tokens",
            usage_dict.get("total_token_count"),
        )
        set_span_attribute(
            span,
            "gen_ai.usage.cache_read_input_tokens",
            usage_dict.get("cached_content_token_count"),
        )
        set_span_attribute(
            span,
            "gen_ai.usage.reasoning_tokens",
            thoughts_token_count,
        )

    if should_send_prompts():
        set_span_attribute(span, "gen_ai.completion.0.role", "model")
        candidates_list = candidates if isinstance(candidates, list) else [candidates]
        i = 0
        for candidate in candidates_list:
            has_content = False
            processed_content = process_content_union(candidate.content)
            content_payload = get_content(processed_content)
            if isinstance(content_payload, dict):
                content_payload = [content_payload]

            set_span_attribute(span, f"gen_ai.completion.{i}.role", "model")
            if content_payload:
                has_content = True
                set_span_attribute(
                    span,
                    f"gen_ai.completion.{i}.content",
                    (
                        content_payload
                        if isinstance(content_payload, str)
                        else json_dumps(content_payload)
                    ),
                )
            blocks = (
                processed_content
                if isinstance(processed_content, list)
                else [processed_content]
            )

            tool_call_index = 0
            for block in blocks:
                block_dict = to_dict(block)
                if not block_dict.get("function_call"):
                    continue
                function_call = to_dict(block_dict.get("function_call", {}))
                has_content = True
                set_span_attribute(
                    span,
                    f"gen_ai.completion.{i}.tool_calls.{tool_call_index}.name",
                    function_call.get("name"),
                )
                set_span_attribute(
                    span,
                    f"gen_ai.completion.{i}.tool_calls.{tool_call_index}.id",
                    (
                        function_call.get("id")
                        if function_call.get("id") is not None
                        else function_call.get("name")
                    ),  # google genai doesn't support tool call ids
                )
                set_span_attribute(
                    span,
                    f"gen_ai.completion.{i}.tool_calls.{tool_call_index}.arguments",
                    json_dumps(function_call.get("arguments")),
                )
                tool_call_index += 1
            if has_content:
                i += 1


@dont_throw
def _set_raw_response_attribute(
    span, response: types.GenerateContentResponse, record_raw_response: bool = False
):
    set_span_attribute(
        span,
        "gen_ai.output.messages",
        json_dumps(
            [
                candidate.model_dump(mode="json", exclude_unset=True)
                for candidate in (response.candidates or [])
            ]
        ),
    )
    if record_raw_response:
        try:
            set_span_attribute(
                span,
                "lmnr.sdk.raw.response",
                response.model_dump_json(),
            )
        except Exception:
            logger.debug("Failed to set lmnr.sdk.raw.response attribute", exc_info=True)


@dont_throw
def _build_from_streaming_response(
    span: Span,
    response: Generator[types.GenerateContentResponse, None, None],
    record_raw_response: bool = False,
) -> Generator[types.GenerateContentResponse, None, None]:
    final_parts = []
    role = "model"
    aggregated_usage_metadata = defaultdict(int)
    model_version = None
    for chunk in response:
        try:
            span.add_event("llm.content.completion.chunk")
        except Exception:
            pass
        # Important: do all processing in a separate sync function, that is
        # wrapped in @dont_throw. If we did it here, the @dont_throw on top of
        # this function would not be able to catch the errors, as they are
        # raised later, after the generator is returned, and when it is being
        # consumed.
        chunk_result = process_stream_chunk(
            chunk,
            role,
            model_version,
            aggregated_usage_metadata,
            final_parts,
        )
        # even though process_stream_chunk can't return None, the result can be
        # None, if the processing throws an error (see @dont_throw)
        if chunk_result:
            role = chunk_result["role"]
            model_version = chunk_result["model_version"]
        yield chunk

    try:
        compound_response = types.GenerateContentResponse(
            candidates=[
                {
                    "content": {
                        "parts": merge_text_parts(final_parts),
                        "role": role,
                    },
                }
            ],
            usage_metadata=types.GenerateContentResponseUsageMetadataDict(
                **aggregated_usage_metadata
            ),
            model_version=model_version,
        )
        if span.is_recording():
            _set_raw_response_attribute(
                span, compound_response, record_raw_response=record_raw_response
            )
            _set_response_attributes(span, compound_response)
    finally:
        if span.is_recording():
            span.end()


@dont_throw
async def _abuild_from_streaming_response(
    span: Span,
    response: AsyncGenerator[types.GenerateContentResponse, None],
    record_raw_response: bool = False,
) -> AsyncGenerator[types.GenerateContentResponse, None]:
    final_parts = []
    role = "model"
    aggregated_usage_metadata = defaultdict(int)
    model_version = None
    async for chunk in response:
        try:
            span.add_event("llm.content.completion.chunk")
        except Exception:
            pass
        # Important: do all processing in a separate sync function, that is
        # wrapped in @dont_throw. If we did it here, the @dont_throw on top of
        # this function would not be able to catch the errors, as they are
        # raised later, after the generator is returned, and when it is being
        # consumed.
        chunk_result = process_stream_chunk(
            chunk,
            role,
            model_version,
            aggregated_usage_metadata,
            final_parts,
        )
        # even though process_stream_chunk can't return None, the result can be
        # None, if the processing throws an error (see @dont_throw)
        if chunk_result:
            role = chunk_result["role"]
            model_version = chunk_result["model_version"]
        yield chunk

    try:
        compound_response = types.GenerateContentResponse(
            candidates=[
                {
                    "content": {
                        "parts": merge_text_parts(final_parts),
                        "role": role,
                    },
                }
            ],
            usage_metadata=types.GenerateContentResponseUsageMetadataDict(
                **aggregated_usage_metadata
            ),
            model_version=model_version,
        )
        if span.is_recording():
            _set_raw_response_attribute(
                span, compound_response, record_raw_response=record_raw_response
            )
            _set_response_attributes(span, compound_response)
    finally:
        if span.is_recording():
            span.end()


@with_tracer_wrapper
def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
        return wrapped(*args, **kwargs)

    span = safe_start_span(
        name=to_wrap.get("span_name"),
        attributes={"gen_ai.system": "gemini"},
        span_type="LLM",
    )
    if not span:
        logger.warning("Failed to start span for google genai")
        return wrapped(*args, **kwargs)

    _set_request_attributes(span, args, kwargs)

    try:
        # Check for rollout mode and apply caching/overrides
        from lmnr.sdk.rollout_control import is_rollout_mode

        is_rollout = is_rollout_mode()
    except Exception:
        is_rollout = False

    try:
        if is_rollout:
            from lmnr.opentelemetry_lib.opentelemetry.instrumentation.google_genai.rollout import (
                get_google_genai_rollout_wrapper,
            )

            rollout_wrapper = get_google_genai_rollout_wrapper()
            if rollout_wrapper:
                with Laminar.use_span(span):
                    response = rollout_wrapper.wrap_generate_content(
                        wrapped,
                        instance,
                        args,
                        kwargs,
                        is_streaming=to_wrap.get("is_streaming", False),
                        is_async=False,
                    )
            else:
                response = wrapped(*args, **kwargs)
        else:
            response = wrapped(*args, **kwargs)

        if to_wrap.get("is_streaming"):
            return _build_from_streaming_response(
                span, response, record_raw_response=is_rollout
            )
        if span.is_recording():
            _set_raw_response_attribute(span, response, record_raw_response=is_rollout)
            _set_response_attributes(span, response)
        span.end()
        return response
    except Exception as e:
        attributes = get_event_attributes_from_context()
        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
        span.record_exception(e, attributes=attributes)
        span.set_status(Status(StatusCode.ERROR, str(e)))
        span.end()
        raise


@with_tracer_wrapper
async def _awrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
        return await wrapped(*args, **kwargs)

    span = safe_start_span(
        name=to_wrap.get("span_name"),
        attributes={"gen_ai.system": "gemini"},
        span_type="LLM",
    )
    if not span:
        logger.warning("Failed to start span for async google genai")
        return await wrapped(*args, **kwargs)

    _set_request_attributes(span, args, kwargs)

    try:
        # Check for rollout mode and apply caching/overrides
        from lmnr.sdk.rollout_control import is_rollout_mode

        is_rollout = is_rollout_mode()
    except Exception:
        is_rollout = False

    try:
        if is_rollout:
            from lmnr.opentelemetry_lib.opentelemetry.instrumentation.google_genai.rollout import (
                get_google_genai_rollout_wrapper,
            )

            rollout_wrapper = get_google_genai_rollout_wrapper()
            if rollout_wrapper:
                # For async, we need to handle both sync and async wrapped functions
                # In rollout mode, wrapped might return cached (sync) or live (async)
                with Laminar.use_span(span):
                    result = rollout_wrapper.wrap_generate_content(
                        wrapped,
                        instance,
                        args,
                        kwargs,
                        is_streaming=to_wrap.get("is_streaming", False),
                        is_async=True,
                    )
                # If result is a coroutine or async generator, await/iterate it
                import inspect

                if inspect.iscoroutine(result):
                    response = await result
                elif inspect.isasyncgen(result):
                    # It's an async generator (cached streaming response)
                    response = result
                else:
                    # It's a sync response (cached non-streaming)
                    response = result
            else:
                response = await wrapped(*args, **kwargs)
        else:
            response = await wrapped(*args, **kwargs)

        if to_wrap.get("is_streaming"):
            return _abuild_from_streaming_response(
                span, response, record_raw_response=is_rollout
            )
        else:
            if span.is_recording():
                _set_raw_response_attribute(
                    span, response, record_raw_response=is_rollout
                )
                _set_response_attributes(span, response)

            span.end()
            return response
    except Exception as e:
        attributes = get_event_attributes_from_context()
        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
        span.record_exception(e, attributes=attributes)
        span.set_status(Status(StatusCode.ERROR, str(e)))
        span.end()
        raise


class GoogleGenAiSdkInstrumentor(BaseInstrumentor):
    """An instrumentor for Google GenAI's client library."""

    def __init__(
        self,
        exception_logger=None,
        upload_base64_image=None,
        convert_image_to_openai_format=True,
    ):
        super().__init__()
        Config.exception_logger = exception_logger
        Config.upload_base64_image = upload_base64_image
        Config.convert_image_to_openai_format = convert_image_to_openai_format

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs):
        tracer_provider = kwargs.get("tracer_provider")
        tracer = get_tracer(__name__, "0.0.1a1", tracer_provider)

        for wrapped_method in WRAPPED_METHODS:
            wrap_function_wrapper(
                wrapped_method.get("package"),
                f"{wrapped_method.get('object')}.{wrapped_method.get('method')}",
                (
                    _awrap(tracer, wrapped_method)
                    if wrapped_method.get("is_async")
                    else _wrap(tracer, wrapped_method)
                ),
            )

    def _uninstrument(self, **kwargs):
        for wrapped_method in WRAPPED_METHODS:
            unwrap(
                f"{wrapped_method.get('package')}.{wrapped_method.get('object')}",
                wrapped_method.get("method"),
            )
