"""Auto-instrumentation manager for Robin AI.

Handles automatic instrumentation of supported libraries (FastAPI, SQLAlchemy, Redis, etc.)
based on configuration flags.
"""

from typing import Any

from fastapi import FastAPI
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode

from robin_commons.log import logger
from robin_commons.telemetry.config import get_observability_config


class InstrumentationManager:
    """Manager for setting up library auto-instrumentation."""

    def __init__(self) -> None:
        """Initialize instrumentation manager with configuration."""
        self.config = get_observability_config()
        self._instrumented_services: set[str] = set()

    def setup_all_instrumentation(self, app: FastAPI | None = None) -> None:
        """Set up all enabled instrumentations.

        Args:
            app: Optional FastAPI application instance to instrument
        """
        if not self.config.is_observability_enabled():
            logger.info("Observability disabled, skipping instrumentation")
            return

        self.setup_fastapi_instrumentation(app)
        self.setup_sqlalchemy_instrumentation()
        self.setup_redis_instrumentation()
        self.setup_httpx_instrumentation()
        self.setup_nats_instrumentation()
        self.setup_grpc_instrumentation()

        logger.info(
            "Instrumentation setup complete",
            instrumented_services=list(self._instrumented_services),
        )

    def setup_fastapi_instrumentation(self, app: FastAPI | None) -> None:
        """Instrument FastAPI application."""
        if not app or not self.config.is_fastapi_instrumentation_enabled():
            return

        try:
            from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

            FastAPIInstrumentor.instrument_app(
                app,
                tracer_provider=trace.get_tracer_provider(),
                excluded_urls="health,metrics,docs,redoc,openapi.json",
            )
            self._instrumented_services.add("fastapi")
        except Exception as e:
            logger.error(f"Failed to setup FastAPI instrumentation: {e}")

    def setup_sqlalchemy_instrumentation(self) -> None:
        """Instrument SQLAlchemy engine."""
        if not self.config.is_sqlalchemy_instrumentation_enabled():
            return

        try:
            from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor

            SQLAlchemyInstrumentor().instrument(
                enable_commenter=True,
                commenter_options={"db_driver": True, "db_framework": True},
            )
            self._instrumented_services.add("sqlalchemy")
        except Exception as e:
            logger.error(f"Failed to setup SQLAlchemy instrumentation: {e}")

    def setup_redis_instrumentation(self) -> None:
        """Instrument Redis client."""
        if not self.config.is_redis_instrumentation_enabled():
            return

        try:
            from opentelemetry.instrumentation.redis import RedisInstrumentor

            RedisInstrumentor().instrument()
            self._instrumented_services.add("redis")
        except Exception as e:
            logger.error(f"Failed to setup Redis instrumentation: {e}")

    def setup_httpx_instrumentation(self) -> None:
        """Instrument HTTPX client."""
        if not self.config.is_httpx_instrumentation_enabled():
            return

        try:
            from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor

            HTTPXClientInstrumentor().instrument(
                request_hook=self._http_client_request_hook,
                response_hook=self._http_client_response_hook,
            )
            self._instrumented_services.add("httpx")
        except Exception as e:
            logger.warning(
                f"HTTPX instrumentation not available: {e}. Trying requests fallback."
            )
            # Fallback to requests if httpx fails (optional)
            self.setup_requests_instrumentation()

    def setup_requests_instrumentation(self) -> None:
        """Instrument requests library (fallback)."""
        try:
            from opentelemetry.instrumentation.requests import RequestsInstrumentor

            RequestsInstrumentor().instrument()
            self._instrumented_services.add("requests")
        except Exception as e:
            logger.error(f"Failed to setup requests instrumentation: {e}")

    def setup_nats_instrumentation(self) -> None:
        """Instrument NATS client."""
        if not self.config.is_nats_instrumentation_enabled():
            return

        # NATS instrumentation might require custom handling or external libraries
        # Placeholder for future implementation or specific library usage
        pass

    def setup_grpc_instrumentation(self) -> None:
        """Instrument gRPC client/server."""
        if not self.config.is_grpc_instrumentation_enabled():
            return

        try:
            from opentelemetry.instrumentation.grpc import (
                GrpcAioInstrumentorClient,
                GrpcAioInstrumentorServer,
                GrpcInstrumentorClient,
                GrpcInstrumentorServer,
            )

            # Instrument both sync and async
            GrpcInstrumentorClient().instrument()
            GrpcInstrumentorServer().instrument()
            GrpcAioInstrumentorClient().instrument()
            GrpcAioInstrumentorServer().instrument()
            self._instrumented_services.add("grpc")
        except ImportError:
            # gRPC dependencies might not be installed
            pass
        except Exception as e:
            logger.error(f"Failed to setup gRPC instrumentation: {e}")

    @staticmethod
    def _http_client_request_hook(span: trace.Span, request: Any) -> None:
        """Hook to add request attributes to HTTP client spans."""
        if hasattr(request, "headers"):
            # Don't log sensitive headers
            pass

    @staticmethod
    def _http_client_response_hook(
        span: trace.Span, request: Any, response: Any
    ) -> None:
        """Hook to add response attributes to HTTP client spans."""
        if hasattr(response, "status_code"):
            span.set_attribute("http.status_code", response.status_code)
            if response.status_code >= 400:
                span.set_status(Status(StatusCode.ERROR))


# Global instrumentation instance
_instrumentation_manager: InstrumentationManager | None = None


def get_instrumentation_manager() -> InstrumentationManager:
    """Get the global instrumentation manager instance."""
    global _instrumentation_manager
    if _instrumentation_manager is None:
        _instrumentation_manager = InstrumentationManager()
    return _instrumentation_manager
