import logging
from abc import ABC
from typing import Any, Callable, List, Optional, Union

from sqlalchemy import URL

from dbgpt.component import BaseComponent, ComponentType, SystemApp
from dbgpt.storage.metadata import DatabaseManager

logger = logging.getLogger(__name__)


class BaseServe(BaseComponent, ABC):
    """Base serve component for DB-GPT"""

    name = "dbgpt_serve_base"

    def __init__(
        self,
        system_app: SystemApp,
        api_prefix: str | List[str],
        api_tags: List[str],
        db_url_or_db: Union[str, URL, DatabaseManager] = None,
        try_create_tables: Optional[bool] = False,
    ):
        self._system_app = system_app
        self._api_prefix = api_prefix
        self._api_tags = api_tags
        self._db_url_or_db = db_url_or_db
        self._try_create_tables = try_create_tables
        self._not_create_table = True
        self._app_has_initiated = False

    def create_or_get_db_manager(self) -> DatabaseManager:
        """Create or get the database manager.
        This method must be called after the application is initialized

        Returns:
            DatabaseManager: The database manager
        """
        from dbgpt.storage.metadata import Model, UnifiedDBManagerFactory, db

        # If you need to use the database, you can get the database manager here
        db_manager_factory: UnifiedDBManagerFactory = self._system_app.get_component(
            ComponentType.UNIFIED_METADATA_DB_MANAGER_FACTORY,
            UnifiedDBManagerFactory,
            default_component=None,
        )
        if db_manager_factory is not None and db_manager_factory.create():
            init_db = db_manager_factory.create()
        else:
            init_db = self._db_url_or_db or db
            init_db = DatabaseManager.build_from(init_db, base=Model)

        if self._try_create_tables and self._not_create_table:
            try:
                init_db.create_all()
            except Exception as e:
                logger.warning(f"Failed to create tables: {e}")
            finally:
                self._not_create_table = False
        return init_db

    @classmethod
    def get_current_serve(cls, system_app: SystemApp) -> Optional["BaseServe"]:
        """Get the current serve component.

        None if the serve component is not exist.

        Args:
            system_app (SystemApp): The system app

        Returns:
            Optional[BaseServe]: The current serve component.
        """
        return cls.get_instance(system_app, default_component=None)

    @classmethod
    def call_on_current_serve(
        cls,
        system_app: SystemApp,
        func: Callable[["BaseServe"], Optional[Any]],
        default_value: Optional[Any] = None,
    ) -> Optional[Any]:
        """Call the function on the current serve component.

        Return default_value if the serve component is not exist or the function return
        None.

        Args:
            system_app (SystemApp): The system app
            func (Callable[[BaseServe], Any]): The function to call
            default_value (Optional[Any], optional): The default value. Defaults to
                None.

        Returns:
            Optional[Any]: The result of the function
        """
        serve = cls.get_current_serve(system_app)
        if not serve:
            return default_value
        result = func(serve)
        if not result:
            result = default_value
        return result
