# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/classes/50_DomoLineage_v2.ipynb.

# %% auto 0
__all__ = ['DomoLineage_Link', 'DomoLineageLink_Dataflow', 'DomoLineageLink_Publication', 'DomoLineageLink_Card',
           'DomoLineageLink_Dataset', 'DomoLineageLinkTypeFactory_Enum', 'DomoLineage_ParentTypeEnum', 'DomoLineage',
           'DomoLineage_Publication']

# %% ../../nbs/classes/50_DomoLineage_v2.ipynb 2
from functools import partial
from dataclasses import dataclass, field
from enum import Enum
from abc import ABC, abstractmethod
from typing import List, Any

import httpx

import domolibrary.routes.datacenter as datacenter_routes

import domolibrary.utils.chunk_execution as dmce
import domolibrary.utils.convert as dmcv

import domolibrary.client.DomoAuth as dmda
import domolibrary.client.DomoError as dmde

# %% ../../nbs/classes/50_DomoLineage_v2.ipynb 5
@dataclass
class DomoLineage_Link(ABC):
    auth: dmda.DomoAuth = field(repr=False)
    type: str
    id: str

    entity: Any = field(repr = False) # DomoDataset, DomoDataflow, DomoPublication

    parents: List["DomoLineage_Link"] = field(default_factory=list)
    children: List["DomoLineage_Link"] = field(default_factory=list)

    @staticmethod
    @abstractmethod
    async def get_entity(entity_id, auth):
        """
        Get the entity associated with this lineage link.
        This method should be implemented by subclasses to return the appropriate entity.
        """
        raise NotImplementedError("Subclasses must implement this method.")
    
    @classmethod
    @abstractmethod
    async def get_from_json(cls, obj, auth):
        """
        Create a DomoLineage_Link instance from a JSON object.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    def _init_children(self):
        if not self.children:
            return []
        
        self.children = [DomoLineageLinkTypeFactory_Enum[c['type']].value(
                auth=self.auth,
                id=c['id'],
                type=c['type'],
                entity=None
            ) for c in self.children]
        return self.children
    
    def _init_parents(self):
        if not self.parents:
            return []
        
        self.parents = [DomoLineageLinkTypeFactory_Enum[p['type']].value(
                auth=self.auth,
                id=p['id'],
                type = p['type'],
                entity=None
            ) for p in self.parents]
        return self.parents


    def __post_init__(self):
        if self.children:
            self._init_children()
    
        if self.parents:
            self._init_parents()
        

@dataclass
class DomoLineageLink_Dataflow(DomoLineage_Link):

    @staticmethod
    async def get_entity(entity_id, auth):
        import domolibrary.classes.DomoDataflow as dmdf

        return await dmdf.DomoDataflow.get_by_id(dataflow_id=entity_id, auth=auth)

    @classmethod
    async def get_from_json(cls, obj, auth):
        
        entity = await cls.get_entity(entity_id= obj['id'], auth=auth)

        return cls(
            id=obj["id"],
            auth=auth,
            type="DATAFLOW",
            entity=entity,
            children=obj.get("children", []),
            parents=obj.get("parents", []),
        )

    def __eq__(self, other):
        if not isinstance(other, DomoLineageLink_Dataflow):
            return False
        return self.id == other.id and self.type == other.type


@dataclass
class DomoLineageLink_Publication(DomoLineage_Link):

    @staticmethod
    async def get_entity(entity_id, auth):
        """
        Get the entity associated with this lineage link.
        This method should be implemented by subclasses to return the appropriate entity.
        """
        import domolibrary.classes.DomoPublish as dmpb

        return await dmpb.DomoPublication.get_by_id(
            publication_id=entity_id, auth=auth
        )

    @classmethod
    async def get_from_json(cls, obj, auth):
        """
        Initialize a DomoLineage instance for a publication.
        """

        entity = await cls.get_entity(entity_id=obj["id"], auth=auth)

        return cls(
            id=obj["id"],
            auth=auth,
            type="PUBLICATION",
            entity=entity,
            children=obj.get("children", []),
            parents=obj.get("parents", []),
        )


@dataclass
class DomoLineageLink_Card(DomoLineage_Link):

    @staticmethod
    async def get_entity(entity_id, auth):
        """
        Get the entity associated with this lineage link.
        This method should be implemented by subclasses to return the appropriate entity.
        """
        import domolibrary.classes.DomoCard as dmcd

        return await dmcd.DomoCard.get_by_id(
            publication_id=entity_id, auth=auth
        )

    @classmethod
    async def get_from_json(cls, obj, auth):
        """
        Initialize a DomoLineage instance for a publication.
        """

        entity = await cls.get_entity(entity_id=obj["id"], auth=auth)

        return cls(
            id=obj["id"],
            auth=auth,
            type="CARD",
            entity=entity,
            children=obj.get("children", []),
            parents=obj.get("parents", []),
        )

@dataclass
class DomoLineageLink_Dataset(DomoLineage_Link):

    @staticmethod
    async def get_entity(entity_id, auth):
        """
        Get the entity associated with this lineage link.
        This method should be implemented by subclasses to return the appropriate entity.
        """
        import domolibrary.classes.DomoDataset as dmds

        return await dmds.DomoDataset.get_by_id(dataset_id=entity_id, auth=auth)
                                                
    @classmethod
    async def get_from_json(cls, obj, auth):
        """
        Initialize a DomoLineage instance for a publication.
        """

        entity = await cls.get_entity(entity_id=obj["id"], auth=auth)

        return cls(
            id=obj["id"],
            auth=auth,
            type="DATA_SOURCE",
            entity=entity,
            children=obj.get("children", []),
            parents=obj.get("parents", []),
        )


class DomoLineageLinkTypeFactory_Enum(Enum):
    DATAFLOW = DomoLineageLink_Dataflow
    PUBLICATION = DomoLineageLink_Publication
    DATA_SOURCE = DomoLineageLink_Dataset
    CARD = DomoLineageLink_Card

    @classmethod
    def init(cls, type: str):
        """
        Get the DomoLineage_Link subclass based on the type string.
        """
        try:
            return cls[type]
        
        except KeyError:
            raise dmde.DomoError(f"Unknown lineage link type: {type}")


# %% ../../nbs/classes/50_DomoLineage_v2.ipynb 6
class DomoLineage_ParentTypeEnum(Enum):
    DomoDataflow = "DATAFLOW"
    DomoPublication = "PUBLICATION"
    DomoDataset = "DATA_SOURCE"
    DomoPage = "PAGE"
    DomoCard = "CARD"


@dataclass
class DomoLineage:
    auth: dmda.DomoAuth = field(repr=False)

    parent_id: Any = field(repr=False)
    parent_type: DomoLineage_ParentTypeEnum = field(repr=False)

    parent: Any = field(repr=False, default=None)

    lineage: List[DomoLineage_Link] = field(repr=False, default_factory=list)

    raw_datacenter: dict = field(repr=False, default_factory=dict)

    @classmethod
    def from_parent(cls, parent, auth: dmda.DomoAuth = None):
        """
        Create a DomoLineage instance from a parent entity.
        The parent can be a DomoDataflow, DomoPublication, or DomoDataset.
        """

        return cls(
            auth=auth, parent_id=parent.id, parent_type=DomoLineage_ParentTypeEnum[parent.__class__.__name__], parent=parent
        )

    async def get_datacenter(
        self,
        session: httpx.AsyncClient = None,
        debug_api: bool = False,
        return_raw: bool = False,
    ):

        parent_type = (
            self.parent_type.value
            if isinstance(self.parent_type, DomoLineage_ParentTypeEnum)
            else self.parent_type
        )

        res = await datacenter_routes.get_lineage_upstream(
            auth=self.auth,
            entity_type=parent_type,
            entity_id=self.parent_id or self.parent.id,
            session=session,
            debug_api=debug_api,
        )

        if return_raw:
            return res

        dmcv.merge_dict(res.response, self.raw_datacenter)

        self.lineage +=  await dmce.gather_with_concurrency(
            *[
                DomoLineageLinkTypeFactory_Enum.init(
                    type=obj["type"]
                ).value.get_from_json(obj=obj, auth=self.auth)
                for _, obj in self.raw_datacenter.items()
                if str(obj["id"]) != str(self.parent_id)
            ],
            n=10,
        )

        return self.lineage

    async def get(
        self,
        session: httpx.AsyncClient = None,
        debug_api: bool = False,
        return_raw: bool = False,
    ):
        return await self.get_datacenter(
            session=session, debug_api=debug_api, return_raw=return_raw
        )

# %% ../../nbs/classes/50_DomoLineage_v2.ipynb 11
@dataclass
class DomoLineage_Publication(DomoLineage):
    parent_type: str = "PUBLICATION"

    datasources: List[Any] = field(repr=False, default=None)

    

    async def get(self, return_raw: bool = False,
                  session: httpx.AsyncClient = None, 
                  debug_api: bool = False
                  ):
        
        session = session or httpx.AsyncClient()

        if not self.parent:
            import domolibrary.classes.DomoPublish as dmpb

            self.parent = await dmpb.DomoPublication.get_by_id(
                publication_id=self.parent_id, auth=self.auth
            )

        if return_raw:
            return self.parent.content

        return await dmce.gather_with_concurrency(
            *[
                pc.Lineage.get()
                for pc in self.parent.content
            ],
            n=10
        )
