import importlib
import logging
from collections.abc import Callable
from typing import Literal, TypedDict

from ..sdk import VerseSDK

ANTHROPIC_MODULE = "anthropic"
LANGCHAIN_MODULE = "langchain"
LITELLM_MODULE = "litellm"
GOOGLE_MODULE = "google"
PYDANTIC_AI_MODULE = "pydantic_ai"

VENDOR_ALLOWLIST = Literal[
    ANTHROPIC_MODULE,
    LANGCHAIN_MODULE,
    LITELLM_MODULE,
    GOOGLE_MODULE,
    PYDANTIC_AI_MODULE,
]

VENDOR_ALLOWLIST_CONFIG = list[VENDOR_ALLOWLIST] | VENDOR_ALLOWLIST | None


class VendorIntegration(TypedDict):
    """Configuration for a vendor integration."""

    module_name: str
    configure_func: str
    cleanup_func: str | None


INTEGRATIONS: dict[str, VendorIntegration] = {
    ANTHROPIC_MODULE: {
        "module_name": ANTHROPIC_MODULE,
        "configure_func": "configure_anthropic",
        "cleanup_func": None,
    },
    LANGCHAIN_MODULE: {
        "module_name": LANGCHAIN_MODULE,
        "configure_func": "configure_langchain",
        "cleanup_func": None,
    },
    LITELLM_MODULE: {
        "module_name": LITELLM_MODULE,
        "configure_func": "configure_litellm",
        "cleanup_func": "cleanup_litellm",
    },
    GOOGLE_MODULE: {
        "module_name": GOOGLE_MODULE,
        "configure_func": "configure_google",
        "cleanup_func": None,
    },
    PYDANTIC_AI_MODULE: {
        "module_name": PYDANTIC_AI_MODULE,
        "configure_func": "patch_pydantic_ai",
        "cleanup_func": None,
    },
}


def _load_integration(integration: VendorIntegration, sdk: VerseSDK) -> Callable | None:
    """
    Load and apply a vendor integration.

    Returns the cleanup function if one exists, otherwise None.
    """
    # Use importlib.import_module instead of __import__ for better relative import handling
    module_name = f"verse_sdk.integrations.{integration['module_name']}"
    module = importlib.import_module(module_name)

    getattr(module, integration["configure_func"])(sdk)

    if integration["cleanup_func"]:
        return getattr(module, integration["cleanup_func"])
    return None


def apply_integrations(
    sdk: VerseSDK,
    vendors: VENDOR_ALLOWLIST_CONFIG,
) -> Callable:
    """
    Apply vendor integrations to the VerseSDK.

    Args:
        sdk: The VerseSDK instance to configure
        vendors: Vendor name(s) to integrate, or None for no integrations

    Returns:
        A cleanup function that should be called to teardown integrations
    """
    cleanup_functions = []

    def _cleanup_integrations():
        """Cleanup all applied integrations."""
        for cleanup_fn in cleanup_functions:
            try:
                cleanup_fn()
            except Exception as e:
                logging.warning(
                    f"Error cleaning up `{cleanup_fn.__name__}`", exc_info=e
                )

    if not vendors:
        return _cleanup_integrations

    vendor_list = [vendors] if isinstance(vendors, str) else vendors

    for vendor in vendor_list:
        vendor = vendor.lower()

        integration = INTEGRATIONS.get(vendor)
        if not integration:
            logging.warning(f"Unknown vendor integration: `{vendor}`")
            continue

        try:
            cleanup_fn = _load_integration(integration, sdk)
            if cleanup_fn:
                cleanup_functions.append(cleanup_fn)
            logging.info(f"`{vendor}` integration applied to `VerseSDK`")
        except Exception as e:
            logging.warning(f"Error applying `{vendor}` integration", exc_info=e)

    return _cleanup_integrations


__all__ = [
    "VENDOR_ALLOWLIST_CONFIG",
    "apply_integrations",
]
