#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   links.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK asset links module.
"""

import re

from vi.api.resources.datasets.utils.helper import build_dataset_id
from vi.client.http.links import LinkBuilder, ResourceLinkParser

ASSET_LINK_RE = re.compile(
    r"^/workspaces/(?P<organization_id>[^/]+)/datasets/(?P=organization_id)_"
    r"(?P<dataset_id>[^/]+)/assets(?:/(?P<asset_id>[^/]+))?$"
)

ASSET_INGESTION_SESSION_LINK_RE = re.compile(
    r"^/workspaces/(?P<organization_id>[^/]+)/datasets/(?P=organization_id)_"
    r"(?P<dataset_id>[^/]+)/assetSources"
    r"/(?P<asset_source_class>[^/]+)/(?P<asset_source_provider>[^/]+)/"
    r"(?P<asset_source_id>[^/]+)/assetIngestionSessions(?:/(?P<asset_ingestion_session_id>[^/]+))?$"
)

MULTIPART_UPLOAD_LINK_RE = re.compile(
    r"^/workspaces/(?P<organization_id>[^/]+)/datasets/(?P=organization_id)_"
    r"(?P<dataset_id>[^/]+)/assetSources"
    r"/(?P<asset_source_class>[^/]+)/(?P<asset_source_provider>[^/]+)/"
    r"(?P<asset_source_id>[^/]+)/multipartUploads(?:/(?P<multipart_upload_id>[^/]+))?$"
)


class AssetLinkParser(ResourceLinkParser):
    """Parse an asset link."""

    _organization_id: str
    _dataset_id: str

    def __init__(self, organization_id: str, dataset_id: str):
        """Initialize the asset link parser.

        Args:
            organization_id: Organization ID for constructing asset links.
            dataset_id: Dataset ID for constructing asset links.

        """
        self._organization_id = organization_id
        self._dataset_id = dataset_id
        self._base_link = (
            LinkBuilder("/workspaces")
            .add_segment(self._organization_id)
            .add_segment("datasets")
            .add_segment(build_dataset_id(self._organization_id, self._dataset_id))
            .add_segment("assets")
            .build()
        )

    def __call__(self, asset_id_or_link: str = "") -> str:
        """Parse the asset link.

        Args:
            asset_id_or_link: Asset ID or full link path.

        Returns:
            The parsed asset link path.

        """
        if ASSET_LINK_RE.match(asset_id_or_link):
            return asset_id_or_link

        return LinkBuilder(self._base_link).add_segment(asset_id_or_link).build()


class AssetSourceBaseLinkParser(ResourceLinkParser):
    """Base parser for asset source-related links.

    This base class handles the common pattern of building links for
    asset source operations (ingestion sessions, multipart uploads, etc.).
    """

    _organization_id: str
    _dataset_id: str
    _asset_source_class: str
    _asset_source_provider: str
    _asset_source_id: str

    def __init__(
        self,
        organization_id: str,
        dataset_id: str,
        asset_source_class: str,
        asset_source_provider: str,
        asset_source_id: str,
    ):
        """Initialize the asset source base link parser.

        Args:
            organization_id: Organization ID.
            dataset_id: Dataset ID.
            asset_source_class: Asset source class identifier.
            asset_source_provider: Asset source provider identifier.
            asset_source_id: Asset source ID.

        """
        self._organization_id = organization_id
        self._dataset_id = dataset_id
        self._asset_source_class = asset_source_class
        self._asset_source_provider = asset_source_provider
        self._asset_source_id = asset_source_id

    def _build_base_link(self, resource_segment: str) -> str:
        """Build the base link with common asset source path.

        Args:
            resource_segment: The final resource segment (e.g., "assetIngestionSessions").

        Returns:
            The constructed base link path.

        """
        return (
            LinkBuilder("/workspaces")
            .add_segment(self._organization_id)
            .add_segment("datasets")
            .add_segment(build_dataset_id(self._organization_id, self._dataset_id))
            .add_segment("assetSources")
            .add_segment(self._asset_source_class)
            .add_segment(self._asset_source_provider)
            .add_segment(self._asset_source_id)
            .add_segment(resource_segment)
            .build()
        )

    def _parse_link(
        self,
        id_or_link: str,
        link_regex: re.Pattern[str],
        base_link: str,
    ) -> str:
        """Parse a link, returning it as-is if it matches the regex pattern.

        Args:
            id_or_link: Resource ID or full link path.
            link_regex: Regex pattern to match valid links.
            base_link: Base link to append the ID to if not a full link.

        Returns:
            The parsed link path.

        """
        if link_regex.match(id_or_link):
            return id_or_link

        return LinkBuilder(base_link).add_segment(id_or_link).build()


class AssetIngestionSessionLinkParser(AssetSourceBaseLinkParser):
    """Parse an asset ingestion session link."""

    def __init__(
        self,
        organization_id: str,
        dataset_id: str,
        asset_source_class: str,
        asset_source_provider: str,
        asset_source_id: str,
    ):
        """Initialize the asset ingestion session link parser.

        Args:
            organization_id: Organization ID.
            dataset_id: Dataset ID.
            asset_source_class: Asset source class identifier.
            asset_source_provider: Asset source provider identifier.
            asset_source_id: Asset source ID.

        """
        super().__init__(
            organization_id,
            dataset_id,
            asset_source_class,
            asset_source_provider,
            asset_source_id,
        )
        self._base_link = self._build_base_link("assetIngestionSessions")

    def __call__(self, asset_ingestion_session_id_or_link: str = "") -> str:
        """Parse the asset ingestion session link.

        Args:
            asset_ingestion_session_id_or_link: Session ID or full link path.

        Returns:
            The parsed asset ingestion session link path.

        """
        return self._parse_link(
            asset_ingestion_session_id_or_link,
            ASSET_INGESTION_SESSION_LINK_RE,
            self._base_link,
        )


class MultipartUploadLinkParser(AssetSourceBaseLinkParser):
    """Parse a multipart upload link."""

    def __init__(
        self,
        organization_id: str,
        dataset_id: str,
        asset_source_class: str,
        asset_source_provider: str,
        asset_source_id: str,
    ):
        """Initialize the multipart upload link parser.

        Args:
            organization_id: Organization ID.
            dataset_id: Dataset ID.
            asset_source_class: Asset source class identifier.
            asset_source_provider: Asset source provider identifier.
            asset_source_id: Asset source ID.

        """
        super().__init__(
            organization_id,
            dataset_id,
            asset_source_class,
            asset_source_provider,
            asset_source_id,
        )
        self._base_link = self._build_base_link("multipartUploads")

    def __call__(self, multipart_upload_id_or_link: str = "") -> str:
        """Parse the multipart upload link.

        Args:
            multipart_upload_id_or_link: Upload ID or full link path.

        Returns:
            The parsed multipart upload link path.

        """
        return self._parse_link(
            multipart_upload_id_or_link,
            MULTIPART_UPLOAD_LINK_RE,
            self._base_link,
        )
