"""Observability integrations for ieum.

Provides integration with observability platforms:
- Langfuse: LLM observability with tracing, metrics, and analytics

Environment Variables:
    LANGFUSE_PUBLIC_KEY: Your Langfuse public key
    LANGFUSE_SECRET_KEY: Your Langfuse secret key
    LANGFUSE_BASE_URL: Optional custom host (default: https://cloud.langfuse.com)

Usage:
    The Langfuse callback handler is automatically added when the
    environment variables are configured.

    # In .env file:
    LANGFUSE_PUBLIC_KEY=pk-lf-...
    LANGFUSE_SECRET_KEY=sk-lf-...
    # Optional: for self-hosted Langfuse
    LANGFUSE_BASE_URL=http://localhost:3000
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from ieum.config import settings

if TYPE_CHECKING:
    from langchain_core.callbacks import BaseCallbackHandler


def get_langfuse_handler(
    *,
    session_id: str | None = None,
    user_id: str | None = None,
    tags: list[str] | None = None,
    metadata: dict | None = None,
) -> "BaseCallbackHandler | None":
    """Get a Langfuse callback handler if configured.

    Creates and returns a Langfuse callback handler for LangChain integration.
    Returns None if Langfuse is not configured.

    Args:
        session_id: Optional session ID for grouping traces
        user_id: Optional user ID for analytics
        tags: Optional tags for filtering traces
        metadata: Optional metadata to attach to traces

    Returns:
        CallbackHandler if Langfuse is configured, None otherwise

    Example:
        ```python
        from ieum.observability import get_langfuse_handler

        handler = get_langfuse_handler(
            session_id="thread-123",
            tags=["production", "chat"],
        )

        if handler:
            agent.invoke({"messages": messages}, config={"callbacks": [handler]})
        ```
    """
    if not settings.has_langfuse:
        return None

    try:
        import os

        # Try new import path first (langfuse >= 2.0)
        try:
            from langfuse.langchain import CallbackHandler
        except ImportError:
            # Fallback to old import path
            from langfuse.callback import CallbackHandler

        # Ensure LANGFUSE_BASE_URL is set from our detected host
        # This handles LANGFUSE_BASEURL (no underscore) and LANGFUSE_HOST variants
        if settings.langfuse_host:
            # Set both for compatibility
            if not os.environ.get("LANGFUSE_BASE_URL"):
                os.environ["LANGFUSE_BASE_URL"] = settings.langfuse_host
            if not os.environ.get("LANGFUSE_HOST"):
                os.environ["LANGFUSE_HOST"] = settings.langfuse_host

        # Langfuse CallbackHandler reads from environment variables:
        # LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_HOST
        handler = CallbackHandler()

        # Set session-specific attributes if the handler supports them
        if session_id and hasattr(handler, "session_id"):
            handler.session_id = session_id
        if user_id and hasattr(handler, "user_id"):
            handler.user_id = user_id
        if tags and hasattr(handler, "tags"):
            handler.tags = tags
        if metadata and hasattr(handler, "metadata"):
            handler.metadata = metadata

        return handler
    except ImportError:
        # langfuse package not installed
        return None
    except Exception as e:
        # Failed to create handler - log the error for debugging
        import sys
        print(f"[Langfuse] Failed to create handler: {e}", file=sys.stderr)
        return None


def create_trace_metadata(
    *,
    assistant_id: str,
    model_name: str | None = None,
    model_provider: str | None = None,
) -> dict:
    """Create standard metadata for Langfuse traces.

    Args:
        assistant_id: Agent identifier
        model_name: Name of the model being used
        model_provider: Provider (ollama, lmstudio, vllm)

    Returns:
        Dictionary of metadata suitable for Langfuse traces
    """
    metadata = {
        "assistant_id": assistant_id,
        "cli_version": _get_version(),
    }

    if model_name:
        metadata["model_name"] = model_name
    if model_provider:
        metadata["model_provider"] = model_provider

    return metadata


def _get_version() -> str:
    """Get the CLI version."""
    try:
        from ieum._version import __version__

        return __version__
    except ImportError:
        return "unknown"


class LangfuseFlushOnExit:
    """Context manager to ensure Langfuse flushes on exit.

    Usage:
        ```python
        with LangfuseFlushOnExit(handler):
            # Run agent
            agent.invoke(...)
        # Handler is flushed automatically
        ```
    """

    def __init__(self, handler: "BaseCallbackHandler | None") -> None:
        """Initialize with a Langfuse handler."""
        self._handler = handler

    def __enter__(self) -> "LangfuseFlushOnExit":
        """Enter context."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Flush the handler on exit."""
        if self._handler is not None:
            try:
                # New langfuse >= 2.0: use client.flush()
                if hasattr(self._handler, "client") and self._handler.client:
                    self._handler.client.flush()
                # Old langfuse: handler.flush() or langfuse.flush()
                elif hasattr(self._handler, "flush"):
                    self._handler.flush()
                elif hasattr(self._handler, "langfuse") and self._handler.langfuse:
                    self._handler.langfuse.flush()
            except Exception:
                pass  # Don't fail on flush errors


def flush_langfuse_handler(handler: "BaseCallbackHandler | None") -> None:
    """Flush a Langfuse handler to send all pending traces.

    Args:
        handler: The Langfuse handler to flush (or None)
    """
    if handler is None:
        return

    try:
        # New langfuse >= 2.0: use client.flush()
        if hasattr(handler, "client") and handler.client:
            handler.client.flush()
        # Old langfuse: handler.flush() or langfuse.flush()
        elif hasattr(handler, "flush"):
            handler.flush()
        elif hasattr(handler, "langfuse") and handler.langfuse:
            handler.langfuse.flush()
    except Exception:
        pass  # Don't fail on flush errors


def show_langfuse_status() -> None:
    """Print Langfuse configuration status to console."""
    from ieum.config import console

    if settings.has_langfuse:
        host = settings.langfuse_host or "https://cloud.langfuse.com"
        console.print(f"[green]Langfuse:[/green] Enabled ({host})")
    else:
        console.print("[dim]Langfuse: Not configured[/dim]")
        console.print("[dim]  Set LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY to enable[/dim]")
