"""Serializable component infrastructure for Agentbyte."""

from __future__ import annotations

import importlib
from typing import Any, ClassVar, Generic, TypeVar, cast

from pydantic import BaseModel, ConfigDict

ConfigT = TypeVar("ConfigT", bound=BaseModel)
ExpectedT = TypeVar("ExpectedT")


class ComponentModel(BaseModel):
    """Portable envelope for serializing loadable Agentbyte components."""

    provider: str
    component_type: str | None = None
    version: int = 1
    component_version: int | None = None
    description: str | None = None
    label: str | None = None
    config: dict[str, Any]

    model_config = ConfigDict(extra="forbid")


WELL_KNOWN_PROVIDERS: dict[str, str] = {
    # LLM clients
    "model_client": "agentbyte.llm.OpenAIChatCompletionClient",
    "openai_client": "agentbyte.llm.OpenAIChatCompletionClient",
    "OpenAIChatCompletionClient": "agentbyte.llm.OpenAIChatCompletionClient",
    "azure_openai_client": "agentbyte.llm.AzureOpenAIChatCompletionClient",
    "AzureOpenAIChatCompletionClient": "agentbyte.llm.AzureOpenAIChatCompletionClient",
}


class ComponentFromConfig(Generic[ConfigT]):
    """Mixin for components that can be restored from validated config."""

    @classmethod
    def _from_config(cls, config: ConfigT):
        raise NotImplementedError


class ComponentToConfig(Generic[ConfigT]):
    """Mixin for components that can be dumped to validated config."""

    component_type: ClassVar[str]
    component_version: ClassVar[int] = 1
    component_provider_override: ClassVar[str | None] = None
    component_description: ClassVar[str | None] = None
    component_label: ClassVar[str | None] = None

    def _to_config(self) -> ConfigT:
        raise NotImplementedError

    def dump_component(self) -> ComponentModel:
        provider = self.component_provider_override or (
            f"{self.__class__.__module__}.{self.__class__.__qualname__}"
        )
        if "<locals>" in provider:
            msg = "Cannot serialize locally defined classes"
            raise TypeError(msg)

        description = self.component_description
        if description is None and self.__class__.__doc__:
            description = self.__class__.__doc__.strip().split("\n\n")[0].strip()

        return ComponentModel(
            provider=provider,
            component_type=self.component_type,
            version=self.component_version,
            component_version=self.component_version,
            description=description,
            label=self.component_label or self.__class__.__name__,
            config=self._to_config().model_dump(exclude_none=True),
        )


class ComponentLoader:
    """Loader for `ComponentModel` envelopes."""

    @classmethod
    def load_component(
        cls,
        model: ComponentModel | dict[str, Any],
        expected: type[ExpectedT] | None = None,
    ):
        loaded = model if isinstance(model, ComponentModel) else ComponentModel(**model)
        provider = WELL_KNOWN_PROVIDERS.get(loaded.provider, loaded.provider)
        module_path, _, class_name = provider.rpartition(".")
        if not module_path or not class_name:
            msg = f"Invalid provider string: {provider}"
            raise ValueError(msg)

        module = importlib.import_module(module_path)
        component_class = getattr(module, class_name)

        if not hasattr(component_class, "component_config_schema"):
            msg = f"Component class {component_class} is missing component_config_schema"
            raise AttributeError(msg)
        if not hasattr(component_class, "_from_config"):
            msg = f"Component class {component_class} does not support loading"
            raise TypeError(msg)

        schema = component_class.component_config_schema
        validated_config = schema.model_validate(loaded.config)
        instance = component_class._from_config(validated_config)

        if expected is not None and not isinstance(instance, expected):
            msg = f"Loaded component {type(instance)} does not match expected type {expected}"
            raise TypeError(msg)
        if expected is None and not isinstance(instance, cls):
            return instance
        return cast(ExpectedT, instance)


class Component(
    Generic[ConfigT],
    ComponentFromConfig[ConfigT],
    ComponentToConfig[ConfigT],
    ComponentLoader,
):
    """Base class for loadable/dumpable Agentbyte components."""

    component_config_schema: ClassVar[type[ConfigT]]

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        if cls.__name__ == "Component":
            return
        if not hasattr(cls, "component_type"):
            msg = f"{cls.__name__} must define component_type"
            raise TypeError(msg)
        if not hasattr(cls, "component_config_schema"):
            msg = f"{cls.__name__} must define component_config_schema"
            raise TypeError(msg)


__all__ = [
    "Component",
    "ComponentFromConfig",
    "ComponentLoader",
    "ComponentModel",
    "ComponentToConfig",
    "WELL_KNOWN_PROVIDERS",
]
