import logging
import time
from typing import Any, Dict, List, Optional, Tuple, Union

import google.adk.agents
from google.adk.agents import callback_context
from google.adk import models
from google.adk.tools import base_tool
from google.adk.tools import tool_context

import opik
from opik import context_storage
from opik.api_objects import opik_client, span, trace
from opik.types import DistributedTraceHeadersDict
from opik.decorator import span_creation_handler, arguments_helpers

from . import (
    helpers as adk_helpers,
    callback_context_info_extractors,
    patchers,
)
from .patchers import (
    litellm_wrappers,
    llm_response_wrapper,
)
from .patchers.adk_otel_tracer import llm_span_helpers
from .graph import mermaid_graph_builder

LOGGER = logging.getLogger(__name__)

SpanOrTraceData = Union[span.SpanData, trace.TraceData]


class OpikTracer:
    """
    Opik tracer for google-adk.
    """

    def __init__(
        self,
        name: Optional[str] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        project_name: Optional[str] = None,
        distributed_headers: Optional[DistributedTraceHeadersDict] = None,
    ):
        """
        Initialize OpikTracer.

        Arguments:
            name: The default name for root span or trace created by the tracer.
            tags: The default tags for all the traces and spans created by the tracer.
            metadata: The default metadata for all the traces and spans created by the tracer.
            project_name: The name of the project for tracing.
            distributed_headers: The distributed trace headers.
        """
        self.name = name
        self.tags = tags
        self.metadata = metadata or {}
        self.metadata["created_from"] = "google-adk"
        self.project_name = project_name
        self._distributed_headers = distributed_headers

        self._init_internal_attributes()

    def _init_internal_attributes(self) -> None:
        self._last_model_output: Optional[Dict[str, Any]] = None
        self._opik_client = opik_client.get_client_cached()
        # Track time-to-first-token: map span_id -> (request_start_time, first_token_time)
        self._ttft_tracking: Dict[str, Tuple[float, Optional[float]]] = {}

        patchers.patch_adk(
            self._opik_client, distributed_headers=self._distributed_headers
        )

    def _has_response_content(self, llm_response: models.LlmResponse) -> bool:
        """
        Check if the LlmResponse contains actual content (text or function calls).

        Arguments:
            llm_response: The LLM response to check.

        Returns:
            True if the response contains text content or function calls, False otherwise.
        """
        try:
            # Check the LlmResponse object directly for content structure
            if llm_response.content is not None and llm_response.content.parts:
                for part in llm_response.content.parts:
                    # Check for text content
                    if part.text and part.text.strip():
                        return True
                    # Check for function call content (tool calls)
                    if part.function_call:
                        return True
            return False
        except Exception as e:
            LOGGER.debug(
                f"Error checking LlmResponse.content.parts for TTFT: {e}",
                exc_info=True,
            )
            return False

    def _safe_ttft_tracking(
        self, span_id: Optional[str], pop: bool = False
    ) -> Tuple[Optional[float], Optional[float]]:
        """
        Safely retrieve time-to-first-token tracking data for a span.

        Arguments:
            span_id: The span ID to look up in tracking.
            pop: If True, remove the entry after fetching. If False, keep it.

        Returns:
            Tuple of (request_start_time, first_token_time). Returns (None, None) if
            span_id is None or not found in tracking.
        """
        if span_id is None or span_id not in self._ttft_tracking:
            return (None, None)
        if pop:
            return self._ttft_tracking.pop(span_id)
        return self._ttft_tracking[span_id]

    def flush(self) -> None:
        self._opik_client.flush()

    def before_agent_callback(
        self,
        callback_context: callback_context.CallbackContext,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        try:
            current_trace = context_storage.get_trace_data()
            current_span = context_storage.top_span_data()

            thread_id, session_metadata = (
                callback_context_info_extractors.try_get_session_info(callback_context)
            )

            agent_metadata = self.metadata.copy()
            agent_metadata["adk_invocation_id"] = callback_context.invocation_id
            agent_metadata.update(session_metadata)

            _try_add_agent_graph_to_metadata(agent_metadata, callback_context)

            if callback_context.user_content is not None:
                user_input = adk_helpers.convert_adk_base_model_to_dict(
                    callback_context.user_content
                )
            else:
                user_input = None

            name = self.name or callback_context.agent_name

            if current_span is not None:
                current_span.update(
                    name=name,
                    metadata={**agent_metadata},
                    input=user_input,
                    tags=self.tags,
                    project_name=self.project_name,
                )
            elif current_trace is not None:
                current_trace.update(
                    name=name,
                    metadata={**agent_metadata},
                    input=user_input,
                    tags=self.tags,
                    thread_id=thread_id,
                    project_name=self.project_name,
                )
            else:
                LOGGER.warning(
                    f"No current span or trace found in context for agent: {callback_context.agent_name}"
                )

        except Exception as e:
            LOGGER.error(f"Failed during before_agent_callback(): {e}", exc_info=True)

    def after_agent_callback(
        self,
        callback_context: callback_context.CallbackContext,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        try:
            output = self._last_model_output
            current_span = context_storage.top_span_data()
            current_trace = context_storage.get_trace_data()
            if current_span is not None:
                current_span.update(
                    output=output,
                    project_name=self.project_name,
                )
            elif current_trace is not None:
                current_trace.update(
                    output=output,
                    project_name=self.project_name,
                )
                self._last_model_output = None
            else:
                LOGGER.warning(
                    "No current span or trace found in context for agent output update"
                )
        except Exception as e:
            LOGGER.error(f"Failed during after_agent_callback(): {e}", exc_info=True)

    def before_model_callback(
        self,
        callback_context: callback_context.CallbackContext,
        llm_request: models.LlmRequest,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        try:
            input = adk_helpers.convert_adk_base_model_to_dict(llm_request)

            provider, model = litellm_wrappers.parse_provider_and_model(
                llm_request.model
            )
            if provider is None:
                provider = adk_helpers.get_adk_provider()

            # ADK runs `before_model_callback` before running `start_as_current_span` function for the LLM call,
            # which makes it impossible to update the Opik span from this method.
            # So we create a span manually here. This flow is handled inside ADKTracerWrapper.
            result = span_creation_handler.create_span_respecting_context(
                start_span_arguments=arguments_helpers.StartSpanParameters(
                    name=model,
                    project_name=self.project_name,
                    metadata={
                        **self.metadata,
                        llm_span_helpers.SPAN_STATUS: llm_span_helpers.LLMSpanStatus.STARTED,
                    },
                    type="llm",
                    model=model,
                    provider=provider,
                    input=input,
                ),
                distributed_trace_headers=None,
            )

            context_storage.add_span_data(result.span_data)

            # Track request start time for time-to-first-token calculation
            request_start_time = time.time()
            self._ttft_tracking[result.span_data.id] = (request_start_time, None)
        except Exception as e:
            LOGGER.error(f"Failed during before_model_callback(): {e}", exc_info=True)

    def after_model_callback(
        self,
        callback_context: callback_context.CallbackContext,
        llm_response: models.LlmResponse,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        try:
            is_partial = llm_response.partial is True
        except Exception:
            LOGGER.debug("Error checking for partial chunks", exc_info=True)
            is_partial = False

        span_id: Optional[str] = None
        exception_occurred = False
        try:
            model = None
            usage = None
            output = None

            if adk_helpers.has_empty_text_part_content(llm_response):
                # Clean up TTFT tracking if it exists before early return
                current_span = context_storage.top_span_data()
                if current_span is not None and current_span.id is not None:
                    self._safe_ttft_tracking(current_span.id, pop=True)
                return

            current_span = context_storage.top_span_data()
            if current_span is None:
                LOGGER.warning(
                    "No current span found in context for model output update"
                )
                return

            # Store span_id early for cleanup on all exit paths
            span_id = current_span.id

            # Track time-to-first-token: detect first token arrival
            # We check for first token on EVERY callback (including partial chunks)
            # to catch the first moment content appears
            request_start_time, first_token_time = self._safe_ttft_tracking(
                span_id, pop=False
            )
            if (
                first_token_time is None
                and request_start_time is not None
                and span_id is not None
            ):
                # Check if this response contains actual content (first token)
                # Content can be text or function calls (tool calls)
                if self._has_response_content(llm_response):
                    # First token detected - record the time
                    first_token_time = time.time()
                    self._ttft_tracking[span_id] = (
                        request_start_time,
                        first_token_time,
                    )

            # Ignore partial chunks for final processing, ADK will call this method with the full response at the end
            # Note: We intentionally keep the TTFT tracking entry for partial chunks since ADK will call
            # this method again with the final non-partial response, where we'll properly clean it up
            if is_partial:
                return

            try:
                output = adk_helpers.convert_adk_base_model_to_dict(llm_response)
                usage_data = llm_response_wrapper.pop_llm_usage_data(
                    output, current_span.provider
                )
                if usage_data is not None:
                    model = usage_data.model
                    usage = usage_data.opik_usage
            except Exception as e:
                LOGGER.debug(
                    f"Error converting LlmResponse to dict or extracting usage data, reason: {e}",
                    exc_info=True,
                )

            # Calculate time-to-first-token and add to metadata
            metadata_update = {}
            request_start_time, first_token_time = self._safe_ttft_tracking(
                span_id, pop=True
            )
            if first_token_time is not None and request_start_time is not None:
                time_to_first_token = first_token_time - request_start_time
                metadata_update["time_to_first_token"] = time_to_first_token

            # Merge with existing metadata
            if current_span.metadata is None:
                current_span.metadata = {}
            current_span.metadata.update(metadata_update)
            current_span.metadata[llm_span_helpers.SPAN_STATUS] = (
                llm_span_helpers.LLMSpanStatus.READY_FOR_FINALIZATION.value
            )

            current_span.update(
                output=output,
                name=model or current_span.model,
                type="llm",
                model=model,
                usage=usage,
                metadata=current_span.metadata,
                project_name=self.project_name,
            )

            context_storage.pop_span_data(ensure_id=current_span.id)
            current_span.init_end_time()
            # We close this span manually because otherwise ADK will close it too late,
            # and it will also add tool spans inside of it, which we want to avoid.
            if opik.is_tracing_active():
                self._opik_client.__internal_api__span__(**current_span.as_parameters)
            self._last_model_output = output

        except Exception as e:
            exception_occurred = True
            LOGGER.error(f"Failed during after_model_callback(): {e}", exc_info=True)
        finally:
            # Clean up TTFT tracking entry on all exit paths to prevent memory leak
            # Skip cleanup for partial chunks (normal return) since ADK will call again with final response
            # For final responses, entry is already popped at line 325, so this is a no-op
            # For errors (exception_occurred=True) or early returns, this ensures cleanup happens
            if span_id is not None and (exception_occurred or not is_partial):
                self._ttft_tracking.pop(span_id, None)

    def before_tool_callback(
        self,
        tool: base_tool.BaseTool,
        args: Dict[str, Any],
        tool_context: tool_context.ToolContext,
        *other_args: Any,
        **kwargs: Any,
    ) -> None:
        try:
            current_span = context_storage.top_span_data()

            tool_metadata = {
                "function_call_id": tool_context.function_call_id,
                **self.metadata,
            }

            # Update existing span with tool information
            if current_span is not None:
                current_span.update(
                    name=tool.name,
                    type="tool",
                    input=args,
                    metadata={**tool_metadata},
                    project_name=self.project_name,
                )
            else:
                LOGGER.warning(
                    f"No current span found in context for tool: {tool.name}"
                )
                _log_tool_context_warning(context=tool_context)

        except Exception as e:
            LOGGER.error(f"Failed during before_tool_callback(): {e}", exc_info=True)

    def after_tool_callback(
        self,
        tool: base_tool.BaseTool,
        args: Dict[str, Any],
        tool_context: tool_context.ToolContext,
        tool_response: Any,
        *other_args: Any,
        **kwargs: Any,
    ) -> None:
        try:
            # Debug logging for callback invocation
            current_span = context_storage.top_span_data()

            output = (
                tool_response
                if isinstance(tool_response, dict)
                else {"output": tool_response}
            )

            # Update existing span with tool output
            if current_span is not None:
                current_span.update(
                    output=output,
                    project_name=self.project_name,
                )
            else:
                LOGGER.warning(
                    f"No current span found in context for tool output update: {tool.name}"
                )
                _log_tool_context_warning(context=tool_context)
        except Exception as e:
            LOGGER.error(f"Failed during after_tool_callback(): {e}", exc_info=True)

    def __getstate__(self) -> Dict[str, Any]:
        state = self.__dict__.copy()
        state.pop("_opik_client", None)
        # Don't serialize TTFT tracking as it's runtime state
        state.pop("_ttft_tracking", None)
        return state

    def __setstate__(self, state: Dict[str, Any]) -> None:
        self.__dict__.update(state)
        self._init_internal_attributes()


def _try_add_agent_graph_to_metadata(
    metadata: Dict[str, Any], callback_context: callback_context.CallbackContext
) -> None:
    current_agent: Optional[google.adk.agents.BaseAgent] = (
        callback_context_info_extractors.try_get_current_agent_instance(
            callback_context
        )
    )

    if current_agent is None:
        return

    try:
        metadata["_opik_graph_definition"] = {
            "format": "mermaid",
            "data": mermaid_graph_builder.build_mermaid_graph_definition(
                current_agent.root_agent
            ),
        }
    except Exception:
        LOGGER.error("Failed to build mermaid graph for agent.", exc_info=True)


def _log_tool_context_warning(context: tool_context.ToolContext) -> None:
    if context is not None:
        warning = f"Function call id: {context.function_call_id}, agent name: {context.agent_name}"
        if context.actions is not None:
            warning += f", is escalate: {context.actions.escalate}, transfer to: {context.actions.transfer_to_agent}"

        LOGGER.warning(warning)
