#  Copyright (C) 2023
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Mike Orlov <m.orlov@abm-jsc.ru>
from typing import MutableMapping, Callable, Optional, cast
from weakref import WeakKeyDictionary

from sqlalchemy.orm import Relationship, InstrumentedAttribute

relation_to_descend: MutableMapping[Relationship, Callable[[Optional[str]], 'EntityProxy']] = WeakKeyDictionary()
relation_to_ascend: MutableMapping[Relationship, Callable[[str], 'EntityProxy']] = WeakKeyDictionary()


def descend(attribute: InstrumentedAttribute, alias: Optional[str] = None) -> 'EntityProxy':
    relation: Relationship = cast(Relationship, attribute.property)
    if descend_ := relation_to_descend.get(relation):
        return descend_(alias)
    else:
        raise KeyError(f"{relation} not found in relation_to_descend mapping")


def ascend(attribute: InstrumentedAttribute, alias: str) -> 'EntityProxy':
    relation: Relationship = cast(Relationship, attribute.property)
    if ascend_ := relation_to_ascend.get(relation):
        return ascend_(alias)
    else:
        raise KeyError(f"{relation} not found in relation_to_descend mapping")
