#  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.
"""Utility functions for logging."""

import logging
import os
import threading
from contextvars import ContextVar
from datetime import datetime
from types import TracebackType
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union, cast
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field

from zenml.client import Client
from zenml.config.pipeline_configurations import PipelineConfiguration
from zenml.config.step_configurations import StepConfiguration
from zenml.constants import (
    ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE,
    ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
    ENV_ZENML_SERVER,
    handle_bool_env_var,
)
from zenml.enums import LoggingLevels, StackComponentType
from zenml.exceptions import DoesNotExistException
from zenml.logger import get_logger
from zenml.models import (
    LogsRequest,
    LogsResponse,
    LogsUpdate,
    PipelineRunResponse,
    StepRunResponse,
)
from zenml.utils import context_utils

if TYPE_CHECKING:
    from zenml.log_stores.base_log_store import BaseLogStoreOrigin
    from zenml.zen_stores.base_zen_store import BaseZenStore

logger = get_logger(__name__)


class LogEntry(BaseModel):
    """A structured log entry with parsed information.

    This is used in two distinct ways:
        1. If we are using the artifact log store, we save the
        entries as JSON-serialized LogEntry's in the artifact store.
        2. When queried, the server returns logs as a list of LogEntry's.
    """

    message: str = Field(description="The log message content")
    name: Optional[str] = Field(
        default=None,
        description="The name of the logger",
    )
    level: Optional[LoggingLevels] = Field(
        default=None,
        description="The log level",
    )
    timestamp: Optional[datetime] = Field(
        default=None,
        description="When the log was created",
    )
    module: Optional[str] = Field(
        default=None, description="The module that generated this log entry"
    )
    filename: Optional[str] = Field(
        default=None,
        description="The name of the file that generated this log entry",
    )
    lineno: Optional[int] = Field(
        default=None, description="The fileno that generated this log entry"
    )
    chunk_index: int = Field(
        default=0,
        description="The index of the chunk in the log entry",
    )
    total_chunks: int = Field(
        default=1,
        description="The total number of chunks in the log entry",
    )
    id: UUID = Field(
        default_factory=uuid4,
        description="The unique identifier of the log entry",
    )

    model_config = ConfigDict(
        # ignore extra attributes during model initialization
        extra="ignore",
    )


class LoggingContext(context_utils.BaseContext):
    """Context manager which collects logs using a LogStore."""

    __context_var__ = ContextVar("logging_context")

    def __init__(
        self,
        name: str,
        log_model: "LogsResponse",
        block_on_exit: bool = True,
        **metadata: Any,
    ) -> None:
        """Initialize the logging context.

        Args:
            name: The name of the logging context.
            log_model: The logs response model for this context.
            block_on_exit: Whether to block until all logs are flushed when the
                context is exited, if there are no more logging contexts active.
            **metadata: Additional metadata to attach to the log entry.
        """
        self.log_model = log_model
        self._lock = threading.Lock()
        self._disabled = False
        self._log_store = Client().active_stack.log_store
        self._metadata = metadata
        self._origin: Optional["BaseLogStoreOrigin"] = None
        self._name = name
        self._block_on_exit = block_on_exit

    @property
    def name(self) -> str:
        """The name of the logging context.

        Returns:
            The name of the logging context.
        """
        return self._name

    @classmethod
    def emit(
        cls,
        record: logging.LogRecord,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Emit a log record using the active logging context.

        This class method is called by stdout/stderr wrappers and logging
        handlers to route logs to the active log store.

        Args:
            record: The log record to emit.
            metadata: Additional metadata to attach to the log entry.
        """
        if context := LoggingContext.get():
            if context._disabled:
                return
            context._disabled = True
            try:
                message = record.getMessage()
                if message and message.strip():
                    if context._origin:
                        context._log_store.emit(
                            origin=context._origin,
                            record=record,
                            metadata=metadata,
                        )
            except Exception:
                logger.debug("Failed to emit log record", exc_info=True)
            finally:
                context._disabled = False

    def update(
        self,
        step_run: Optional["StepRunResponse"] = None,
        pipeline_run: Optional["PipelineRunResponse"] = None,
    ) -> None:
        """Update the metadata of the logging context.

        Args:
            step_run: The step run.
            pipeline_run: The pipeline run.
        """
        if step_run and self._origin:
            self._origin.metadata.update(
                get_step_log_metadata(step_run=step_run)
            )
        if pipeline_run and self._origin:
            self._origin.metadata.update(
                get_run_log_metadata(pipeline_run=pipeline_run)
            )

    def attach(
        self, response: Union["StepRunResponse", "PipelineRunResponse"]
    ) -> None:
        """Attach the logging context to a step run or pipeline run.

        Args:
            response: The step run or pipeline run.

        Raises:
            ValueError: If the response is not a step run or pipeline run.
        """
        if isinstance(response, StepRunResponse):
            log_update = LogsUpdate(step_run_id=response.id)
        elif isinstance(response, PipelineRunResponse):
            log_update = LogsUpdate(pipeline_run_id=response.id)
        else:
            raise ValueError("Response must be a step run or pipeline run.")

        Client().zen_store.update_logs(
            logs_id=self.log_model.id, logs_update=log_update
        )

    def __enter__(self) -> "LoggingContext":
        """Enter the context and set as active.

        Returns:
            self
        """
        with self._lock:
            super().__enter__()
            self._origin = self._log_store.register_origin(
                name=self.name,
                log_model=self.log_model,
                metadata=self._metadata,
            )

        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        """Exit the context and restore previous context.

        Args:
            exc_type: The class of the exception.
            exc_val: The instance of the exception.
            exc_tb: The traceback of the exception.
        """
        if exc_type is not None:
            LoggingContext.emit(
                logging.LogRecord(
                    name="",
                    level=logging.ERROR,
                    msg="An exception has occurred.",
                    args=(),
                    exc_info=(exc_type, exc_val, exc_tb) if exc_val else None,
                    func=None,
                    pathname="",
                    lineno=0,
                ),
                metadata={"zenml.event.type": "exception"},
            )

        with self._lock:
            super().__exit__(exc_type, exc_val, exc_tb)
            if self._origin:
                self._log_store.deregister_origin(
                    self._origin, blocking=self._block_on_exit
                )
                self._origin = None


def generate_logs_request(source: str) -> LogsRequest:
    """Generate a LogsRequest for logging.

    Args:
        source: The source of the logs (e.g., "client", "orchestrator", "step").

    Returns:
        A LogsRequest object.
    """
    from zenml.log_stores.artifact.artifact_log_store import (
        ArtifactLogStore,
        prepare_logs_uri,
    )

    client = Client()
    log_store = client.active_stack.log_store
    log_id = uuid4()

    if isinstance(log_store, ArtifactLogStore):
        artifact_store = client.active_stack.artifact_store
        return LogsRequest(
            id=log_id,
            project=client.active_project.id,
            source=source,
            uri=prepare_logs_uri(
                artifact_store=artifact_store,
                log_id=log_id,
            ),
            artifact_store_id=artifact_store.id,
        )
    else:
        return LogsRequest(
            id=log_id,
            project=client.active_project.id,
            source=source,
            log_store_id=log_store.id if log_store else None,
        )


def is_pipeline_logging_enabled(
    pipeline_configuration: PipelineConfiguration,
) -> bool:
    """Check if logging is enabled for a pipeline configuration.

    Args:
        pipeline_configuration: The pipeline configuration.

    Returns:
        True if logging is enabled, False if disabled.
    """
    if handle_bool_env_var(ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE, False):
        return False
    elif pipeline_configuration.enable_pipeline_logs is not None:
        return pipeline_configuration.enable_pipeline_logs
    else:
        return True


def is_step_logging_enabled(
    step_configuration: StepConfiguration,
    pipeline_configuration: PipelineConfiguration,
) -> bool:
    """Check if logging is enabled for a step configuration.

    Args:
        step_configuration: The step configuration.
        pipeline_configuration: The pipeline configuration.

    Returns:
        True if logging is enabled, False if disabled.
    """
    from zenml.orchestrators.utils import is_setting_enabled

    if handle_bool_env_var(ENV_ZENML_DISABLE_STEP_LOGS_STORAGE, False):
        return False
    else:
        is_enabled_on_step = step_configuration.enable_step_logs
        is_enabled_on_pipeline = pipeline_configuration.enable_step_logs
        return is_setting_enabled(is_enabled_on_step, is_enabled_on_pipeline)


def search_logs_by_source(
    logs_collection: List[LogsResponse], source: str
) -> Optional[LogsResponse]:
    """Get the logs response for a given source.

    Args:
        logs_collection: The logs collection.
        source: The source of the logs.

    Returns:
        The logs response for the given source.
    """
    for log in logs_collection:
        if log.source == source:
            return log
    return None


def get_run_log_metadata(
    pipeline_run: "PipelineRunResponse",
) -> Dict[str, Any]:
    """Get the log metadata for a pipeline run.

    Args:
        pipeline_run: The pipeline run.

    Returns:
        The log metadata.
    """
    log_metadata = {
        "zenml.pipeline.run.id": str(pipeline_run.id),
        "zenml.pipeline.run.name": pipeline_run.name,
        "zenml.project.id": str(pipeline_run.project.id),
        "zenml.project.name": pipeline_run.project.name,
    }

    client = Client()
    server_info = client.zen_store.get_store_info()

    if server_info.is_pro_server():
        if server_info.pro_workspace_id:
            log_metadata.update(
                {
                    "zenml.workspace.id": str(server_info.pro_workspace_id),
                }
            )
        if server_info.pro_workspace_name:
            log_metadata.update(
                {
                    "zenml.workspace.name": server_info.pro_workspace_name,
                }
            )

    if pipeline_run.pipeline is not None:
        log_metadata.update(
            {
                "zenml.pipeline.id": str(pipeline_run.pipeline.id),
                "zenml.pipeline.name": pipeline_run.pipeline.name,
            }
        )

    if pipeline_run.stack is not None:
        log_metadata.update(
            {
                "zenml.stack.id": str(pipeline_run.stack.id),
                "zenml.stack.name": pipeline_run.stack.name,
            }
        )

    if pipeline_run.user is not None:
        log_metadata.update(
            {
                "zenml.user.id": str(pipeline_run.user.id),
                "zenml.user.name": pipeline_run.user.name,
            }
        )

    return log_metadata


def get_step_log_metadata(step_run: "StepRunResponse") -> Dict[str, Any]:
    """Get the log metadata for a step run.

    Args:
        step_run: The step run.

    Returns:
        The log metadata.
    """
    return {
        "zenml.step.run.id": str(step_run.id),
        "zenml.step.run.name": step_run.name,
    }


def fetch_logs(
    logs: "LogsResponse",
    zen_store: "BaseZenStore",
    limit: int,
) -> List["LogEntry"]:
    """Fetch logs from the log store.

    This function is designed to be called from the server side where we can't
    always instantiate the full Stack object due to missing integration dependencies.
    Instead, it directly instantiates the appropriate log store based on the logs model.

    Args:
        logs: The logs response model containing metadata about the logs.
        zen_store: The zen store instance.
        limit: Maximum number of log entries to return.

    Returns:
        List of log entries.

    Raises:
        DoesNotExistException: If the log store doesn't exist or is not the right type.
        NotImplementedError: If the log store's dependencies are not installed.
        RuntimeError: If the function is called from the client environment.
    """
    from zenml.log_stores.base_log_store import BaseLogStore
    from zenml.stack import StackComponent
    from zenml.zen_server.rbac.endpoint_utils import (
        verify_permissions_and_get_entity,
    )

    if ENV_ZENML_SERVER not in os.environ:
        # This utility function should not be called from the client environment
        # because it would cause instantiating the active log store again.
        raise RuntimeError(
            "This utility function is only supported in the server "
            "environment. Use the log store directly instead."
        )

    log_store: Optional[BaseLogStore] = None

    if logs.log_store_id:
        try:
            log_store_model = verify_permissions_and_get_entity(
                id=logs.log_store_id,
                get_method=zen_store.get_stack_component,
            )
        except KeyError:
            raise DoesNotExistException(
                f"Log store '{logs.log_store_id}' does not exist."
            )

        if not log_store_model.type == StackComponentType.LOG_STORE:
            raise DoesNotExistException(
                f"Stack component '{logs.log_store_id}' is not a log store."
            )

        try:
            log_store = cast(
                BaseLogStore,
                StackComponent.from_model(log_store_model),
            )
        except ImportError:
            raise NotImplementedError(
                f"Log store '{log_store_model.name}' could not be "
                "instantiated."
            )
    elif logs.artifact_store_id:
        from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
        from zenml.log_stores.artifact.artifact_log_store import (
            ArtifactLogStore,
        )

        try:
            artifact_store_model = verify_permissions_and_get_entity(
                id=logs.artifact_store_id,
                get_method=zen_store.get_stack_component,
            )
        except KeyError:
            raise DoesNotExistException(
                f"Artifact store '{logs.artifact_store_id}' does not exist."
            )
        if not artifact_store_model.type == StackComponentType.ARTIFACT_STORE:
            raise DoesNotExistException(
                f"Stack component '{logs.artifact_store_id}' is not an artifact store."
            )
        artifact_store = cast(
            "BaseArtifactStore",
            StackComponent.from_model(artifact_store_model),
        )
        log_store = ArtifactLogStore.from_artifact_store(
            artifact_store=artifact_store
        )

    else:
        return []

    try:
        return log_store.fetch(logs_model=logs, limit=limit)
    finally:
        log_store.cleanup()


def setup_logging_context(
    source: str,
    pipeline_run: Optional["PipelineRunResponse"] = None,
    step_run: Optional["StepRunResponse"] = None,
    block_on_exit: bool = True,
) -> LoggingContext:
    """Setup a logging context for a given source.

    Args:
        source: The source of the logs.
        pipeline_run: The pipeline run.
        step_run: The step run.
        block_on_exit: Whether to block until all logs are flushed when the
            context is exited, if there are no more logging contexts active.

    Returns:
        A LoggingContext object.
    """
    log_metadata = {}

    logs_response: Optional[LogsResponse] = None

    if step_run and step_run.log_collection:
        logs_response = search_logs_by_source(
            logs_collection=step_run.log_collection,
            source=source,
        )
    elif pipeline_run and pipeline_run.log_collection:
        logs_response = search_logs_by_source(
            logs_collection=pipeline_run.log_collection,
            source=source,
        )

    logs_request = generate_logs_request(source=source)

    if step_run:
        log_metadata.update(get_step_log_metadata(step_run=step_run))
        logs_request.step_run_id = step_run.id

    if pipeline_run:
        log_metadata.update(get_run_log_metadata(pipeline_run=pipeline_run))
        if step_run is None:
            logs_request.pipeline_run_id = pipeline_run.id

    if logs_response is None:
        logs_response = Client().zen_store.create_logs(logs_request)

    return LoggingContext(
        name=str(logs_response.id),
        log_model=logs_response,
        block_on_exit=block_on_exit,
        **log_metadata,
    )
