#  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.
"""SQLModel implementation of artifact table."""

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

from pydantic import ValidationError
from sqlalchemy import TEXT, Column, UniqueConstraint
from sqlalchemy.orm import joinedload, object_session
from sqlalchemy.sql.base import ExecutableOption
from sqlmodel import Field, Relationship, asc, col, desc, select

from zenml.config.source import Source
from zenml.enums import (
    ArtifactSaveType,
    ArtifactType,
    ExecutionStatus,
    MetadataResourceTypes,
    TaggableResourceTypes,
)
from zenml.models import (
    ArtifactResponse,
    ArtifactResponseBody,
    ArtifactResponseMetadata,
    ArtifactResponseResources,
    ArtifactUpdate,
    ArtifactVersionRequest,
    ArtifactVersionResponse,
    ArtifactVersionResponseBody,
    ArtifactVersionResponseMetadata,
    ArtifactVersionResponseResources,
    ArtifactVersionUpdate,
)
from zenml.models.v2.core.artifact import ArtifactRequest
from zenml.utils.time_utils import utc_now
from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
from zenml.zen_stores.schemas.project_schemas import ProjectSchema
from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
from zenml.zen_stores.schemas.user_schemas import UserSchema
from zenml.zen_stores.schemas.utils import (
    RunMetadataInterface,
    jl_arg,
)

if TYPE_CHECKING:
    from zenml.zen_stores.schemas.artifact_visualization_schemas import (
        ArtifactVisualizationSchema,
    )
    from zenml.zen_stores.schemas.model_schemas import (
        ModelVersionArtifactSchema,
    )
    from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
    from zenml.zen_stores.schemas.step_run_schemas import (
        StepRunInputArtifactSchema,
        StepRunOutputArtifactSchema,
    )
    from zenml.zen_stores.schemas.tag_schemas import TagSchema


class ArtifactSchema(NamedSchema, table=True):
    """SQL Model for artifacts."""

    __tablename__ = "artifact"
    __table_args__ = (
        UniqueConstraint(
            "name",
            "project_id",
            name="unique_artifact_name_in_project",
        ),
    )

    # Fields
    has_custom_name: bool
    versions: List["ArtifactVersionSchema"] = Relationship(
        back_populates="artifact",
        sa_relationship_kwargs={"cascade": "delete"},
    )
    tags: List["TagSchema"] = Relationship(
        sa_relationship_kwargs=dict(
            primaryjoin=f"and_(foreign(TagResourceSchema.resource_type)=='{TaggableResourceTypes.ARTIFACT.value}', foreign(TagResourceSchema.resource_id)==ArtifactSchema.id)",
            secondary="tag_resource",
            secondaryjoin="TagSchema.id == foreign(TagResourceSchema.tag_id)",
            order_by="TagSchema.name",
            overlaps="tags",
        ),
    )

    project_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=ProjectSchema.__tablename__,
        source_column="project_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=False,
    )
    project: "ProjectSchema" = Relationship()

    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,
    )
    user: Optional["UserSchema"] = Relationship(
        back_populates="artifacts",
    )

    @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_resources:
            options.extend(
                [
                    joinedload(jl_arg(ArtifactSchema.user)),
                    # joinedload(jl_arg(ArtifactSchema.tags)),
                ]
            )

        return options

    @property
    def latest_version(self) -> Optional["ArtifactVersionSchema"]:
        """Fetch the latest version for this artifact.

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

        Returns:
            The latest version for this artifact.
        """
        if session := object_session(self):
            return (
                session.execute(
                    select(ArtifactVersionSchema)
                    .where(ArtifactVersionSchema.artifact_id == self.id)
                    .order_by(desc(ArtifactVersionSchema.created))
                    .limit(1)
                )
                .scalars()
                .one_or_none()
            )
        else:
            raise RuntimeError(
                "Missing DB session to fetch latest version for artifact."
            )

    @classmethod
    def from_request(
        cls,
        artifact_request: ArtifactRequest,
    ) -> "ArtifactSchema":
        """Convert an `ArtifactRequest` to an `ArtifactSchema`.

        Args:
            artifact_request: The request model to convert.

        Returns:
            The converted schema.
        """
        return cls(
            name=artifact_request.name,
            has_custom_name=artifact_request.has_custom_name,
            project_id=artifact_request.project,
            user_id=artifact_request.user,
        )

    def to_model(
        self,
        include_metadata: bool = False,
        include_resources: bool = False,
        **kwargs: Any,
    ) -> ArtifactResponse:
        """Convert an `ArtifactSchema` to an `ArtifactResponse`.

        Args:
            include_metadata: Whether the metadata will be filled.
            include_resources: Whether the resources will be filled.
            **kwargs: Keyword arguments to allow schema specific logic



        Returns:
            The created `ArtifactResponse`.
        """
        # Create the body of the model
        body = ArtifactResponseBody(
            user_id=self.user_id,
            project_id=self.project_id,
            created=self.created,
            updated=self.updated,
        )

        # Create the metadata of the model
        metadata = None
        if include_metadata:
            metadata = ArtifactResponseMetadata(
                has_custom_name=self.has_custom_name,
            )

        resources = None
        if include_resources:
            latest_id, latest_name = None, None
            if latest_version := self.latest_version:
                latest_id = latest_version.id
                latest_name = latest_version.version

            resources = ArtifactResponseResources(
                user=self.user.to_model() if self.user else None,
                tags=[tag.to_model() for tag in self.tags],
                latest_version_id=latest_id,
                latest_version_name=latest_name,
            )

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

    def update(self, artifact_update: ArtifactUpdate) -> "ArtifactSchema":
        """Update an `ArtifactSchema` with an `ArtifactUpdate`.

        Args:
            artifact_update: The update model to apply.

        Returns:
            The updated `ArtifactSchema`.
        """
        self.updated = utc_now()
        if artifact_update.name:
            self.name = artifact_update.name
            self.has_custom_name = True
        if artifact_update.has_custom_name is not None:
            self.has_custom_name = artifact_update.has_custom_name
        return self


class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
    """SQL Model for artifact versions."""

    __tablename__ = "artifact_version"
    __table_args__ = (
        UniqueConstraint(
            "version",
            "artifact_id",
            name="unique_version_for_artifact_id",
        ),
    )

    # Fields
    version: str
    version_number: Optional[int]
    type: str
    uri: str = Field(sa_column=Column(TEXT, nullable=False))
    materializer: str = Field(sa_column=Column(TEXT, nullable=False))
    data_type: str = Field(sa_column=Column(TEXT, nullable=False))
    tags: List["TagSchema"] = Relationship(
        sa_relationship_kwargs=dict(
            primaryjoin=f"and_(foreign(TagResourceSchema.resource_type)=='{TaggableResourceTypes.ARTIFACT_VERSION.value}', foreign(TagResourceSchema.resource_id)==ArtifactVersionSchema.id)",
            secondary="tag_resource",
            secondaryjoin="TagSchema.id == foreign(TagResourceSchema.tag_id)",
            order_by="TagSchema.name",
            overlaps="tags",
        ),
    )
    save_type: str = Field(sa_column=Column(TEXT, nullable=False))
    content_hash: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
    item_count: Optional[int] = None

    # Foreign keys
    artifact_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=ArtifactSchema.__tablename__,
        source_column="artifact_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=False,
    )
    artifact_store_id: Optional[UUID] = build_foreign_key_field(
        source=__tablename__,
        target=StackComponentSchema.__tablename__,
        source_column="artifact_store_id",
        target_column="id",
        ondelete="SET NULL",
        nullable=True,
    )
    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,
    )

    # Relationships
    artifact: "ArtifactSchema" = Relationship(back_populates="versions")
    user: Optional["UserSchema"] = Relationship(
        back_populates="artifact_versions"
    )
    project: "ProjectSchema" = Relationship(back_populates="artifact_versions")
    run_metadata: List["RunMetadataSchema"] = Relationship(
        sa_relationship_kwargs=dict(
            secondary="run_metadata_resource",
            primaryjoin=f"and_(foreign(RunMetadataResourceSchema.resource_type)=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
            secondaryjoin="RunMetadataSchema.id==foreign(RunMetadataResourceSchema.run_metadata_id)",
            overlaps="run_metadata",
        ),
    )
    visualizations: List["ArtifactVisualizationSchema"] = Relationship(
        back_populates="artifact_version",
        sa_relationship_kwargs={"cascade": "delete"},
    )

    # Needed for cascade deletion behavior
    model_versions_artifacts_links: List["ModelVersionArtifactSchema"] = (
        Relationship(
            back_populates="artifact_version",
            sa_relationship_kwargs={"cascade": "delete"},
        )
    )
    output_of_step_runs: List["StepRunOutputArtifactSchema"] = Relationship(
        back_populates="artifact_version",
        sa_relationship_kwargs={"cascade": "delete"},
    )
    input_of_step_runs: List["StepRunInputArtifactSchema"] = Relationship(
        back_populates="artifact_version",
        sa_relationship_kwargs={"cascade": "delete"},
    )

    @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(ArtifactVersionSchema.visualizations)),
        #             joinedload(jl_arg(ArtifactVersionSchema.run_metadata)),
        #         ]
        #     )

        if include_resources:
            options.extend(
                [
                    joinedload(jl_arg(ArtifactVersionSchema.user)),
                    # joinedload(jl_arg(ArtifactVersionSchema.tags)),
                ]
            )

        return options

    @property
    def producer_run_ids(self) -> Optional[Tuple[UUID, UUID]]:
        """Fetch the producer run IDs for this artifact version.

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

        Returns:
            The producer step run ID and pipeline run ID for this artifact
            version.
        """
        from zenml.zen_stores.schemas import (
            StepRunOutputArtifactSchema,
            StepRunSchema,
        )

        if session := object_session(self):
            row = session.execute(
                select(StepRunSchema.id, StepRunSchema.pipeline_run_id)
                .join(
                    StepRunOutputArtifactSchema,
                    col(StepRunOutputArtifactSchema.step_id)
                    == col(StepRunSchema.id),
                )
                .where(col(StepRunOutputArtifactSchema.artifact_id) == self.id)
                .where(
                    col(StepRunSchema.status)
                    == ExecutionStatus.COMPLETED.value
                )
                # Fetch the oldest step run
                .order_by(asc(StepRunSchema.created))
                .limit(1)
            ).one_or_none()

            return (row[0], row[1]) if row else None
        else:
            raise RuntimeError(
                "Missing DB session to fetch producer run for artifact version."
            )

    @classmethod
    def from_request(
        cls,
        artifact_version_request: ArtifactVersionRequest,
    ) -> "ArtifactVersionSchema":
        """Convert an `ArtifactVersionRequest` to an `ArtifactVersionSchema`.

        Args:
            artifact_version_request: The request model to convert.

        Raises:
            ValueError: If the request does not specify a version number.

        Returns:
            The converted schema.
        """
        if not artifact_version_request.version:
            raise ValueError("Missing version for artifact version request.")

        try:
            version_number = int(artifact_version_request.version)
        except ValueError:
            version_number = None
        return cls(
            artifact_id=artifact_version_request.artifact_id,
            version=str(artifact_version_request.version),
            version_number=version_number,
            artifact_store_id=artifact_version_request.artifact_store_id,
            project_id=artifact_version_request.project,
            user_id=artifact_version_request.user,
            type=artifact_version_request.type.value,
            uri=artifact_version_request.uri,
            materializer=artifact_version_request.materializer.model_dump_json(),
            data_type=artifact_version_request.data_type.model_dump_json(),
            save_type=artifact_version_request.save_type.value,
            content_hash=artifact_version_request.content_hash,
            item_count=artifact_version_request.item_count,
        )

    def to_model(
        self,
        include_metadata: bool = False,
        include_resources: bool = False,
        **kwargs: Any,
    ) -> ArtifactVersionResponse:
        """Convert an `ArtifactVersionSchema` to an `ArtifactVersionResponse`.

        Args:
            include_metadata: Whether the metadata will be filled.
            include_resources: Whether the resources will be filled.
            **kwargs: Keyword arguments to allow schema specific logic



        Returns:
            The created `ArtifactVersionResponse`.
        """
        try:
            materializer = Source.model_validate_json(self.materializer)
        except ValidationError:
            # This is an old source which was an importable source path
            materializer = Source.from_import_path(self.materializer)

        try:
            data_type = Source.model_validate_json(self.data_type)
        except ValidationError:
            # This is an old source which was an importable source path
            data_type = Source.from_import_path(self.data_type)

        # Create the body of the model
        artifact = self.artifact.to_model()
        body = ArtifactVersionResponseBody(
            user_id=self.user_id,
            project_id=self.project_id,
            artifact=artifact,
            version=self.version or str(self.version_number),
            uri=self.uri,
            type=ArtifactType(self.type),
            materializer=materializer,
            data_type=data_type,
            created=self.created,
            updated=self.updated,
            save_type=ArtifactSaveType(self.save_type),
            artifact_store_id=self.artifact_store_id,
            content_hash=self.content_hash,
            item_count=self.item_count,
        )

        # Create the metadata of the model
        metadata = None
        if include_metadata:
            metadata = ArtifactVersionResponseMetadata(
                visualizations=[v.to_model() for v in self.visualizations],
                run_metadata=self.fetch_metadata(),
            )

        resources = None
        if include_resources:
            producer_step_run_id, producer_pipeline_run_id = None, None
            if producer_run_ids := self.producer_run_ids:
                # TODO: Why was the producer_pipeline_run_id only set for one
                # of the cases before?
                producer_step_run_id, producer_pipeline_run_id = (
                    producer_run_ids
                )

            resources = ArtifactVersionResponseResources(
                user=self.user.to_model() if self.user else None,
                tags=[tag.to_model() for tag in self.tags],
                producer_step_run_id=producer_step_run_id,
                producer_pipeline_run_id=producer_pipeline_run_id,
            )

        return ArtifactVersionResponse(
            id=self.id,
            body=body,
            metadata=metadata,
            resources=resources,
        )

    def update(
        self, artifact_version_update: ArtifactVersionUpdate
    ) -> "ArtifactVersionSchema":
        """Update an `ArtifactVersionSchema` with an `ArtifactVersionUpdate`.

        Args:
            artifact_version_update: The update model to apply.

        Returns:
            The updated `ArtifactVersionSchema`.
        """
        self.updated = utc_now()
        return self
