"""
Langfuse Integration for Sentrial SDK

Enables dual-write mode where data is sent to both Sentrial and Langfuse simultaneously.
This allows customers to evaluate Sentrial while keeping their existing Langfuse setup.

Usage:
    from sentrial import SentrialClient
    from sentrial.langfuse_integration import LangfuseIntegration

    # Option 1: Pass Langfuse client directly
    from langfuse import Langfuse
    langfuse = Langfuse()
    
    client = SentrialClient(
        langfuse=LangfuseIntegration(langfuse_client=langfuse)
    )

    # Option 2: Pass credentials (client created automatically)
    client = SentrialClient(
        langfuse=LangfuseIntegration(
            public_key="pk-lf-...",
            secret_key="sk-lf-...",
        )
    )

    # Then use Sentrial normally - data flows to both systems
    with client.begin(user_id="user_123", event="chat", input="Hello") as interaction:
        interaction.track_tool_call(...)
        interaction.set_output("World")
"""

import os
import logging
from typing import Any, Callable, Optional, Dict, TYPE_CHECKING

if TYPE_CHECKING:
    from langfuse import Langfuse
    from langfuse.client import StatefulTraceClient, StatefulSpanClient, StatefulGenerationClient

_logger = logging.getLogger("sentrial.langfuse")


class LangfuseIntegration:
    """
    Langfuse integration for dual-write mode.
    
    When attached to a SentrialClient, all tracking calls are mirrored to Langfuse.
    This enables side-by-side comparison and gradual migration.
    """

    def __init__(
        self,
        langfuse_client: Optional["Langfuse"] = None,
        public_key: Optional[str] = None,
        secret_key: Optional[str] = None,
        host: Optional[str] = None,
        enabled: bool = True,
        fail_silently: bool = True,
    ):
        """
        Initialize Langfuse integration.

        Args:
            langfuse_client: Existing Langfuse client instance (preferred)
            public_key: Langfuse public key (if not using existing client)
            secret_key: Langfuse secret key (if not using existing client)
            host: Langfuse host URL (defaults to LANGFUSE_HOST env var or cloud)
            enabled: Whether to actually send data to Langfuse (default: True)
            fail_silently: If True, Langfuse errors are logged but don't crash (default: True)
        """
        self.enabled = enabled
        self.fail_silently = fail_silently
        self._client: Optional["Langfuse"] = None
        self._owns_client = False

        if not enabled:
            return

        if langfuse_client is not None:
            self._client = langfuse_client
            self._owns_client = False
        else:
            # Create client from credentials
            try:
                from langfuse import Langfuse
                
                self._client = Langfuse(
                    public_key=public_key or os.environ.get("LANGFUSE_PUBLIC_KEY"),
                    secret_key=secret_key or os.environ.get("LANGFUSE_SECRET_KEY"),
                    host=host or os.environ.get("LANGFUSE_HOST", "https://cloud.langfuse.com"),
                )
                self._owns_client = True
            except ImportError:
                _logger.warning(
                    "Langfuse package not installed. Install with: pip install langfuse"
                )
                self.enabled = False
            except Exception as e:
                if fail_silently:
                    _logger.warning(f"Failed to initialize Langfuse client: {e}")
                    self.enabled = False
                else:
                    raise

    @property
    def client(self) -> Optional["Langfuse"]:
        """Get the underlying Langfuse client."""
        return self._client

    def _safe_call(self, func: Callable, *args, **kwargs) -> Any:
        """Safely call a Langfuse function, handling errors based on fail_silently."""
        if not self.enabled or self._client is None:
            return None

        try:
            return func(*args, **kwargs)
        except Exception as e:
            if self.fail_silently:
                _logger.warning(f"Langfuse call failed: {e}")
                return None
            raise

    def create_trace(
        self,
        session_id: str,
        name: str,
        user_id: str,
        input_data: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Optional["StatefulTraceClient"]:
        """
        Create a Langfuse trace corresponding to a Sentrial session.

        Args:
            session_id: Sentrial session ID (used as trace ID for correlation)
            name: Trace name (typically the agent/event name)
            user_id: User ID
            input_data: Initial input
            metadata: Additional metadata

        Returns:
            Langfuse trace object, or None if disabled/failed
        """
        def _create():
            return self._client.trace(
                id=session_id,  # Use Sentrial session ID for correlation
                name=name,
                user_id=user_id,
                input=input_data,
                metadata={
                    "sentrial_session_id": session_id,
                    **(metadata or {}),
                },
            )

        return self._safe_call(_create)

    def create_span(
        self,
        trace: "StatefulTraceClient",
        name: str,
        input_data: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Optional["StatefulSpanClient"]:
        """
        Create a Langfuse span (for tool calls).

        Args:
            trace: Parent trace
            name: Span name (tool name)
            input_data: Tool input
            metadata: Additional metadata

        Returns:
            Langfuse span object, or None if disabled/failed
        """
        if trace is None:
            return None

        def _create():
            return trace.span(
                name=name,
                input=input_data,
                metadata=metadata,
            )

        return self._safe_call(_create)

    def end_span(
        self,
        span: "StatefulSpanClient",
        output_data: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
    ) -> None:
        """
        End a Langfuse span.

        Args:
            span: The span to end
            output_data: Tool output
            error: Error message if failed
        """
        if span is None:
            return

        def _end():
            span.end(
                output=output_data,
                level="ERROR" if error else "DEFAULT",
                status_message=error,
            )

        self._safe_call(_end)

    def create_generation(
        self,
        trace: "StatefulTraceClient",
        name: str,
        model: Optional[str] = None,
        input_data: Optional[Any] = None,
        output_data: Optional[Any] = None,
        usage: Optional[Dict[str, int]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Optional["StatefulGenerationClient"]:
        """
        Create a Langfuse generation (for LLM calls).

        Args:
            trace: Parent trace
            name: Generation name
            model: Model name
            input_data: LLM input (messages, prompt)
            output_data: LLM output
            usage: Token usage dict with prompt_tokens, completion_tokens, total_tokens
            metadata: Additional metadata

        Returns:
            Langfuse generation object, or None if disabled/failed
        """
        if trace is None:
            return None

        def _create():
            return trace.generation(
                name=name,
                model=model,
                input=input_data,
                output=output_data,
                usage=usage,
                metadata=metadata,
            )

        return self._safe_call(_create)

    def update_trace(
        self,
        trace: "StatefulTraceClient",
        output: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Update a trace with final output.

        Args:
            trace: The trace to update
            output: Final output/response
            metadata: Additional metadata to merge
        """
        if trace is None:
            return

        def _update():
            trace.update(
                output=output,
                metadata=metadata,
            )

        self._safe_call(_update)

    def score_trace(
        self,
        trace: "StatefulTraceClient",
        name: str,
        value: float,
        comment: Optional[str] = None,
    ) -> None:
        """
        Add a score to a trace.

        Args:
            trace: The trace to score
            name: Score name
            value: Score value
            comment: Optional comment/explanation
        """
        if trace is None:
            return

        def _score():
            trace.score(
                name=name,
                value=value,
                comment=comment,
            )

        self._safe_call(_score)

    def flush(self) -> None:
        """Flush pending Langfuse events."""
        if self._client is not None:
            self._safe_call(self._client.flush)

    def shutdown(self) -> None:
        """Shutdown the Langfuse client (only if we own it)."""
        if self._owns_client and self._client is not None:
            self._safe_call(self._client.shutdown)
            self._client = None


class LangfuseTraceContext:
    """
    Context for tracking Langfuse trace/spans alongside a Sentrial interaction.
    
    Used internally by the SentrialClient when langfuse integration is enabled.
    """

    def __init__(
        self,
        integration: LangfuseIntegration,
        session_id: str,
        name: str,
        user_id: str,
        input_data: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ):
        self.integration = integration
        self.session_id = session_id
        self._trace = integration.create_trace(
            session_id=session_id,
            name=name,
            user_id=user_id,
            input_data=input_data,
            metadata=metadata,
        )

    @property
    def trace(self) -> Optional["StatefulTraceClient"]:
        """Get the underlying Langfuse trace."""
        return self._trace

    def track_tool_call(
        self,
        tool_name: str,
        tool_input: Dict[str, Any],
        tool_output: Dict[str, Any],
        tool_error: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Track a tool call as a Langfuse span."""
        span = self.integration.create_span(
            trace=self._trace,
            name=tool_name,
            input_data=tool_input,
            metadata=metadata,
        )
        
        if span is not None:
            self.integration.end_span(
                span=span,
                output_data=tool_output,
                error=tool_error.get("message") if tool_error else None,
            )

    def track_llm_call(
        self,
        name: str,
        model: Optional[str] = None,
        input_data: Optional[Any] = None,
        output_data: Optional[Any] = None,
        prompt_tokens: Optional[int] = None,
        completion_tokens: Optional[int] = None,
        total_tokens: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Track an LLM call as a Langfuse generation."""
        usage = None
        if any([prompt_tokens, completion_tokens, total_tokens]):
            usage = {}
            if prompt_tokens is not None:
                usage["prompt_tokens"] = prompt_tokens
            if completion_tokens is not None:
                usage["completion_tokens"] = completion_tokens
            if total_tokens is not None:
                usage["total_tokens"] = total_tokens

        self.integration.create_generation(
            trace=self._trace,
            name=name,
            model=model,
            input_data=input_data,
            output_data=output_data,
            usage=usage,
            metadata=metadata,
        )

    def finish(
        self,
        output: Optional[str] = None,
        success: bool = True,
        custom_metrics: Optional[Dict[str, float]] = None,
    ) -> None:
        """Finish the trace with final output and scores."""
        # Update trace with output
        self.integration.update_trace(
            trace=self._trace,
            output=output,
            metadata={"success": success},
        )

        # Add scores for custom metrics
        if custom_metrics:
            for metric_name, value in custom_metrics.items():
                self.integration.score_trace(
                    trace=self._trace,
                    name=metric_name,
                    value=value,
                )

        # Add success score
        self.integration.score_trace(
            trace=self._trace,
            name="success",
            value=1.0 if success else 0.0,
        )

        # Flush to ensure data is sent
        self.integration.flush()
