from importlib.metadata import version, PackageNotFoundError
from typing import Dict
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor

from .crewai import CrewAIInstrumentor
from .litellm import LitellmInstrumentor
from .portkey_span_exporter import PortkeySpanExporter
# from .langgraph import LanggraphInstrumentor

__all__ = ["initialize_instrumentation"]

package_instrumentor_map: Dict[str, BaseInstrumentor] = {
    "crewai": CrewAIInstrumentor,
    "litellm": LitellmInstrumentor,
    # "langgraph": LanggraphInstrumentor,
}


def is_package_installed(pkg_name):
    try:
        version(pkg_name)
        return True
    except PackageNotFoundError:
        return False


def initialize_instrumentation(api_key: str, base_url: str):
    tracer_provider = TracerProvider()
    exporter = PortkeySpanExporter(api_key=api_key, base_url=base_url)
    tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
    trace.set_tracer_provider(tracer_provider)
    for package, instrumentor in package_instrumentor_map.items():
        if is_package_installed(package):
            instrumentor().instrument()
            print(f"Obiguard: {package} Instrumentation initialized")
