from __future__ import annotations

from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

from typing_extensions import Self

if TYPE_CHECKING:
    from neo4j import AsyncResult, AsyncSession, AsyncTransaction, Record

    from infrahub.core.schema.schema_branch import SchemaBranch


@runtime_checkable
class Timestamp(Protocol): ...


@runtime_checkable
class NodeSchema(Protocol): ...


@runtime_checkable
class ProfileSchema(Protocol): ...


@runtime_checkable
class TemplateSchema(Protocol): ...


@runtime_checkable
class Branch(Protocol): ...


@runtime_checkable
class InfrahubDatabase(Protocol):
    @property
    def is_session(self) -> bool: ...

    @property
    def is_transaction(self) -> bool: ...

    def add_schema(self, schema: SchemaBranch, name: str | None = None) -> None: ...
    def start_session(self, read_only: bool = False, schemas: list[SchemaBranch] | None = None) -> Self: ...
    def start_transaction(self, schemas: list[SchemaBranch] | None = None) -> Self: ...
    async def session(self) -> AsyncSession: ...
    async def transaction(self, name: str | None) -> AsyncTransaction: ...
    async def close(self) -> None: ...

    async def execute_query(
        self, query: str, params: dict[str, Any] | None = None, name: str = "undefined"
    ) -> list[Record]: ...

    async def execute_query_with_metadata(
        self, query: str, params: dict[str, Any] | None = None, name: str = "undefined"
    ) -> tuple[list[Record], dict[str, Any]]: ...

    async def run_query(
        self, query: str, params: dict[str, Any] | None = None, name: str = "undefined"
    ) -> AsyncResult: ...

    def render_list_comprehension(self, items: str, item_name: str) -> str: ...
    def render_list_comprehension_with_list(self, items: str, item_names: list[str]) -> str: ...
    def render_uuid_generation(self, node_label: str, node_attr: str) -> str: ...
    def get_id_function_name(self) -> str: ...
    def to_database_id(self, db_id: str | int) -> str | int: ...


@runtime_checkable
class CoreNode(Protocol):
    id: str

    def get_id(self) -> str: ...
    def get_kind(self) -> str: ...
    @classmethod
    async def init(
        cls,
        schema: NodeSchema | ProfileSchema | str,
        db: InfrahubDatabase,
        branch: Branch | str | None = None,
        at: Timestamp | str | None = None,
    ) -> Self: ...
    async def new(self, db: InfrahubDatabase, id: str | None = None, **kwargs: Any) -> Self: ...
    async def save(self, db: InfrahubDatabase, at: Timestamp | None = None) -> Self: ...
    async def delete(self, db: InfrahubDatabase, at: Timestamp | None = None) -> None: ...
    async def load(
        self,
        db: InfrahubDatabase,
        id: str | None = None,
        db_id: str | None = None,
        updated_at: Timestamp | str | None = None,
        **kwargs: Any,
    ) -> Self: ...
    async def to_graphql(
        self,
        db: InfrahubDatabase,
        fields: dict | None = None,
        related_node_ids: set | None = None,
        filter_sensitive: bool = False,
        permissions: dict | None = None,
    ) -> dict: ...
    async def render_display_label(self, db: InfrahubDatabase | None = None) -> str: ...
    async def from_graphql(self, data: dict, db: InfrahubDatabase) -> bool: ...
