from typing import Optional, List

from secoda_common.models.table import Badge

import attr
from marshmallow_annotations.ext.attrs import AttrsSchema


@attr.s(auto_attribs=True, kw_only=True)
class LineageItem:
    key: str  # down/upstream table/col/task key
    level: int  # upstream/downstream distance from current resource
    description: Optional[str]
    badges: Optional[List[Badge]] = None
    usage: Optional[int] = None  # statistic to sort lineage items by


class LineageItemSchema(AttrsSchema):
    class Meta:
        target = LineageItem
        register_as_scheme = True


@attr.s(auto_attribs=True, kw_only=True)
class Lineage:
    key: str  # current table/col/task key
    direction: str  # upstream/downstream/both
    depth: int  # how many levels up/down 0 == all
    upstream_entities: List[LineageItem]  # list of upstream entities
    downstream_entities: List[LineageItem]  # list of downstream entities


class LineageSchema(AttrsSchema):
    class Meta:
        target = Lineage
        register_as_scheme = True