#  Copyright (c) ZenML GmbH 2022. 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.
"""Pipeline snapshot schemas."""

import json
from typing import TYPE_CHECKING, Any, List, Optional, Sequence
from uuid import UUID

from sqlalchemy import TEXT, CheckConstraint, Column, String, UniqueConstraint
from sqlalchemy.dialects.mysql import MEDIUMTEXT
from sqlalchemy.orm import joinedload, object_session, selectinload
from sqlalchemy.sql.base import ExecutableOption
from sqlmodel import Field, Relationship, asc, col, desc, select

from zenml.config.pipeline_configurations import PipelineConfiguration
from zenml.config.pipeline_spec import PipelineSpec
from zenml.config.step_configurations import Step
from zenml.constants import MEDIUMTEXT_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
from zenml.enums import TaggableResourceTypes, VisualizationResourceTypes
from zenml.logger import get_logger
from zenml.models import (
    PipelineSnapshotRequest,
    PipelineSnapshotResponse,
    PipelineSnapshotResponseBody,
    PipelineSnapshotResponseMetadata,
    PipelineSnapshotResponseResources,
    PipelineSnapshotUpdate,
)
from zenml.utils.time_utils import utc_now
from zenml.zen_stores.schemas.base_schemas import BaseSchema
from zenml.zen_stores.schemas.code_repository_schemas import (
    CodeReferenceSchema,
)
from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema
from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
from zenml.zen_stores.schemas.project_schemas import ProjectSchema
from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
from zenml.zen_stores.schemas.schema_utils import (
    build_foreign_key_field,
    build_index,
)
from zenml.zen_stores.schemas.stack_schemas import StackSchema
from zenml.zen_stores.schemas.tag_schemas import TagSchema
from zenml.zen_stores.schemas.user_schemas import UserSchema
from zenml.zen_stores.schemas.utils import jl_arg

if TYPE_CHECKING:
    from zenml.zen_stores.schemas.curated_visualization_schemas import (
        CuratedVisualizationSchema,
    )
    from zenml.zen_stores.schemas.deployment_schemas import (
        DeploymentSchema,
    )
    from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
    from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema

logger = get_logger(__name__)


class PipelineSnapshotSchema(BaseSchema, table=True):
    """SQL Model for pipeline snapshots."""

    __tablename__ = "pipeline_snapshot"
    __table_args__ = (
        UniqueConstraint(
            "pipeline_id",
            "name",
            name="unique_name_for_pipeline_id",
        ),
        build_index(
            table_name=__tablename__,
            column_names=["source_snapshot_id"],
        ),
    )

    # Fields
    name: Optional[str] = Field(nullable=True)
    description: Optional[str] = Field(
        sa_column=Column(
            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
                MEDIUMTEXT, "mysql"
            ),
            nullable=True,
        )
    )
    is_dynamic: bool = Field(nullable=False, default=False)

    pipeline_configuration: str = Field(
        sa_column=Column(
            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
                MEDIUMTEXT, "mysql"
            ),
            nullable=False,
        )
    )
    client_environment: str = Field(sa_column=Column(TEXT, nullable=False))
    run_name_template: str = Field(nullable=False)
    client_version: str = Field(nullable=True)
    server_version: str = Field(nullable=True)
    pipeline_version_hash: Optional[str] = Field(nullable=True, default=None)
    pipeline_spec: Optional[str] = Field(
        sa_column=Column(
            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
                MEDIUMTEXT, "mysql"
            ),
            nullable=True,
        )
    )
    source_code: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
    code_path: Optional[str] = Field(nullable=True)

    # Foreign keys
    user_id: Optional[UUID] = build_foreign_key_field(
        source=__tablename__,
        target=UserSchema.__tablename__,
        source_column="user_id",
        target_column="id",
        ondelete="SET NULL",
        nullable=True,
    )
    project_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=ProjectSchema.__tablename__,
        source_column="project_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=False,
    )
    stack_id: Optional[UUID] = build_foreign_key_field(
        source=__tablename__,
        target=StackSchema.__tablename__,
        source_column="stack_id",
        target_column="id",
        ondelete="SET NULL",
        nullable=True,
    )
    pipeline_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=PipelineSchema.__tablename__,
        source_column="pipeline_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=False,
    )
    schedule_id: Optional[UUID] = build_foreign_key_field(
        source=__tablename__,
        target=ScheduleSchema.__tablename__,
        source_column="schedule_id",
        target_column="id",
        ondelete="SET NULL",
        nullable=True,
    )
    build_id: Optional[UUID] = build_foreign_key_field(
        source=__tablename__,
        target=PipelineBuildSchema.__tablename__,
        source_column="build_id",
        target_column="id",
        ondelete="SET NULL",
        nullable=True,
    )
    code_reference_id: Optional[UUID] = build_foreign_key_field(
        source=__tablename__,
        target=CodeReferenceSchema.__tablename__,
        source_column="code_reference_id",
        target_column="id",
        ondelete="SET NULL",
        nullable=True,
    )
    # This is not a foreign key to remove a cycle which messes with our DB
    # backup process
    source_snapshot_id: Optional[UUID] = None

    # Deprecated, remove once we remove run templates entirely
    template_id: Optional[UUID] = None

    # SQLModel Relationships
    source_snapshot: Optional["PipelineSnapshotSchema"] = Relationship(
        sa_relationship_kwargs=dict(
            primaryjoin="PipelineSnapshotSchema.source_snapshot_id == foreign(PipelineSnapshotSchema.id)",
            viewonly=True,
        ),
    )
    user: Optional["UserSchema"] = Relationship(
        back_populates="snapshots",
    )
    project: "ProjectSchema" = Relationship(back_populates="snapshots")
    stack: Optional["StackSchema"] = Relationship()
    pipeline: "PipelineSchema" = Relationship()
    schedule: Optional["ScheduleSchema"] = Relationship()
    build: Optional["PipelineBuildSchema"] = Relationship(
        sa_relationship_kwargs={
            "foreign_keys": "[PipelineSnapshotSchema.build_id]"
        }
    )
    code_reference: Optional["CodeReferenceSchema"] = Relationship()

    pipeline_runs: List["PipelineRunSchema"] = Relationship(
        sa_relationship_kwargs={"cascade": "delete"}
    )
    step_runs: List["StepRunSchema"] = Relationship(
        sa_relationship_kwargs={"cascade": "delete"}
    )
    step_configurations: List["StepConfigurationSchema"] = Relationship(
        sa_relationship_kwargs={
            "cascade": "delete",
            "order_by": "asc(StepConfigurationSchema.index)",
        }
    )
    deployment: Optional["DeploymentSchema"] = Relationship(
        back_populates="snapshot"
    )
    step_count: int
    tags: List["TagSchema"] = Relationship(
        sa_relationship_kwargs=dict(
            primaryjoin=f"and_(foreign(TagResourceSchema.resource_type)=='{TaggableResourceTypes.PIPELINE_SNAPSHOT.value}', foreign(TagResourceSchema.resource_id)==PipelineSnapshotSchema.id)",
            secondary="tag_resource",
            secondaryjoin="TagSchema.id == foreign(TagResourceSchema.tag_id)",
            order_by="TagSchema.name",
            overlaps="tags",
        ),
    )
    visualizations: List["CuratedVisualizationSchema"] = Relationship(
        sa_relationship_kwargs=dict(
            primaryjoin=(
                "and_(CuratedVisualizationSchema.resource_type"
                f"=='{VisualizationResourceTypes.PIPELINE_SNAPSHOT.value}', "
                "foreign(CuratedVisualizationSchema.resource_id)==PipelineSnapshotSchema.id)"
            ),
            overlaps="visualizations",
            cascade="delete",
            order_by="CuratedVisualizationSchema.display_order",
        ),
    )

    @property
    def latest_run(self) -> Optional["PipelineRunSchema"]:
        """Fetch the latest run for this snapshot.

        Raises:
            RuntimeError: If no session for the schema exists.

        Returns:
            The latest run for this snapshot.
        """
        from zenml.zen_stores.schemas import (
            PipelineRunSchema,
            PipelineSnapshotSchema,
        )

        if session := object_session(self):
            return (
                session.execute(
                    select(PipelineRunSchema)
                    .join(
                        PipelineSnapshotSchema,
                        col(PipelineSnapshotSchema.id)
                        == col(PipelineRunSchema.snapshot_id),
                    )
                    .where(
                        # The snapshot for this run used this snapshot as a
                        # source (e.g. run triggered from the server,
                        # invocation of a deployment). We currently do not
                        # include runs created directly from a snapshot (e.g.
                        # run directly, scheduled runs), as these happen before
                        # the user officially creates (= assigns a name to) the
                        # snapshot.
                        col(PipelineSnapshotSchema.source_snapshot_id)
                        == self.id,
                    )
                    .order_by(desc(PipelineRunSchema.created))
                    .limit(1)
                )
                .scalars()
                .one_or_none()
            )
        else:
            raise RuntimeError(
                "Missing DB session to fetch latest run for snapshot."
            )

    def get_step_configurations(
        self, include: Optional[List[str]] = None
    ) -> List["StepConfigurationSchema"]:
        """Get step configurations for the snapshot.

        Args:
            include: List of step names to include. If not given, all step
                configurations will be included.

        Raises:
            RuntimeError: If no session for the schema exists.

        Returns:
            List of step configurations.
        """
        if session := object_session(self):
            query = (
                select(StepConfigurationSchema)
                .where(StepConfigurationSchema.snapshot_id == self.id)
                .order_by(asc(StepConfigurationSchema.index))
            )

            if include:
                query = query.where(
                    col(StepConfigurationSchema.name).in_(include)
                )

            return list(session.execute(query).scalars().all())
        else:
            raise RuntimeError(
                "Missing DB session to fetch step configurations."
            )

    def get_step_configuration(
        self, step_name: str
    ) -> "StepConfigurationSchema":
        """Get a step configuration of the snapshot.

        Args:
            step_name: The name of the step to get the configuration for.

        Raises:
            KeyError: If the step configuration is not found.

        Returns:
            The step configuration.
        """
        step_configs = self.get_step_configurations(include=[step_name])
        if len(step_configs) == 0:
            raise KeyError(
                f"Step configuration for step `{step_name}` not found."
            )
        return step_configs[0]

    @classmethod
    def get_query_options(
        cls,
        include_metadata: bool = False,
        include_resources: bool = False,
        **kwargs: Any,
    ) -> Sequence[ExecutableOption]:
        """Get the query options for the schema.

        Args:
            include_metadata: Whether metadata will be included when converting
                the schema to a model.
            include_resources: Whether resources will be included when
                converting the schema to a model.
            **kwargs: Keyword arguments to allow schema specific logic

        Returns:
            A list of query options.
        """
        options = []

        if include_metadata:
            options.extend(
                [
                    joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                    joinedload(jl_arg(PipelineSnapshotSchema.build)),
                    joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                    joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
                    joinedload(jl_arg(PipelineSnapshotSchema.code_reference)),
                ]
            )

        if include_resources:
            options.extend(
                [
                    joinedload(jl_arg(PipelineSnapshotSchema.user)),
                    selectinload(
                        jl_arg(PipelineSnapshotSchema.visualizations)
                    ),
                ]
            )

        return options

    @classmethod
    def from_request(
        cls,
        request: PipelineSnapshotRequest,
        code_reference_id: Optional[UUID],
    ) -> "PipelineSnapshotSchema":
        """Create schema from request.

        Args:
            request: The request to convert.
            code_reference_id: Optional ID of the code reference for the
                snapshot.

        Returns:
            The created schema.
        """
        client_env = json.dumps(request.client_environment)
        if len(client_env) > TEXT_FIELD_MAX_LENGTH:
            logger.warning(
                "Client environment is too large to be stored in the database. "
                "Skipping."
            )
            client_env = "{}"

        name = None
        if isinstance(request.name, str):
            name = request.name

        return cls(
            name=name,
            description=request.description,
            source_code=request.source_code,
            is_dynamic=request.is_dynamic,
            stack_id=request.stack,
            project_id=request.project,
            pipeline_id=request.pipeline,
            build_id=request.build,
            user_id=request.user,
            schedule_id=request.schedule,
            template_id=request.template,
            source_snapshot_id=request.source_snapshot,
            code_reference_id=code_reference_id,
            run_name_template=request.run_name_template,
            pipeline_configuration=request.pipeline_configuration.model_dump_json(),
            step_count=len(request.step_configurations),
            client_environment=client_env,
            client_version=request.client_version,
            server_version=request.server_version,
            pipeline_version_hash=request.pipeline_version_hash,
            pipeline_spec=json.dumps(
                request.pipeline_spec.model_dump(mode="json"), sort_keys=True
            )
            if request.pipeline_spec
            else None,
            code_path=request.code_path,
        )

    def update(
        self, update: PipelineSnapshotUpdate
    ) -> "PipelineSnapshotSchema":
        """Update the schema.

        Args:
            update: The update to apply.

        Returns:
            The updated schema.
        """
        if isinstance(update.name, str):
            self.name = update.name
        elif update.name is False:
            self.name = None

        if update.description:
            self.description = update.description

        self.updated = utc_now()
        return self

    @property
    def is_runnable(self) -> bool:
        """Implements the `is_runnable` property.

        Returns:
            True if the snapshot is runnable from server.
        """
        return (
            self.build is not None
            and not self.build.is_local
            and self.build.stack_id is not None
        )

    def to_model(
        self,
        include_metadata: bool = False,
        include_resources: bool = False,
        include_python_packages: bool = False,
        include_config_schema: Optional[bool] = None,
        step_configuration_filter: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> PipelineSnapshotResponse:
        """Convert schema to response.

        Args:
            include_metadata: Whether the metadata will be filled.
            include_resources: Whether the resources will be filled.
            include_python_packages: Whether the python packages will be filled.
            include_config_schema: Whether the config schema will be filled.
            step_configuration_filter: List of step configurations to include in
                the response. If not given, all step configurations will be
                included.
            **kwargs: Keyword arguments to allow schema specific logic

        Returns:
            The response.
        """
        deployable = False
        if self.build and self.stack and self.stack.has_deployer:
            deployable = True

        body = PipelineSnapshotResponseBody(
            user_id=self.user_id,
            project_id=self.project_id,
            created=self.created,
            updated=self.updated,
            runnable=self.is_runnable,
            deployable=deployable,
            is_dynamic=self.is_dynamic,
        )
        metadata = None
        if include_metadata:
            pipeline_configuration = PipelineConfiguration.model_validate_json(
                self.pipeline_configuration
            )
            step_configurations = {}
            for step_configuration in self.get_step_configurations(
                include=step_configuration_filter
            ):
                step_configurations[step_configuration.name] = Step.from_dict(
                    json.loads(step_configuration.config),
                    pipeline_configuration,
                )

            client_environment = json.loads(self.client_environment)
            if not include_python_packages:
                client_environment.pop("python_packages", None)

            config_template = None
            config_schema = None

            if include_config_schema and self.build and self.build.stack_id:
                from zenml.zen_stores import template_utils

                if step_configuration_filter:
                    # If only a subset of step configurations is requested,
                    # we still need to get all of them to generate the config
                    # template and schema
                    all_step_configurations = {
                        step_configuration.name: Step.from_dict(
                            json.loads(step_configuration.config),
                            pipeline_configuration,
                        )
                        for step_configuration in self.get_step_configurations()
                    }
                else:
                    all_step_configurations = step_configurations

                config_template = template_utils.generate_config_template(
                    snapshot=self,
                    pipeline_configuration=pipeline_configuration,
                    step_configurations=all_step_configurations,
                )
                config_schema = template_utils.generate_config_schema(
                    snapshot=self,
                    pipeline_configuration=pipeline_configuration,
                    step_configurations=all_step_configurations,
                )

            metadata = PipelineSnapshotResponseMetadata(
                description=self.description,
                source_code=self.source_code,
                run_name_template=self.run_name_template,
                pipeline_configuration=pipeline_configuration,
                step_configurations=step_configurations,
                client_environment=client_environment,
                client_version=self.client_version,
                server_version=self.server_version,
                pipeline_version_hash=self.pipeline_version_hash,
                pipeline_spec=PipelineSpec.model_validate_json(
                    self.pipeline_spec
                )
                if self.pipeline_spec
                else None,
                code_path=self.code_path,
                template_id=self.template_id,
                source_snapshot_id=self.source_snapshot_id,
                config_schema=config_schema,
                config_template=config_template,
            )

        resources = None
        if include_resources:
            latest_run = self.latest_run
            latest_run_user = latest_run.user if latest_run else None

            resources = PipelineSnapshotResponseResources(
                user=self.user.to_model() if self.user else None,
                pipeline=self.pipeline.to_model(),
                stack=self.stack.to_model() if self.stack else None,
                build=self.build.to_model() if self.build else None,
                schedule=self.schedule.to_model() if self.schedule else None,
                code_reference=self.code_reference.to_model()
                if self.code_reference
                else None,
                deployment=self.deployment.to_model()
                if self.deployment
                else None,
                tags=[tag.to_model() for tag in self.tags],
                latest_run_id=latest_run.id if latest_run else None,
                latest_run_status=latest_run.status if latest_run else None,
                latest_run_user=latest_run_user.to_model()
                if latest_run_user
                else None,
                visualizations=[
                    visualization.to_model(
                        include_metadata=False,
                        include_resources=False,
                    )
                    for visualization in self.visualizations
                ],
            )

        return PipelineSnapshotResponse(
            id=self.id,
            name=self.name,
            body=body,
            metadata=metadata,
            resources=resources,
        )


class StepConfigurationSchema(BaseSchema, table=True):
    """SQL Model for step configurations."""

    __tablename__ = "step_configuration"
    __table_args__ = (
        UniqueConstraint(
            "snapshot_id",
            "step_run_id",
            "name",
            name="unique_step_configuration_for_snapshot_or_step_run",
        ),
        CheckConstraint(
            "(snapshot_id IS NULL AND step_run_id IS NOT NULL) OR "
            "(snapshot_id IS NOT NULL AND step_run_id IS NULL)",
            name="ck_step_configuration_snapshot_step_run_exclusivity",
        ),
    )

    index: int
    name: str
    config: str = Field(
        sa_column=Column(
            String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
                MEDIUMTEXT, "mysql"
            ),
            nullable=False,
        )
    )

    snapshot_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=PipelineSnapshotSchema.__tablename__,
        source_column="snapshot_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=True,
    )
    step_run_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target="step_run",
        source_column="step_run_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=True,
    )
