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

# %% auto 0
__all__ = ['DomoLineage_Type', 'DomoLineage_Instantiation_Error', 'DomoLineage', 'DomoLineage_Card', 'DomoLineage_Publication']

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

import httpx


import domolibrary.routes.datacenter as datacenter_routes

import domolibrary.utils.chunk_execution as dmce

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

from nbdev.showdoc import patch_to

# %% ../../nbs/classes/50_DomoLineage.ipynb 5
async def _get_dataflow( entity_id, auth):
    import domolibrary.classes.DomoDataflow as dmdf
    return await dmdf.DomoDataflow.get_by_id(dataflow_id = entity_id, auth = auth)


async def _get_dataset( entity_id, auth):
    import domolibrary.classes.DomoDataset as dmds
    return await dmds.DomoDataset.get_by_id(dataset_id = entity_id, auth = auth)

async def _get_publication( entity_id, auth):
    import domolibrary.classes.DomoPublish as dmpb
    return await dmpb.DomoPublication.get_by_id(publication_id = entity_id, auth = auth)

# %% ../../nbs/classes/50_DomoLineage.ipynb 6
class DomoLineage_Type(Enum):
    DATAFLOW = partial(_get_dataflow)
    DATA_SOURCE = partial(_get_dataset)
    PUBLICATION = partial(_get_publication)


class DomoLineage_Instantiation_Error(dmde.ClassError):
    def __init__(self, cls_instance, message : str = None):
        super().__init__(
            cls_instance = cls_instance,
            entity_id = self.auth.domo_instance,
            message = "must provide parent or parent_id" or message
        )

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

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

    upstream_lineage_ls: List[Any] = field(repr = False, default=None)
    
    # page_id_ls: List[str] = field(default_factory=list)
    
    # card_id_ls: List[str] = field(default_factory=list)

    # dataflow_id_ls: List[str] = field(default_factory=list)
    # dataset_id_ls: List[str] = field(default_factory=list)

    @classmethod
    def _for_dataset(cls, auth, parent = None, parent_id = None):
        return cls(
            auth = auth,
            parent = parent,
            parent_id = parent_id,
            parent_type = 'DATA_SOURCE'
        )

    @classmethod
    def _for_dataflow(cls, auth, parent = None, parent_id = None):
        return cls(
            auth = auth,
            parent = parent,
            parent_id = parent_id,
            parent_type = 'DATAFLOW'
        )

    @classmethod
    def _for_publication(cls, auth, parent = None, parent_id = None):
        return cls(
            auth = auth,
            parent = parent,
            parent_id = parent_id,
            parent_type = 'PUBLICATION'
        )
    

    def __post_init__(self):
        self.auth = self.auth or (self.parent and self.parent.auth)
        
        if not self.parent_id and self.parent:
            self.parent_id = self.parent.id

        if not self.parent_id and not self.parent:
            raise DomoLineage_Instantiation_Error(cls_instance = self)

# %% ../../nbs/classes/50_DomoLineage.ipynb 7
@patch_to(DomoLineage)
async def get_entity_lineage_upstream(
    self: DomoLineage,
    session: httpx.AsyncClient = None,
    debug_api: bool = False,
    return_raw: bool = False,
):

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

    if return_raw:
        return res

    self.upstream_lineage_ls = await dmce.gather_with_concurrency(
        *[
            DomoLineage_Type[obj["type"]].value(entity_id=obj["id"], auth=self.auth)
            for key, obj in res.response.items()
        ],
        n=10
    )

    return self.upstream_lineage_ls


@patch_to(DomoLineage)
async def get(
    self: DomoLineage,
    total_list: List[Any] = None,
    debug_api: bool = False,
    session: httpx.AsyncClient = None,
    return_raw: bool = False,
    tested_ids = None
):
    tested_ids = tested_ids or []

    res = await self.get_entity_lineage_upstream(debug_api=debug_api, session=session)

    if return_raw:
        return res

    if not total_list:
        total_list = self.upstream_lineage_ls

    for entity in self.upstream_lineage_ls:
        if self.parent_id in tested_ids or self.parent_id == entity.id:
            continue
        
        tested_ids.append(self.parent_id)

        print(tested_ids)


        await entity.Lineage.get(
            total_list=total_list, debug_api=debug_api, session=session, tested_ids=tested_ids
        )

        for child in entity.Lineage.upstream_lineage_ls:
            if child not in total_list:
                total_list.append(child)

    return total_list

# %% ../../nbs/classes/50_DomoLineage.ipynb 10
@dataclass
class DomoLineage_Page(DomoLineage):
    parent_type:str = 'PAGE'
    cards: List[Any] = field(repr=False, default=None)

    @classmethod
    def _for_page(cls, auth, parent = None, parent_id = None):
        return cls(
            auth = auth,
            parent = parent,
            parent_id = parent_id,
            parent_type = 'PAGE'
        )

    async def get(self):
        if not self.parent:
            import domolibrary.classes.DomoPage as dmpg

            self.parent= await dmpg.DomoPage.get_by_id(
                page_id=self.parent_id, auth=self.auth
            )

        self.cards = await self.parent.get_cards()

        return self.cards

# %% ../../nbs/classes/50_DomoLineage.ipynb 12
@dataclass
class DomoLineage_Card(DomoLineage):
    parent_type: str = "CARD"
    datasources: List[Any] = field(repr=False, default=None)

    async def get(self):
        if not self.parent:
            import domolibrary.classes.DomoCard as dmcd

            self.parent = await dmcd.DomoCard.get_by_id(
                card_id=self.parent_id, auth=self.auth
            )

        self.datasources = await self.parent.get_datasources()

        return self.datasources

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

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

    async def get(self, 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
            )

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