# 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.

"""FrameworkRegistry and LLMClientRegistry.

Uses regular dict (not WeakValueDictionary) for LLMClientRegistry (PROD-236).
"""

from __future__ import annotations

import logging

from klira.sdk.adapters.base_framework import BaseFrameworkAdapter
from klira.sdk.adapters.base_llm import BaseLLMAdapter

logger = logging.getLogger(__name__)


class FrameworkRegistry:
    """Registry for framework adapters."""

    def __init__(self) -> None:
        self._adapters: dict[str, BaseFrameworkAdapter] = {}

    def register(self, name: str, adapter: BaseFrameworkAdapter) -> None:
        """Register a framework adapter."""
        self._adapters[name] = adapter
        logger.debug("Registered framework adapter: %s", name)

    def get(self, name: str) -> BaseFrameworkAdapter | None:
        """Get a registered framework adapter by name."""
        return self._adapters.get(name)

    def get_all(self) -> dict[str, BaseFrameworkAdapter]:
        """Get all registered framework adapters."""
        return dict(self._adapters)

    def clear(self) -> None:
        """Clear all registrations."""
        self._adapters.clear()


class LLMClientRegistry:
    """Registry for LLM client adapters.

    Uses regular dict, NOT WeakValueDictionary (PROD-236).
    """

    def __init__(self) -> None:
        self._adapters: dict[str, BaseLLMAdapter] = {}

    def register(self, name: str, adapter: BaseLLMAdapter) -> None:
        """Register an LLM client adapter."""
        self._adapters[name] = adapter
        logger.debug("Registered LLM adapter: %s", name)

    def get(self, name: str) -> BaseLLMAdapter | None:
        """Get a registered LLM adapter by name."""
        return self._adapters.get(name)

    def get_all(self) -> dict[str, BaseLLMAdapter]:
        """Get all registered LLM adapters."""
        return dict(self._adapters)

    def clear(self) -> None:
        """Clear all registrations."""
        self._adapters.clear()


# Singleton instances
_framework_registry = FrameworkRegistry()
_llm_registry = LLMClientRegistry()


def get_framework_registry() -> FrameworkRegistry:
    return _framework_registry


def get_llm_registry() -> LLMClientRegistry:
    return _llm_registry
