"""Service for managing Traces via the Galtea API.

This module provides both:
1. Direct API operations (create, list, get, delete)
2. Tracing utilities (trace decorator, start_trace, set_context, clear_context)

The tracing utilities delegate to galtea.utils.tracing which is the single source of truth.
"""

import logging
from contextlib import contextmanager
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union

from galtea.domain.models.trace import Trace, TraceBase, TraceType
from galtea.infrastructure.clients.http_client import Client
from galtea.utils.string import build_query_params, is_valid_id
from galtea.utils.timestamp import normalize_timestamp

F = TypeVar("F", bound=Callable[..., Any])


class TraceService:
    """Service for managing Traces via the Galtea API.

    A Trace represents a single tool or function call that occurred during
    an inference result, capturing input, output, timing, and error information.

    Note: For automatic tracing, use the @traced decorator or start_trace()
    context manager from galtea.utils.tracing. This service is for direct
    API operations only.
    """

    def __init__(self, client: Client):
        """Initialize the TraceService with the provided HTTP client.

        Args:
            client: The HTTP client for making API requests.
        """
        self._client: Client = client
        self._logger: logging.Logger = logging.getLogger(__name__)

    def create(
        self,
        inference_result_id: str,
        name: str,
        type: Optional[TraceType] = None,
        description: Optional[str] = None,
        input_data: Optional[Any] = None,
        output_data: Optional[Any] = None,
        error: Optional[str] = None,
        latency_ms: Optional[float] = None,
        metadata: Optional[Any] = None,
        start_time: Optional[str] = None,
        end_time: Optional[str] = None,
    ) -> Trace:
        """Create a new trace for a tool/function call.

        Args:
            inference_result_id: The ID of the inference result this trace belongs to.
            name: The name of the tool or function that was called.
            type: TraceType value (e.g., TOOL, GENERATION, AGENT, CHAIN, RETRIEVER).
            description: Human-readable description of what this operation does. Max 10KB.
            input_data: The input data passed to the tool.
            output_data: The output data returned by the tool.
            error: Error message if the tool call failed.
            latency_ms: The time in milliseconds the tool call took.
            metadata: Additional metadata about the trace.
            start_time: ISO 8601 timestamp when the trace started.
            end_time: ISO 8601 timestamp when the trace ended.

        Returns:
            The created trace object.

        Raises:
            ValueError: If inference_result_id is invalid or name is empty.
        """
        if not is_valid_id(inference_result_id):
            raise ValueError("The inference_result_id provided is not valid.")

        if not name or not isinstance(name, str) or not name.strip():
            raise ValueError("A valid name must be provided.")

        trace_base: TraceBase = TraceBase(
            inference_result_id=inference_result_id,
            name=name,
            type=type,
            description=description,
            input_data=input_data,
            output_data=output_data,
            error=error,
            latency_ms=latency_ms,
            metadata=metadata,
            start_time=start_time,
            end_time=end_time,
        )

        trace_base.model_validate(trace_base.model_dump())

        response = self._client.post(
            "traces",
            json=trace_base.model_dump(by_alias=True, exclude_none=True),
        )

        return Trace(**response.json())

    def create_batch(self, traces: List[TraceBase]) -> List[Trace]:
        """Create multiple traces in a single API call.

        Args:
            traces: List of trace objects to create.

        Returns:
            List of created trace objects.

        Raises:
            ValueError: If traces list is empty.
        """
        if not traces or not isinstance(traces, list):
            raise ValueError("A non-empty list of traces must be provided.")

        traces_payload: List[Dict[str, Any]] = [trace.model_dump(by_alias=True, exclude_none=True) for trace in traces]

        response = self._client.post(
            "traces/batch",
            json={"traces": traces_payload},
        )

        return [Trace(**trace_data) for trace_data in response.json()]

    def get(self, trace_id: str) -> Trace:
        """Retrieve a trace by its ID.

        Args:
            trace_id: The ID of the trace to retrieve.

        Returns:
            The retrieved trace object.

        Raises:
            ValueError: If trace_id is invalid.
        """
        if not is_valid_id(trace_id):
            raise ValueError("A valid trace ID must be provided.")

        response = self._client.get(f"traces/{trace_id}")
        return Trace(**response.json())

    def list(
        self,
        inference_result_id: Optional[Union[str, List[str]]] = None,
        session_id: Optional[Union[str, List[str]]] = None,
        types: Optional[Union[List[str], List[TraceType]]] = None,
        names: Optional[List[str]] = None,
        from_start_time: Optional[Union[str, int]] = None,
        to_start_time: Optional[Union[str, int]] = None,
        sort_by_start_time: Optional[str] = None,
        offset: Optional[int] = None,
        limit: Optional[int] = None,
    ) -> List[Trace]:
        """List traces for given inference result(s).

        Args:
            inference_result_id: The inference result ID or list of IDs.
            session_id: The session ID or list of IDs to filter traces.
            types: List of TraceType values to filter by.
            names: List of trace names to filter by.
            from_start_time (str | int, optional): Filter traces started at or after this timestamp.
                Accepts ISO 8601 string (e.g., '2024-01-01T00:00:00Z') or Unix timestamp (seconds).
            to_start_time (str | int, optional): Filter traces started at or before this timestamp.
                Accepts ISO 8601 string (e.g., '2024-12-31T23:59:59Z') or Unix timestamp (seconds).
            sort_by_start_time: Sort order ('asc' or 'desc').
            offset: Offset for pagination.
            limit: Limit for pagination.

        Returns:
            List of trace objects.

        Raises:
            ValueError: If neither inference_result_id nor session_id is provided.
        """
        if inference_result_id is not None and not isinstance(inference_result_id, (str, list)):
            raise ValueError("The inference result ID must be a string or a list of strings.")

        if session_id is not None and not isinstance(session_id, (str, list)):
            raise ValueError("The session ID must be a string or a list of strings.")

        if (session_id is None and inference_result_id is None) or (session_id == [] and inference_result_id == []):
            raise ValueError("At least one of session_id or inference_result_id must be provided.")

        inference_result_ids: Optional[List[str]] = (
            [inference_result_id] if isinstance(inference_result_id, str) else inference_result_id
        )
        session_ids: Optional[List[str]] = [session_id] if isinstance(session_id, str) else session_id

        if sort_by_start_time is not None and sort_by_start_time not in ["asc", "desc"]:
            raise ValueError("Sort by start_time must be 'asc' or 'desc'.")

        # Normalize types filter (convert enums to string values)
        normalized_types: Optional[List[str]] = None
        if types is not None:
            normalized_types = [t.value if isinstance(t, TraceType) else t for t in types]

        # Normalize date filter parameters
        from_start_time_normalized = normalize_timestamp(from_start_time)
        to_start_time_normalized = normalize_timestamp(to_start_time)

        query_params: str = build_query_params(
            inferenceResultIds=inference_result_ids,
            sessionIds=session_ids,
            types=normalized_types,
            names=names,
            fromStartTime=from_start_time_normalized,
            toStartTime=to_start_time_normalized,
            offset=offset,
            limit=limit,
            sort=["startTime", sort_by_start_time] if sort_by_start_time else None,
        )

        response = self._client.get(f"traces?{query_params}")
        return [Trace(**trace_data) for trace_data in response.json()]

    def delete(self, trace_id: str) -> None:
        """Delete a trace by its ID.

        Args:
            trace_id: The ID of the trace to delete.

        Raises:
            ValueError: If trace_id is invalid.
        """
        if not is_valid_id(trace_id):
            raise ValueError("A valid trace ID must be provided.")

        self._client.delete(f"traces/{trace_id}")

    # =========================================================================
    # Tracing Utilities (delegates to galtea.utils.tracing)
    # =========================================================================
    # These methods provide access to tracing utilities via galtea.traces.*
    # The single source of truth is galtea.utils.tracing

    def trace(
        self,
        name: Optional[str] = None,
        type: Optional[str] = None,
        log_args: bool = True,
        log_results: bool = True,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> Callable[[F], F]:
        """Decorator to automatically trace a function.

        All traced functions automatically create OpenTelemetry spans that are
        exported to Galtea.

        Args:
            name: Custom span name. Defaults to function name.
            type: TraceType value (e.g., TOOL, GENERATION, AGENT, CHAIN, RETRIEVER).
            log_args: Whether to log function arguments as input.
            log_results: Whether to log return value as output.
            attributes: Custom attributes to add to the span.

        Returns:
            Decorated function with automatic tracing.

        Example:
            @galtea.traces.trace(type="TOOL")
            def fetch_user(user_id: str) -> dict:
                return db.get_user(user_id)
        """
        from galtea.utils.tracing import trace as _trace

        return _trace(
            name=name,
            type=type,
            log_args=log_args,
            log_results=log_results,
            attributes=attributes,
        )

    @contextmanager
    def start_trace(
        self,
        name: str,
        type: str = TraceType.SPAN,
        input: Optional[Any] = None,
        metadata: Optional[Any] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ):
        """Context manager for creating a trace.

        Use this for fine-grained tracing of specific code blocks.

        Args:
            name: Name of the trace.
            type: TraceType value (e.g., TOOL, GENERATION, AGENT, CHAIN, RETRIEVER).
            input: Input data for the trace (Galtea semantic field).
            metadata: Metadata for the trace (Galtea semantic field).
            attributes: Custom OpenTelemetry attributes to add to the span.

        Yields:
            A GalteaSpan object with update() and other helper methods.

        Example:
            with galtea.traces.start_trace("database_query", type="TOOL", input={"query": "SELECT * FROM users"}) as s:
                db_query = "SELECT * FROM users"
                result = db.query(db_query)
                s.update(output=result, metadata={"query": db_query})
        """
        from galtea.utils.tracing import start_trace as _start_trace

        with _start_trace(
            name=name,
            type=type,
            input=input,
            metadata=metadata,
            attributes=attributes,
        ) as span:
            yield span

    def set_context(
        self,
        inference_result_id: Optional[str] = None,
        session_id: Optional[str] = None,
    ) -> object:
        """Set Galtea context values for trace correlation.

        Call this before traced functions to associate spans with an inference result.

        Args:
            inference_result_id: The inference result ID to associate traces with.
            session_id: The session ID.

        Returns:
            A context token that can be used with clear_context().

        Example:
            token = galtea.traces.set_context(inference_result_id="inf_123")
            try:
                my_traced_function()
            finally:
                galtea.traces.clear_context(token)
        """
        from galtea.utils.tracing import set_context as _set_context

        return _set_context(
            inference_result_id=inference_result_id,
        )

    def clear_context(self, token: object, flush: bool = True) -> None:
        """Clear a previously attached context.

        Args:
            token: The token returned by set_context().
            flush: If True (default), flush pending traces for the current inference result.
        """
        from galtea.utils.tracing import clear_context as _clear_context

        _clear_context(token=token, flush=flush)
