import os
from typing import Any, Literal

from kedro.io import AbstractDataset, DatasetError

REQUIRED_LANGFUSE_CREDENTIALS = {"public_key", "secret_key"}
REQUIRED_LANGFUSE_CREDENTIALS_AUTOGEN = {"endpoint"}
OPTIONAL_LANGFUSE_CREDENTIALS = {"host"}


class LangfuseTraceDataset(AbstractDataset):
    """Kedro dataset for managing Langfuse tracing clients and callbacks.

    This dataset provides appropriate tracing objects based on mode configuration,
    enabling seamless integration with different AI frameworks and direct SDK usage.
    Environment variables are automatically configured during initialization.

    **Modes:**

    - **langchain:** Returns a `CallbackHandler` for LangChain integration.
    - **openai:** Returns a wrapped OpenAI client with automatic tracing.
    - **autogen:** Returns a configured `Tracer` for AutoGen integration via OTLP.
      Note: Langfuse's graph visualisation is in beta and may not render
      complex multi-agent workflows correctly.
    - **sdk:** Returns a raw Langfuse client for manual tracing.

    Examples:
        Using catalog YAML configuration:

        ```yaml
        langfuse_trace:
          type: kedro_datasets_experimental.langfuse.LangfuseTraceDataset
          credentials: langfuse_credentials
          mode: openai
        ```

        Using Python API:

        ```python
        from kedro_datasets_experimental.langfuse import LangfuseTraceDataset

        # Basic usage (using default Langfuse cloud)
        dataset = LangfuseTraceDataset(
            credentials={
                "public_key": "pk_...",
                "secret_key": "sk_...",  # pragma: allowlist secret
                "openai": {"api_key": "sk-..."},  # pragma: allowlist secret
            },
            mode="openai",
        )

        # With custom host
        dataset = LangfuseTraceDataset(
            credentials={
                "public_key": "pk_...",
                "secret_key": "sk_...",  # pragma: allowlist secret
                "host": "https://custom.langfuse.com",
                "openai": {"api_key": "sk-..."},  # pragma: allowlist secret
            },
            mode="openai",
        )

        # Load tracing client
        client = dataset.load()
        response = client.chat.completions.create(...)  # Automatically traced

        # AutoGen mode Langfuse cloud
        dataset = LangfuseTraceDataset(
            credentials={
                "public_key": "pk_...",
                "secret_key": "sk_...",  # pragma: allowlist secret
                "endpoint": "https://cloud.langfuse.com/api/public/otel/v1/traces",
            },
            mode="autogen",
        )
        tracer = dataset.load()

        # AutoGen mode self-hosted
        dataset = LangfuseTraceDataset(
            credentials={
                "public_key": "pk_...",
                "secret_key": "sk_...",  # pragma: allowlist secret
                "host": "http://localhost:3000",
                "endpoint": "http://localhost:3000/api/public/otel/v1/traces",
            },
            mode="autogen",
        )
        tracer = dataset.load()
        # Use with AutoGen's runtime logging
        ```
    """

    def __init__(
        self,
        credentials: dict[str, Any],
        mode: Literal["langchain", "openai", "autogen", "sdk"] = "sdk",
        **trace_kwargs: Any
    ):
        """Initialize LangfuseTraceDataset and configure environment variables.

        Validates credentials and sets up appropriate environment variables for
        Langfuse tracing integration. Environment variables are set immediately
        during initialization for use by all tracing modes.

        Args:
            credentials: Dictionary with Langfuse credentials. Required: {public_key, secret_key}.
                Optional: {host} (defaults to Langfuse cloud if not provided).
                For autogen mode, {endpoint} is required — the full OTLP endpoint URL
                (e.g. https://cloud.langfuse.com/api/public/otel/v1/traces).
                For OpenAI mode, include openai section with {api_key, base_url}.
            mode: Tracing mode - "langchain", "openai", "autogen", or "sdk" (default).
            **trace_kwargs: Additional kwargs passed to the tracing client.

        Raises:
            DatasetError: If required Langfuse credentials are missing or empty.

        Examples:
            >>> # Basic SDK mode (using default Langfuse cloud)
            >>> dataset = LangfuseTraceDataset(
            ...     credentials={"public_key": "pk_...", "secret_key": "sk_..."}  # pragma: allowlist secret
            ... )

            >>> # With custom host
            >>> dataset = LangfuseTraceDataset(
            ...     credentials={
            ...         "public_key": "pk_...",
            ...         "secret_key": "sk_...",  # pragma: allowlist secret
            ...         "host": "https://custom.langfuse.com"
            ...     }
            ... )

            >>> # OpenAI mode with API key
            >>> dataset = LangfuseTraceDataset(
            ...     credentials={
            ...         "public_key": "pk_...",
            ...         "secret_key": "sk_...",  # pragma: allowlist secret
            ...         "openai": {"api_key": "sk-...", "base_url": "..."}  # pragma: allowlist secret
            ...     },
            ...     mode="openai"
            ... )

            >>> # AutoGen mode cloud
            >>> dataset = LangfuseTraceDataset(
            ...     credentials={
            ...         "public_key": "pk_...",
            ...         "secret_key": "sk_...",  # pragma: allowlist secret
            ...         "endpoint": "https://cloud.langfuse.com/api/public/otel/v1/traces",
            ...     },
            ...     mode="autogen"
            ... )

            >>> # AutoGen mode self-hosted
            >>> dataset = LangfuseTraceDataset(
            ...     credentials={
            ...         "public_key": "pk_...",
            ...         "secret_key": "sk_...",  # pragma: allowlist secret
            ...         "host": "http://localhost:3000",
            ...         "endpoint": "http://localhost:3000/api/public/otel/v1/traces",
            ...     },
            ...     mode="autogen"
            ... )

        Note:
            Sets LANGFUSE_SECRET_KEY, LANGFUSE_PUBLIC_KEY, and LANGFUSE_HOST
            environment variables from the provided credentials.
        """
        self._credentials = credentials
        self._mode = mode
        self._trace_kwargs = trace_kwargs
        self._cached_client = None

        # Validate Langfuse credentials before setting environment variables
        self._validate_langfuse_credentials()

        # Set Langfuse environment variables from credentials
        os.environ["LANGFUSE_SECRET_KEY"] = self._credentials["secret_key"]
        os.environ["LANGFUSE_PUBLIC_KEY"] = self._credentials["public_key"]

        if "host" in self._credentials:
            os.environ["LANGFUSE_HOST"] = self._credentials["host"]

    def _validate_langfuse_credentials(self) -> None:
        """Validate Langfuse credentials before setting environment variables.

        Raises:
            DatasetError: If Langfuse credentials are missing or invalid.
        """
        # Validate required keys
        for key in REQUIRED_LANGFUSE_CREDENTIALS:
            if key not in self._credentials:
                raise DatasetError(f"Missing required Langfuse credential: '{key}'")

            # Validate that credential is not empty
            if not self._credentials[key] or not str(self._credentials[key]).strip():
                raise DatasetError(f"Langfuse credential '{key}' cannot be empty")

        # Validate optional keys if present
        for key in OPTIONAL_LANGFUSE_CREDENTIALS:
            if key in self._credentials:
                # If host is provided, it cannot be empty
                if not self._credentials[key] or not str(self._credentials[key]).strip():
                    raise DatasetError(f"Langfuse credential '{key}' cannot be empty if provided")

        # AutoGen mode has additional required credentials
        if self._mode == "autogen":
            for key in REQUIRED_LANGFUSE_CREDENTIALS_AUTOGEN:
                if not self._credentials.get(key):
                    raise DatasetError(
                        f"AutoGen mode requires '{key}' in credentials "
                        f"(e.g. 'https://cloud.langfuse.com/api/public/otel/v1/traces'). "
                        f"Provide the full OTLP endpoint URL for trace export."
                    )

    def _describe(self) -> dict[str, Any]:
        """Return a description of the dataset for Kedro's internal use.

        Returns:
            Dictionary containing dataset description with mode and masked credentials.
        """
        return {"mode": self._mode, "credentials": "***"}

    def _validate_openai_client_params(self) -> None:
        """Validate OpenAI credentials in the 'openai' section.

        Raises:
            DatasetError: If OpenAI credentials are missing or invalid.
        """
        # Check if openai section exists
        if "openai" not in self._credentials:
            raise DatasetError("OpenAI mode requires 'openai' section in credentials")

        openai_creds = self._credentials["openai"]

        # Check for required API key
        if "api_key" not in openai_creds:
            raise DatasetError("Missing required OpenAI credential: 'api_key'")

        # Validate that API key is not empty
        if not openai_creds["api_key"] or not openai_creds["api_key"].strip():
            raise DatasetError("OpenAI API key cannot be empty")

        # Validate base_url is not empty if provided
        if "base_url" in openai_creds and not str(openai_creds["base_url"]).strip():
            raise DatasetError("OpenAI credential 'base_url' cannot be empty if provided")

    def _build_autogen_tracer(self) -> Any:
        """Build and return a configured Tracer for AutoGen integration with Langfuse.

        Sets up OpenTelemetry TracerProvider with OTLP exporter to Langfuse,
        configures it as the global provider, and returns a ready-to-use Tracer.

        Returns:
            Tracer configured to export traces to Langfuse.

        Raises:
            DatasetError: If required OpenTelemetry dependencies are not installed.
        """
        try:
            from opentelemetry import trace  # noqa: PLC0415
            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (  # noqa: PLC0415
                OTLPSpanExporter,
            )
            from opentelemetry.sdk.trace import TracerProvider  # noqa: PLC0415
            from opentelemetry.sdk.trace.export import (  # noqa: PLC0415
                BatchSpanProcessor,
            )
        except ImportError as exc:
            raise DatasetError(
                "AutoGen mode requires OpenTelemetry. "
                "Install with: pip install opentelemetry-sdk opentelemetry-exporter-otlp-proto-http"
            ) from exc

        import base64  # noqa: PLC0415

        auth = base64.b64encode(
            f"{self._credentials['public_key']}:{self._credentials['secret_key']}".encode()
        ).decode()

        # Endpoint is provided by user and validated in _validate_langfuse_credentials
        endpoint = self._credentials["endpoint"]

        exporter = OTLPSpanExporter(
            endpoint=endpoint,
            headers={"Authorization": f"Basic {auth}"}
        )

        processor = BatchSpanProcessor(exporter)

        # Use existing provider if already set, otherwise create a new one.
        existing_provider = trace.get_tracer_provider()
        if hasattr(existing_provider, "add_span_processor"):
            existing_provider.add_span_processor(processor)
        else:
            provider = TracerProvider()
            provider.add_span_processor(processor)
            trace.set_tracer_provider(provider)

        return trace.get_tracer("langfuse.autogen")

    def load(self) -> Any:
        """Load appropriate tracing client based on configured mode.

        Creates and returns the appropriate tracing client for the specified mode.
        The client is cached after first load to avoid repeated initialisation.
        All clients use environment variables set during initialisation for authentication.

        Returns:
            Tracing client object based on mode:
            - langchain mode: CallbackHandler for LangChain integration
            - openai mode: Wrapped OpenAI client with automatic tracing
            - autogen mode: Configured Tracer for OpenTelemetry integration
            - sdk mode: Raw Langfuse client for manual tracing

        Raises:
            DatasetError: If mode-specific dependencies are missing or credentials are invalid.

        Examples:
            # LangChain mode
                dataset = LangfuseTraceDataset(credentials=creds, mode="langchain")
                callback = dataset.load()
                chain.invoke(input, config={"callbacks": [callback]})

            # OpenAI mode
                dataset = LangfuseTraceDataset(credentials=creds, mode="openai")
                client = dataset.load()
                response = client.chat.completions.create(model="gpt-4", messages=[...])

            # AutoGen mode
                dataset = LangfuseTraceDataset(credentials=creds, mode="autogen")
                tracer = dataset.load()  # Returns configured Tracer

                # Option 1: Automatic tracing (LLM calls traced automatically)
                agent.invoke(context)  # Traces sent to Langfuse

                # Option 2: Add custom spans with context
                with tracer.start_as_current_span("response_generation") as span:
                    span.set_attribute("intent", "claim_new")
                    agent.invoke(context)  # Child spans nested under parent

            # SDK mode
                dataset = LangfuseTraceDataset(credentials=creds, mode="sdk")
                langfuse = dataset.load()
                trace = langfuse.trace(name="my-trace")
        """
        # Return cached client if available
        if self._cached_client is not None:
            return self._cached_client

        # Create and cache the appropriate client
        if self._mode == "langchain":
            from langfuse.langchain import CallbackHandler  # noqa: PLC0415
            self._cached_client = CallbackHandler(**self._trace_kwargs)
        elif self._mode == "openai":
            from langfuse.openai import OpenAI  # noqa: PLC0415
            self._validate_openai_client_params()
            self._cached_client = OpenAI(**self._credentials["openai"])
        elif self._mode == "autogen":
            self._cached_client = self._build_autogen_tracer()
        else:
            try:
                from langfuse import get_client  # noqa: PLC0415
                self._cached_client = get_client()
            except ImportError:
                from langfuse import Langfuse  # noqa: PLC0415
                self._cached_client = Langfuse(**self._trace_kwargs)

        return self._cached_client

    def save(self, data: Any) -> None:
        """Save operation is not supported for tracing datasets.

        Args:
            data: Data to save (not used).

        Raises:
            NotImplementedError: Always raised as tracing datasets are read-only.

        Note:
            LangfuseTraceDataset is designed for providing tracing clients,
            not for data storage. Use the returned tracing clients to automatically
            log traces, spans, and generations to Langfuse.
        """
        raise NotImplementedError("LangfuseTraceDataset is read-only - it provides tracing clients, not data storage")
