# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Klira SDK v2 — LLM observability, tracing, and policy enforcement."""

from __future__ import annotations

import atexit
import logging
from contextlib import contextmanager
from typing import Any, Iterator

from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

from klira.sdk.config import KliraConfig
from klira.sdk.exceptions import KliraInitError
from klira.sdk.telemetry.exporter import KliraOTLPSpanExporter
from klira.sdk.telemetry.filter import FilteringSpanExporter
from klira.sdk.telemetry.processor import NoneAttributeFilterProcessor

logger = logging.getLogger(__name__)


class Klira:
    """Main SDK entry point.

    Usage:
        Klira.init(api_key="klira_...")

    The init() method sets up the telemetry pipeline:
        TracerProvider → NoneAttributeFilterProcessor → BatchSpanProcessor
        → FilteringSpanExporter → KliraOTLPSpanExporter → Platform
    """

    _config: KliraConfig | None = None
    _provider: TracerProvider | None = None
    _initialized: bool = False
    _atexit_registered: bool = False

    @classmethod
    def init(cls, **kwargs: Any) -> KliraConfig:
        """Initialize the Klira SDK.

        Args:
            api_key: Klira API key (must start with 'klira_').
            app_name: Application name for resource identification.
            endpoint: Klira platform endpoint.
            tracing_enabled: Whether to enable tracing (default True).
            anonymization: PHI/PII de-identification method ('redact', 'mask',
                'replace', 'hash'). None to disable (default).
            clinical_domain: Clinical domain for healthcare features.
            framework: Orchestration framework name (e.g., 'langchain', 'crewai').
            evals_run: Eval run identifier (enables eval mode).
            dataset_id: Dataset identifier for evals.
            use_remote_policies: Whether to fetch policies from platform.
            policies_path: Local policy file path.
            debug: Enable debug logging.
            **kwargs: Additional config parameters.

        Returns:
            The created KliraConfig instance.

        Raises:
            KliraInitError: If initialization fails.
        """
        try:
            config = KliraConfig.from_env(**kwargs)
        except Exception as e:
            raise KliraInitError(f"Failed to create config: {e}") from e

        if config.debug:
            logging.getLogger("klira").setLevel(logging.DEBUG)

        cls._config = config

        if config.tracing_enabled:
            cls._setup_telemetry(config)
            if not cls._initialized:
                cls._auto_patch_llm_clients()

        cls._load_policies(config)

        cls._initialized = True
        logger.info(
            "Klira SDK initialized (app=%s, tracing=%s)",
            config.app_name,
            config.tracing_enabled,
        )
        return config

    @classmethod
    def _setup_telemetry(cls, config: KliraConfig) -> None:
        """Set up the OpenTelemetry pipeline."""
        import os

        # Prevent CrewAI from overriding our TracerProvider with its own.
        # CrewAI's Telemetry.set_tracer() calls trace.set_tracer_provider()
        # which would redirect all Klira spans to CrewAI's exporter.
        os.environ.setdefault("CREWAI_DISABLE_TELEMETRY", "true")

        # If a third-party framework (e.g. CrewAI) already set a TracerProvider
        # at import time, OTel's Once guard prevents us from replacing it.
        # Reset the guard so we can install our own provider.
        # NOTE: This uses private OTel attributes — no public API exists for
        # this. We guard with hasattr + try/except to handle future OTel versions.
        current = trace.get_tracer_provider()
        if isinstance(current, TracerProvider):
            try:
                if hasattr(trace, "_TRACER_PROVIDER_SET_ONCE") and hasattr(
                    trace, "_TRACER_PROVIDER"
                ):
                    trace._TRACER_PROVIDER_SET_ONCE = trace.Once()
                    trace._TRACER_PROVIDER = None
                    logger.debug(
                        "Reset existing TracerProvider to install Klira pipeline"
                    )
                else:
                    logger.warning(
                        "Cannot reset OTel TracerProvider: expected private "
                        "attributes not found (opentelemetry-api version may "
                        "have changed). Klira telemetry may not work correctly."
                    )
            except Exception:
                logger.warning(
                    "Failed to reset OTel TracerProvider. "
                    "Klira telemetry may not work correctly.",
                    exc_info=True,
                )

        resource = Resource.create(
            {
                "service.name": config.app_name,
                "klira.sdk.version": "2.0.0",
            }
        )

        provider = TracerProvider(resource=resource, shutdown_on_exit=False)

        # Build the exporter chain:
        # KliraOTLPSpanExporter → FilteringSpanExporter → BatchSpanProcessor
        # → NoneAttributeFilterProcessor
        klira_exporter = KliraOTLPSpanExporter(
            endpoint=config.get_telemetry_endpoint(),
            api_key=config.api_key,
            anonymization=config.anonymization is not None,
            phi_method=config.anonymization,
            phi_export_entity_details=config.phi_export_entity_details,
        )
        filtering_exporter = FilteringSpanExporter(klira_exporter)
        batch_processor = BatchSpanProcessor(
            filtering_exporter, schedule_delay_millis=500
        )
        none_filter = NoneAttributeFilterProcessor(batch_processor)

        provider.add_span_processor(none_filter)
        trace.set_tracer_provider(provider)

        # Verify the provider was actually installed
        active = trace.get_tracer_provider()
        if active is not provider:
            logger.warning(
                "TracerProvider installation may have failed: "
                "active provider (%s) differs from Klira provider. "
                "Another SDK may have set a provider first.",
                type(active).__name__,
            )

        cls._provider = provider

        if not cls._atexit_registered:
            atexit.register(cls._atexit_shutdown)
            cls._atexit_registered = True

    @classmethod
    def _load_policies(cls, config: KliraConfig) -> None:
        """Load guardrails policies into the GuardrailsEngine singleton.

        Uses remote loader by default, falling back to local filesystem.
        Fail-open: if loading fails, logs warning and continues with empty policies.
        """
        from klira.sdk.guardrails.engine import GuardrailsEngine
        from klira.sdk.guardrails.policy_loader import (
            FileSystemPolicyLoader,
            RemotePolicyLoader,
        )

        policies: list = []
        try:
            if config.use_remote_policies:
                policies = RemotePolicyLoader(
                    api_key=config.api_key,
                    endpoint=config.policies_endpoint,
                    fallback_path=config.policies_path,
                ).load()
            elif config.policies_path:
                policies = FileSystemPolicyLoader(config.policies_path).load()
        except Exception:
            logger.warning("Failed to load policies, continuing with empty set", exc_info=True)

        # Construct LLM fallback evaluator from config if fully configured
        llm_evaluator = None
        if (
            config.llm_fallback_enabled
            and config.llm_fallback_provider
            and config.llm_fallback_model
            and config.llm_fallback_api_key
        ):
            try:
                from klira.sdk.guardrails.llm_evaluator import (
                    BuiltInLLMFallbackEvaluator,
                )

                llm_evaluator = BuiltInLLMFallbackEvaluator.from_config(
                    provider=config.llm_fallback_provider,
                    model=config.llm_fallback_model,
                    api_key=config.llm_fallback_api_key,
                )
            except (ImportError, ValueError) as e:
                logger.warning("Failed to create LLM fallback evaluator: %s", e)
        elif config.llm_fallback_enabled:
            logger.warning(
                "llm_fallback_enabled=True but not fully configured. "
                "Required: llm_fallback_provider, llm_fallback_model, llm_fallback_api_key."
            )

        engine = GuardrailsEngine(
            policies=policies,
            llm_fallback_enabled=config.llm_fallback_enabled,
            llm_evaluator=llm_evaluator,
            llm_fallback_on_error=config.llm_fallback_on_error,
        )
        GuardrailsEngine.set_instance(engine)
        logger.debug("Loaded %d guardrails policies", len(policies))

    @classmethod
    def _auto_patch_llm_clients(cls) -> None:
        """Auto-detect and patch available LLM client libraries.

        Tries to import each supported library and patches it so that
        LLM API calls are wrapped with klira.llm.{provider} spans.
        Libraries that are not installed are silently skipped.
        """
        # OpenAI Chat Completions
        try:
            import openai

            from klira.sdk.adapters.openai_completion import OpenAICompletionAdapter

            OpenAICompletionAdapter().patch(openai.OpenAI)
            logger.debug("Patched openai.OpenAI.chat.completions.create")
        except ImportError:
            pass
        except (AttributeError, RuntimeError) as e:
            logger.warning("Failed to patch openai.completion: %s", e)

        # OpenAI Responses API
        try:
            import openai  # noqa: F811

            from klira.sdk.adapters.openai_responses import OpenAIResponsesAdapter

            OpenAIResponsesAdapter().patch(openai.OpenAI)
            logger.debug("Patched openai.OpenAI.responses.create")
        except ImportError:
            pass
        except (AttributeError, RuntimeError) as e:
            logger.warning("Failed to patch openai.responses: %s", e)

        # Anthropic
        try:
            import anthropic

            from klira.sdk.adapters.anthropic import AnthropicAdapter

            AnthropicAdapter().patch(anthropic.Anthropic)
            logger.debug("Patched anthropic.Anthropic.messages.create")
        except ImportError:
            pass
        except (AttributeError, RuntimeError) as e:
            logger.warning("Failed to patch anthropic: %s", e)

        # Google Gemini
        try:
            import google.generativeai as genai

            from klira.sdk.adapters.gemini import GeminiAdapter

            GeminiAdapter().patch(genai.GenerativeModel)
            logger.debug("Patched google.generativeai.GenerativeModel.generate_content")
        except ImportError:
            pass
        except (AttributeError, RuntimeError) as e:
            logger.warning("Failed to patch gemini: %s", e)

        # Ollama
        try:
            import ollama as ollama_mod  # noqa: F811

            from klira.sdk.adapters.ollama import OllamaAdapter

            OllamaAdapter().patch(ollama_mod.Client)
            logger.debug("Patched ollama.Client.chat")
        except ImportError:
            pass
        except (AttributeError, RuntimeError) as e:
            logger.warning("Failed to patch ollama: %s", e)

        # LiteLLM (module-level function, not a class method)
        try:
            import litellm

            from klira.sdk.adapters.litellm import LiteLLMAdapter

            LiteLLMAdapter().patch(litellm)  # type: ignore[arg-type]
            logger.debug("Patched litellm.completion")
        except ImportError:
            pass
        except (AttributeError, RuntimeError) as e:
            logger.warning("Failed to patch litellm: %s", e)

    @classmethod
    def get_config(cls) -> KliraConfig:
        """Get the current SDK configuration.

        Raises:
            KliraInitError: If init() has not been called.
        """
        if cls._config is None:
            raise KliraInitError("Klira.init() must be called before using the SDK")
        return cls._config

    @classmethod
    def is_initialized(cls) -> bool:
        """Whether the SDK has been initialized."""
        return cls._initialized

    @classmethod
    def shutdown(cls) -> None:
        """Shut down the SDK and flush pending traces."""
        if cls._provider is not None:
            cls._provider.force_flush(timeout_millis=5000)
            cls._provider.shutdown()
            cls._provider = None
        cls._config = None
        cls._initialized = False

    @classmethod
    def _atexit_shutdown(cls) -> None:
        """atexit handler: flush and shut down OTel before C-extensions are torn down."""
        if cls._provider is not None:
            try:
                cls._provider.force_flush(timeout_millis=5000)
            except Exception:
                pass  # Best-effort during interpreter shutdown
            try:
                cls._provider.shutdown()
            except Exception:
                pass  # Best-effort during interpreter shutdown
            cls._provider = None
        cls._config = None
        cls._initialized = False

    @classmethod
    @contextmanager
    def session(cls, **kwargs: Any) -> Iterator[KliraConfig]:
        """Context manager for SDK lifecycle.

        Usage:
            with Klira.session(api_key="klira_...") as config:
                # SDK is initialized
                ...
            # SDK is shut down automatically
        """
        config = cls.init(**kwargs)
        try:
            yield config
        finally:
            cls.shutdown()

    @classmethod
    def force_flush(cls, timeout_millis: int = 5000) -> bool:
        """Force flush pending traces."""
        if cls._provider is not None:
            return cls._provider.force_flush(timeout_millis)
        return True

    @classmethod
    def _reset(cls) -> None:
        """Reset the SDK state. For testing only."""
        cls.shutdown()
