import json
from typing import TYPE_CHECKING

import sentry_sdk
from sentry_sdk.ai._openai_responses_api import (
    _get_system_instructions,
    _is_system_instruction,
    _transform_system_instructions,
)
from sentry_sdk.ai.utils import (
    GEN_AI_ALLOWED_MESSAGE_ROLES,
    normalize_message_role,
    normalize_message_roles,
    set_data_normalized,
    truncate_and_annotate_messages,
)
from sentry_sdk.consts import SPANDATA
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing_utils import should_truncate_gen_ai_input
from sentry_sdk.utils import (
    event_from_exception,
    has_data_collection_enabled,
    safe_serialize,
)

if TYPE_CHECKING:
    from typing import Any, Optional, Union

    from agents import TResponseInputItem, Usage

    from sentry_sdk._types import TextPart

try:
    import agents
    from agents import Model
except ImportError:
    raise DidNotEnable("OpenAI Agents not installed")


def _capture_exception(exc: "Any") -> None:
    event, hint = event_from_exception(
        exc,
        client_options=sentry_sdk.get_client().options,
        mechanism={"type": "openai_agents", "handled": False},
    )
    sentry_sdk.capture_event(event, hint=hint)


def _set_agent_data(
    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "agents.Agent"
) -> None:
    set_on_span = (
        span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
    )

    set_on_span(
        SPANDATA.GEN_AI_SYSTEM, "openai"
    )  # See footnote for  https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-system for explanation why.

    set_on_span(SPANDATA.GEN_AI_AGENT_NAME, agent.name)

    if agent.model_settings.max_tokens:
        set_on_span(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, agent.model_settings.max_tokens)

    model_name: "Optional[str]" = None
    if isinstance(agent.model, Model) and hasattr(agent.model, "model"):
        model_name = agent.model.model
    elif isinstance(agent.model, str):
        model_name = agent.model
    elif hasattr(agent, "_sentry_request_model"):
        model_name = agent._sentry_request_model

    if model_name:
        set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)

    if agent.model_settings.presence_penalty:
        set_on_span(
            SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
            agent.model_settings.presence_penalty,
        )

    if agent.model_settings.temperature:
        set_on_span(
            SPANDATA.GEN_AI_REQUEST_TEMPERATURE, agent.model_settings.temperature
        )

    if agent.model_settings.top_p:
        set_on_span(SPANDATA.GEN_AI_REQUEST_TOP_P, agent.model_settings.top_p)

    if agent.model_settings.frequency_penalty:
        set_on_span(
            SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
            agent.model_settings.frequency_penalty,
        )


def _set_usage_data(
    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", usage: "Usage"
) -> None:
    set_on_span = (
        span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
    )
    set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
    set_on_span(
        SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED,
        usage.input_tokens_details.cached_tokens,
    )
    set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
    set_on_span(
        SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
        usage.output_tokens_details.reasoning_tokens,
    )
    set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)


def _set_input_data(
    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
    get_response_kwargs: "dict[str, Any]",
) -> None:
    client = sentry_sdk.get_client()
    if has_data_collection_enabled(client.options):
        if not client.options["data_collection"]["gen_ai"]["inputs"]:
            return
    elif not should_send_default_pii():
        return
    request_messages = []

    messages: "str | list[TResponseInputItem]" = get_response_kwargs.get("input", [])

    instructions_text_parts: "list[TextPart]" = []
    explicit_instructions = get_response_kwargs.get("system_instructions")
    if explicit_instructions is not None:
        instructions_text_parts.append(
            {
                "type": "text",
                "content": explicit_instructions,
            }
        )

    system_instructions = _get_system_instructions(messages)

    instructions_text_parts += _transform_system_instructions(system_instructions)

    if len(instructions_text_parts) > 0:
        if isinstance(span, StreamedSpan):
            span.set_attribute(
                SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
                json.dumps(instructions_text_parts),
            )
        else:
            span.set_data(
                SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
                json.dumps(instructions_text_parts),
            )

    non_system_messages = [
        message
        for message in messages
        if not _is_system_instruction(message)  # type: ignore[arg-type]
    ]
    for message in non_system_messages:
        if "role" in message:
            normalized_role = normalize_message_role(message.get("role"))  # type: ignore
            content = message.get("content")  # type: ignore
            request_messages.append(
                {
                    "role": normalized_role,
                    "content": (
                        [{"type": "text", "text": content}]
                        if isinstance(content, str)
                        else content
                    ),
                }
            )
        else:
            if message.get("type") == "function_call":  # type: ignore
                request_messages.append(
                    {
                        "role": GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT,
                        "content": [message],
                    }
                )
            elif message.get("type") == "function_call_output":  # type: ignore
                request_messages.append(
                    {
                        "role": GEN_AI_ALLOWED_MESSAGE_ROLES.TOOL,
                        "content": [message],
                    }
                )

    normalized_messages = normalize_message_roles(request_messages)
    scope = sentry_sdk.get_current_scope()
    messages_data = (
        truncate_and_annotate_messages(normalized_messages, span, scope)
        if should_truncate_gen_ai_input(client.options)
        else normalized_messages
    )
    if messages_data is not None:
        set_data_normalized(
            span,
            SPANDATA.GEN_AI_REQUEST_MESSAGES,
            messages_data,
            unpack=False,
        )


def _set_output_data(
    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any"
) -> None:
    client = sentry_sdk.get_client()
    record_outputs = False
    if has_data_collection_enabled(client.options):
        record_outputs = client.options["data_collection"]["gen_ai"]["outputs"]
    elif should_send_default_pii():
        record_outputs = True

    if not record_outputs:
        return

    output_messages: "dict[str, list[Any]]" = {
        "response": [],
        "tool": [],
    }

    for output in result.output:
        if output.type == "function_call":
            output_messages["tool"].append(output.dict())
        elif output.type == "message":
            for output_message in output.content:
                try:
                    output_messages["response"].append(output_message.text)
                except AttributeError:
                    # Unknown output message type, just return the json
                    output_messages["response"].append(output_message.dict())

    if record_outputs and len(output_messages["tool"]) > 0:
        if isinstance(span, StreamedSpan):
            span.set_attribute(
                SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
                safe_serialize(output_messages["tool"]),
            )
        else:
            span.set_data(
                SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
                safe_serialize(output_messages["tool"]),
            )

    if record_outputs and len(output_messages["response"]) > 0:
        set_data_normalized(
            span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
        )