#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
#
#  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:
#
#       https://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.
"""Base class for log stores."""

import logging
import threading
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Any, Dict, List, Optional, Type, cast

from zenml.enums import StackComponentType
from zenml.models import LogsResponse
from zenml.stack import Flavor, StackComponent, StackComponentConfig
from zenml.utils.logging_utils import LogEntry


class BaseLogStoreConfig(StackComponentConfig):
    """Base configuration for all log stores."""


class BaseLogStoreOrigin:
    """Base class for all ZenML log store origins.

    The origin is the entry point for all log records to be sent to the log
    store for processing. The process of sending a log record is as follows:

    1. instantiate the log store or use the active log store
    2. register an origin by calling log_store.register_origin() and passing
    the log model and optional metadata to be attached to each log record
    3. emit the log record by calling log_store.emit() and passing the origin
    and log record
    4. deregister the origin when all logs have been emitted by calling
    log_store.deregister(origin)
    """

    def __init__(
        self,
        name: str,
        log_store: "BaseLogStore",
        log_model: LogsResponse,
        metadata: Dict[str, Any],
    ) -> None:
        """Initialize a log store origin.

        Args:
            name: The name of the origin.
            log_store: The log store to emit logs to.
            log_model: The log model associated with the origin.
            metadata: Additional metadata to attach to all log entries that will
                be emitted by this origin.
        """
        self.name = name
        self.log_store = log_store
        self.log_model = log_model
        self.metadata = metadata


class BaseLogStore(StackComponent, ABC):
    """Base class for all ZenML log stores.

    A log store is responsible for collecting, storing, and retrieving logs
    during pipeline and step execution. Different implementations may store
    logs in different backends (artifact store, OpenTelemetry, Datadog, etc.).
    """

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the log store.

        Args:
            *args: Positional arguments for the base class.
            **kwargs: Keyword arguments for the base class.
        """
        super().__init__(*args, **kwargs)
        self._origins: Dict[str, BaseLogStoreOrigin] = {}
        self._lock = threading.RLock()

    @property
    def config(self) -> BaseLogStoreConfig:
        """Returns the configuration of the log store.

        Returns:
            The configuration.
        """
        return cast(BaseLogStoreConfig, self._config)

    @property
    def origin_class(self) -> Type[BaseLogStoreOrigin]:
        """Class of the origin.

        Returns:
            The class of the origin used with this log store.
        """
        return BaseLogStoreOrigin

    def register_origin(
        self, name: str, log_model: LogsResponse, metadata: Dict[str, Any]
    ) -> BaseLogStoreOrigin:
        """Register an origin for the log store.

        Args:
            name: The name of the origin.
            log_model: The log model associated with the origin.
            metadata: Additional metadata to attach to the log entry.

        Returns:
            The origin.
        """
        with self._lock:
            origin = self.origin_class(name, self, log_model, metadata)
            self._origins[name] = origin
            return origin

    def deregister_origin(
        self,
        origin: BaseLogStoreOrigin,
        blocking: bool = True,
    ) -> None:
        """Deregister an origin previously registered with the log store.

        If no other origins are left, the log store will be flushed. The
        `blocking` parameter determines whether to block until the flush is
        complete.

        Args:
            origin: The origin to deregister.
            blocking: Whether to block until the deregistration is complete
                and all logs are flushed if this is the last origin registered.
        """
        with self._lock:
            if origin.name not in self._origins:
                return
            self._release_origin(origin)
            del self._origins[origin.name]
            if len(self._origins) == 0:
                self.flush(blocking=blocking)

    @abstractmethod
    def emit(
        self,
        origin: BaseLogStoreOrigin,
        record: logging.LogRecord,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Process a log record from the logging system.

        Args:
            origin: The origin used to send the log record.
            record: The Python logging.LogRecord to process.
            metadata: Additional metadata to attach to the log entry.
        """

    @abstractmethod
    def _release_origin(
        self,
        origin: BaseLogStoreOrigin,
    ) -> None:
        """Finalize the stream of log records associated with an origin.

        This is used to announce the end of the stream of log records associated
        with an origin and that no more log records will be emitted.

        The implementation should ensure that all log records associated with
        the origin are flushed to the backend and any resources (clients,
        connections, file descriptors, etc.) are released.

        Args:
            origin: The origin to finalize.
        """

    @abstractmethod
    def flush(self, blocking: bool = True) -> None:
        """Flush the log store.

        This method is called to ensure that all logs are flushed to the backend.

        Args:
            blocking: Whether to block until the flush is complete.
        """

    @abstractmethod
    def fetch(
        self,
        logs_model: LogsResponse,
        limit: int,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
    ) -> List[LogEntry]:
        """Fetch logs from the log store.

        This method is called from the server to retrieve logs for display
        on the dashboard or via API. The implementation should not require
        any integration-specific SDKs that aren't available on the server.

        Args:
            logs_model: The logs model containing metadata about the logs.
            start_time: Filter logs after this time.
            end_time: Filter logs before this time.
            limit: Maximum number of log entries to return.

        Returns:
            List of log entries matching the query.
        """


class BaseLogStoreFlavor(Flavor):
    """Base class for all ZenML log store flavors."""

    @property
    def type(self) -> StackComponentType:
        """Type of the flavor.

        Returns:
            The type of the flavor.
        """
        return StackComponentType.LOG_STORE

    @property
    def config_class(self) -> Type[BaseLogStoreConfig]:
        """Config class for the base log store flavor.

        Returns:
            The config class.
        """
        return BaseLogStoreConfig

    @property
    @abstractmethod
    def implementation_class(self) -> Type["BaseLogStore"]:
        """Implementation class for the base log store flavor.

        Returns:
            The implementation class.
        """
