"""
Opik tracing integration for Boring-OS SDK.

Provides automatic tracing of all tool calls using Opik's observability platform.
Traces can be sent to a local Opik server or Comet.com cloud.

Usage:
    from boring_os import BoringOS
    
    # Enable tracing
    client = BoringOS(
        api_key="...",
        enable_tracing=True,
        opik_project="my-agent"
    )
    
    # All tool calls are now traced
    result = client.tools.call("get_customer_profile", customer_id="123")

Requires: pip install boring-os[opik]
"""

from __future__ import annotations

import time
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Generator

if TYPE_CHECKING:
    from boring_os.models import ToolResult

# Lazy-loaded opik module
_opik_module = None


def _get_opik() -> Any:
    """
    Lazy import opik to avoid requiring it as a hard dependency.
    
    Returns:
        The opik module
    
    Raises:
        ImportError: If opik is not installed
    """
    global _opik_module
    if _opik_module is None:
        try:
            import opik
            _opik_module = opik
        except ImportError:
            raise ImportError(
                "opik is required for tracing. "
                "Install with: pip install boring-os[opik]"
            )
    return _opik_module


def is_opik_available() -> bool:
    """
    Check if opik is installed and available.
    
    Returns:
        True if opik can be imported, False otherwise
    """
    try:
        import opik  # noqa: F401
        return True
    except ImportError:
        return False


class OpikTracer:
    """
    Tracer that sends spans to Opik for observability.
    
    Wraps tool calls with Opik spans to capture:
    - Tool name and parameters (input)
    - Execution result (output)
    - Execution time and metadata
    - Error information on failures
    
    Example:
        tracer = OpikTracer(project_name="my-agent")
        
        with tracer.trace_tool("get_customer", {"id": "123"}) as span:
            result = execute_tool(...)
            span.set_output(result)
    """
    
    # Default local Opik server URL
    DEFAULT_LOCAL_URL = "http://localhost:5173"
    
    def __init__(
        self,
        project_name: str = "boring-os",
        use_local: bool = True,
        url: str | None = None,
        api_key: str | None = None,
        workspace: str | None = None,
    ):
        """
        Initialize the Opik tracer.
        
        Args:
            project_name: Project name for grouping traces in Opik
            use_local: If True, use local Opik server; if False, use Comet cloud
            url: Custom URL for Opik server (defaults to http://localhost:5173 for local)
            api_key: Comet API key (required if use_local=False)
            workspace: Comet workspace (required if use_local=False)
        
        Note:
            Configuration is set via session config to avoid Opik's
            connectivity validation during initialization. This allows
            the client to be created even if the Opik server is not
            currently running - traces will simply fail gracefully.
        """
        opik = _get_opik()
        
        # Import Opik's config module to set session configuration
        # This bypasses opik.configure() which validates server connectivity
        from opik.config import update_session_config
        
        # Configure Opik based on mode using session config
        if use_local:
            local_url = url or self.DEFAULT_LOCAL_URL
            # Ensure URL has trailing slash for consistency with Opik internals
            if not local_url.endswith("/"):
                local_url = local_url + "/"
            update_session_config("url_override", local_url)
        else:
            # For cloud mode, set API key and workspace
            if api_key:
                update_session_config("api_key", api_key)
            if workspace:
                update_session_config("workspace", workspace)
        
        # Set the project name in session config
        update_session_config("project_name", project_name)
        
        self._project_name = project_name
        self._enabled = True
        self._opik = opik
        self._url = url
    
    @property
    def enabled(self) -> bool:
        """Check if tracing is enabled."""
        return self._enabled
    
    def enable(self) -> None:
        """Enable tracing."""
        self._enabled = True
    
    def disable(self) -> None:
        """Disable tracing temporarily."""
        self._enabled = False
    
    @contextmanager
    def trace_tool(
        self,
        tool_name: str,
        params: dict[str, Any],
        trace_id: str | None = None,
        timeout: int | None = None,
    ) -> Generator["OpikSpanWrapper", None, None]:
        """
        Create a span for a tool call.
        
        Args:
            tool_name: Name of the tool being called
            params: Tool parameters
            trace_id: Optional trace ID for correlation
            timeout: Tool execution timeout
        
        Yields:
            OpikSpanWrapper for setting output and metadata
        """
        if not self._enabled:
            # Return a no-op wrapper if tracing is disabled
            yield OpikSpanWrapper(None)
            return
        
        # Build span input
        span_input = {
            "tool_name": tool_name,
            "params": params,
        }
        if timeout is not None:
            span_input["timeout"] = timeout
        
        # Build initial metadata
        metadata = {
            "sdk": "boring-os",
            "sdk_version": "0.1.0",
            "project": self._project_name,
        }
        if trace_id:
            metadata["trace_id"] = trace_id
        
        start_time = time.time()
        
        # Create Opik span
        with self._opik.start_span(
            name=f"boring_os.tool.{tool_name}",
            input=span_input,
            metadata=metadata,
        ) as span:
            wrapper = OpikSpanWrapper(span, start_time)
            try:
                yield wrapper
            except Exception as e:
                # Record error in span
                wrapper.set_error(e)
                raise
    
    def track(self, name: str | None = None) -> Any:
        """
        Get the opik.track decorator for manual tracing.
        
        Args:
            name: Optional name for the traced function
        
        Returns:
            The opik.track decorator
        
        Example:
            tracer = OpikTracer()
            
            @tracer.track(name="my_workflow")
            def process_request(query: str):
                return client.tools.call("some_tool", query=query)
        """
        if name:
            return self._opik.track(name=name)
        return self._opik.track


class OpikSpanWrapper:
    """
    Wrapper around an Opik span for setting output and metadata.
    
    Handles both active spans and no-op mode when tracing is disabled.
    """
    
    def __init__(self, span: Any, start_time: float | None = None):
        """
        Initialize the span wrapper.
        
        Args:
            span: The underlying Opik span (or None for no-op mode)
            start_time: Start timestamp for duration calculation
        """
        self._span = span
        self._start_time = start_time or time.time()
    
    def set_output(self, output: Any) -> None:
        """
        Set the span output (tool result).
        
        Args:
            output: The tool execution result
        """
        if self._span is not None:
            self._span.set_output(output)
    
    def set_result(self, result: "ToolResult") -> None:
        """
        Set span output from a ToolResult object.
        
        Extracts result data and metadata from the ToolResult.
        
        Args:
            result: The ToolResult from tool execution
        """
        if self._span is None:
            return
        
        # Set the actual result data
        if result.result is not None:
            self._span.set_output(result.result)
        
        # Add execution metadata
        execution_ms = int((time.time() - self._start_time) * 1000)
        metadata = {
            "execution_ms": execution_ms,
            "status": result.status.value,
            "request_id": result.request_id,
        }
        
        # Add tool metadata if available
        if result.metadata:
            if result.metadata.tool_name:
                metadata["tool_name"] = result.metadata.tool_name
            if result.metadata.tool_version:
                metadata["tool_version"] = result.metadata.tool_version
            if result.metadata.execution_ms:
                metadata["server_execution_ms"] = result.metadata.execution_ms
        
        self._span.set_metadata(metadata)
    
    def set_metadata(self, metadata: dict[str, Any]) -> None:
        """
        Add additional metadata to the span.
        
        Args:
            metadata: Key-value pairs to add
        """
        if self._span is not None:
            self._span.set_metadata(metadata)
    
    def set_error(self, error: Exception) -> None:
        """
        Record an error in the span.
        
        Args:
            error: The exception that occurred
        """
        if self._span is not None:
            execution_ms = int((time.time() - self._start_time) * 1000)
            self._span.set_metadata({
                "error": str(error),
                "error_type": type(error).__name__,
                "execution_ms": execution_ms,
                "status": "error",
            })


def configure_tracing(
    project_name: str = "boring-os",
    use_local: bool = True,
    url: str | None = None,
    api_key: str | None = None,
    workspace: str | None = None,
) -> OpikTracer:
    """
    Configure and return an Opik tracer.
    
    Convenience function for creating a tracer with common settings.
    
    Args:
        project_name: Project name for grouping traces
        use_local: If True, use local Opik server
        url: Custom URL for Opik server (defaults to http://localhost:5173 for local)
        api_key: Comet API key (for cloud mode)
        workspace: Comet workspace (for cloud mode)
    
    Returns:
        Configured OpikTracer instance
    
    Example:
        from boring_os.tracing import configure_tracing
        
        # For local development (uses default http://localhost:5173)
        tracer = configure_tracing(project_name="my-project", use_local=True)
        
        # For local with custom URL
        tracer = configure_tracing(
            project_name="my-project",
            use_local=True,
            url="http://localhost:8000"
        )
        
        # For production with Comet cloud
        tracer = configure_tracing(
            project_name="production",
            use_local=False,
            api_key="your-api-key",
            workspace="your-workspace"
        )
    """
    return OpikTracer(
        project_name=project_name,
        use_local=use_local,
        url=url,
        api_key=api_key,
        workspace=workspace,
    )
