from _typeshed import Incomplete
from collections.abc import Callable
from types import TracebackType
from typing import Final

from ..recorder import AWSXRayRecorder
from .dummy_entities import DummySubsegment
from .entity import Entity
from .segment import Segment

SUBSEGMENT_RECORDING_ATTRIBUTE: Final = "_self___SUBSEGMENT_RECORDING_ATTRIBUTE__"

def set_as_recording(decorated_func, wrapped) -> None: ...
def is_already_recording(func: Callable[..., object]) -> bool: ...
def subsegment_decorator(wrapped, instance, args, kwargs): ...

class SubsegmentContextManager:
    name: str | None
    subsegment_kwargs: dict[str, str]
    recorder: AWSXRayRecorder
    subsegment: Subsegment | None
    def __init__(self, recorder: AWSXRayRecorder, name: str | None = None, *, namespace: str = "local") -> None: ...
    def __call__(self, wrapped, instance, args: list[Incomplete], kwargs: dict[str, Incomplete]): ...
    def __enter__(self) -> DummySubsegment | Subsegment | None: ...
    def __exit__(
        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
    ) -> None: ...

class Subsegment(Entity):
    parent_segment: Segment
    trace_id: str
    type: str
    namespace: str
    sql: dict[str, Incomplete]
    def __init__(self, name: str, namespace: str, segment: Segment) -> None: ...
    def add_subsegment(self, subsegment: Subsegment) -> None: ...
    def remove_subsegment(self, subsegment: Subsegment) -> None: ...
    def close(self, end_time: float | None = None) -> None: ...
    def set_sql(self, sql: dict[str, Incomplete]) -> None: ...
    def to_dict(self) -> dict[str, Incomplete]: ...
