from __future__ import annotations

import asyncio
import inspect
import json
import warnings
from collections.abc import MutableMapping
from enum import Enum
from time import sleep
from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict
from urllib.parse import urlencode

import httpx
from pydantic import BaseModel, Field

from ..exceptions import (
    BranchNotFoundError,
    InvalidResponseError,
    JsonDecodeError,
    SchemaNotFoundError,
    ValidationError,
)
from ..graphql import Mutation
from ..queries import SCHEMA_HASH_SYNC_STATUS
from .export import RESTRICTED_NAMESPACES, NamespaceExport, SchemaExport, schema_to_export_dict
from .main import (
    AttributeSchema,
    AttributeSchemaAPI,
    BranchSchema,
    BranchSupportType,
    GenericSchema,
    GenericSchemaAPI,
    NodeSchema,
    NodeSchemaAPI,
    ProfileSchemaAPI,
    RelationshipCardinality,
    RelationshipKind,
    RelationshipSchema,
    RelationshipSchemaAPI,
    SchemaRoot,
    SchemaRootAPI,
    TemplateSchemaAPI,
)

if TYPE_CHECKING:
    from ..client import InfrahubClient, InfrahubClientSync, SchemaType, SchemaTypeSync
    from ..node import InfrahubNode, InfrahubNodeSync

    InfrahubNodeTypes: TypeAlias = InfrahubNode | InfrahubNodeSync


__all__ = [
    "AttributeSchema",
    "AttributeSchemaAPI",
    "BranchSupportType",
    "GenericSchema",
    "GenericSchemaAPI",
    "NamespaceExport",
    "NodeSchema",
    "NodeSchemaAPI",
    "ProfileSchemaAPI",
    "RelationshipCardinality",
    "RelationshipKind",
    "RelationshipSchema",
    "RelationshipSchemaAPI",
    "SchemaExport",
    "SchemaRoot",
    "SchemaRootAPI",
    "TemplateSchemaAPI",
    "schema_to_export_dict",
]


class DropdownMutationOptionalArgs(TypedDict):
    color: str | None
    description: str | None
    label: str | None


class DropdownMutation(str, Enum):
    add = "SchemaDropdownAdd"
    remove = "SchemaDropdownRemove"


class EnumMutation(str, Enum):
    add = "SchemaEnumAdd"
    remove = "SchemaEnumRemove"


MainSchemaTypes: TypeAlias = NodeSchema | GenericSchema
MainSchemaTypesAPI: TypeAlias = NodeSchemaAPI | GenericSchemaAPI | ProfileSchemaAPI | TemplateSchemaAPI
MainSchemaTypesAll: TypeAlias = (
    NodeSchema | GenericSchema | NodeSchemaAPI | GenericSchemaAPI | ProfileSchemaAPI | TemplateSchemaAPI
)


class SchemaWarningType(Enum):
    DEPRECATION = "deprecation"


class SchemaWarningKind(BaseModel):
    kind: str = Field(..., description="The kind impacted by the warning")
    field: str | None = Field(default=None, description="The attribute or relationship impacted by the warning")

    @property
    def display(self) -> str:
        suffix = f".{self.field}" if self.field else ""
        return f"{self.kind}{suffix}"


class SchemaWarning(BaseModel):
    type: SchemaWarningType = Field(..., description="The type of warning")
    kinds: list[SchemaWarningKind] = Field(default_factory=list, description="The kinds impacted by the warning")
    message: str = Field(..., description="The message that describes the warning")


class InfrahubSchemaBase:
    client: InfrahubClient | InfrahubClientSync
    cache: dict[str, BranchSchema]

    def __init__(self, client: InfrahubClient | InfrahubClientSync) -> None:
        self.client = client
        self.cache = {}

    @staticmethod
    def _build_export_schemas(
        schema_nodes: MutableMapping[str, MainSchemaTypesAPI],
        namespaces: list[str] | None = None,
    ) -> SchemaExport:
        """Organize fetched schemas into a per-namespace export structure.

        Filters out system types (Profile/Template) and restricted namespaces
        (see :data:`RESTRICTED_NAMESPACES`), and optionally limits to specific
        namespaces.  If the caller requests restricted namespaces they are
        silently excluded and a :func:`warnings.warn` is emitted.

        Returns:
            A :class:`SchemaExport` containing user-defined schemas by namespace.
        """
        if namespaces:
            restricted = set(namespaces) & set(RESTRICTED_NAMESPACES)
            if restricted:
                warnings.warn(
                    f"Restricted namespace(s) {sorted(restricted)} requested but will be excluded from export",
                    stacklevel=3,
                )

        ns_map: dict[str, NamespaceExport] = {}
        for schema in schema_nodes.values():
            if isinstance(schema, (ProfileSchemaAPI, TemplateSchemaAPI)):
                continue
            if schema.namespace in RESTRICTED_NAMESPACES:
                continue
            if namespaces and schema.namespace not in namespaces:
                continue
            ns = schema.namespace
            if ns not in ns_map:
                ns_map[ns] = NamespaceExport()
            schema_dict = schema_to_export_dict(schema)
            if isinstance(schema, GenericSchemaAPI):
                ns_map[ns].generics.append(schema_dict)
            else:
                ns_map[ns].nodes.append(schema_dict)
        return SchemaExport(namespaces=ns_map)

    def validate(self, data: dict[str, Any]) -> None:
        SchemaRoot(**data)

    def validate_data_against_schema(self, schema: MainSchemaTypesAPI, data: dict) -> None:
        for key in data:
            if key not in schema.relationship_names + schema.attribute_names:
                identifier = f"{schema.kind}"
                raise ValidationError(
                    identifier=identifier,
                    message=f"{key} is not a valid value for {identifier}",
                )

    def set_cache(self, schema: dict[str, Any] | SchemaRootAPI | BranchSchema, branch: str | None = None) -> None:
        """
        Set the cache manually (primarily for unit testing)

        Args:
            schema: The schema to set the cache as provided by the /api/schema endpoint either in dict or SchemaRootAPI format
            branch: The name of the branch to set the cache for.
        """
        branch = branch or self.client.default_branch

        if isinstance(schema, SchemaRootAPI):
            schema = BranchSchema.from_schema_root_api(data=schema)
        elif isinstance(schema, dict):
            schema = BranchSchema.from_api_response(data=schema)

        self.cache[branch] = schema

    def generate_payload_create(
        self,
        schema: MainSchemaTypesAPI,
        data: dict,
        source: str | None = None,
        owner: str | None = None,
        is_protected: bool | None = None,
    ) -> dict[str, Any]:
        obj_data: dict[str, Any] = {}
        item_metadata: dict[str, Any] = {}
        if source:
            item_metadata["source"] = str(source)
        if owner:
            item_metadata["owner"] = str(owner)
        if is_protected is not None:
            item_metadata["is_protected"] = is_protected

        for key, value in data.items():
            obj_data[key] = {}
            if key in schema.attribute_names:
                obj_data[key] = {"value": value}
                obj_data[key].update(item_metadata)
            elif key in schema.relationship_names:
                rel = schema.get_relationship(name=key)
                if rel:
                    if rel.cardinality == "one":
                        obj_data[key] = {"id": str(value)}
                        obj_data[key].update(item_metadata)
                    elif rel.cardinality == "many":
                        obj_data[key] = [{"id": str(item)} for item in value]
                        for item in obj_data[key]:
                            item.update(item_metadata)

        return obj_data

    @staticmethod
    def _validate_load_schema_response(response: httpx.Response) -> SchemaLoadResponse:
        if response.status_code == httpx.codes.OK:
            status = response.json()
            return SchemaLoadResponse(
                hash=status["hash"], previous_hash=status["previous_hash"], warnings=status.get("warnings") or []
            )

        if response.status_code in {
            httpx.codes.BAD_REQUEST,
            httpx.codes.UNPROCESSABLE_ENTITY,
            httpx.codes.UNAUTHORIZED,
            httpx.codes.FORBIDDEN,
        }:
            return SchemaLoadResponse(errors=response.json())

        response.raise_for_status()

        raise InvalidResponseError(message=f"Invalid response received from server HTTP {response.status_code}")

    @staticmethod
    def _get_schema_name(schema: type[SchemaType | SchemaTypeSync] | str) -> str:
        if isinstance(schema, str):
            return schema

        if hasattr(schema, "_is_runtime_protocol") and getattr(schema, "_is_runtime_protocol", None):
            if inspect.iscoroutinefunction(schema.save):
                return schema.__name__
            if schema.__name__[-4:] == "Sync":
                return schema.__name__[:-4]
            return schema.__name__

        raise ValueError("schema must be a protocol or a string")

    @staticmethod
    def _parse_schema_response(response: httpx.Response, branch: str) -> MutableMapping[str, Any]:
        if response.status_code == 400:
            raise BranchNotFoundError(
                identifier=branch, message=f"The requested branch was not found on the server [{branch}]"
            )
        response.raise_for_status()

        try:
            data: MutableMapping[str, Any] = response.json()
        except json.decoder.JSONDecodeError as exc:
            raise JsonDecodeError(
                message=f"Invalid Schema response received from the server at {response.url}: {response.text} [{response.status_code}] ",
                content=response.text,
                url=str(response.url),
            ) from exc

        return data

    @staticmethod
    def _deprecated_schema_timeout() -> None:
        warnings.warn(
            "The 'timeout' parameter is deprecated while fetching the schema and will be removed version 2.0.0 of the Infrahub Python SDK. "
            "Use client.default_timeout instead.",
            DeprecationWarning,
            stacklevel=2,
        )


class InfrahubSchema(InfrahubSchemaBase):
    client: InfrahubClient

    async def get(
        self,
        kind: type[SchemaType | SchemaTypeSync] | str,
        branch: str | None = None,
        refresh: bool = False,
        timeout: int | None = None,
    ) -> MainSchemaTypesAPI:
        branch = branch or self.client.default_branch

        kind_str = self._get_schema_name(schema=kind)

        if timeout:
            self._deprecated_schema_timeout()

        if refresh:
            self.cache[branch] = await self._fetch(branch=branch)

        if branch in self.cache and kind_str in self.cache[branch].nodes:
            return self.cache[branch].nodes[kind_str]

        # Fetching the latest schema from the server if we didn't fetch it earlier
        #   because we coulnd't find the object on the local cache
        if not refresh:
            self.cache[branch] = await self._fetch(branch=branch)

        if branch in self.cache and kind_str in self.cache[branch].nodes:
            return self.cache[branch].nodes[kind_str]

        raise SchemaNotFoundError(identifier=kind_str)

    async def all(
        self,
        branch: str | None = None,
        refresh: bool = False,
        namespaces: list[str] | None = None,
        schema_hash: str | None = None,
    ) -> MutableMapping[str, MainSchemaTypesAPI]:
        """Retrieve the entire schema for a given branch.

        if present in cache, the schema will be served from the cache, unless refresh is set to True
        if the schema is not present in the cache, it will be fetched automatically from the server

        Args:
            branch (str, optional): Name of the branch to query. Defaults to default_branch.
            refresh (bool, optional): Force a refresh of the schema. Defaults to False.
            schema_hash (str, optional): Only refresh if the current schema doesn't match this hash.

        Returns:
            dict[str, MainSchemaTypes]: Dictionary of all schema organized by kind
        """
        branch = branch or self.client.default_branch
        if refresh and branch in self.cache and schema_hash and self.cache[branch].hash == schema_hash:
            refresh = False

        if refresh or branch not in self.cache:
            self.cache[branch] = await self._fetch(branch=branch, namespaces=namespaces)

        return self.cache[branch].nodes

    async def load(
        self, schemas: list[dict], branch: str | None = None, wait_until_converged: bool = False
    ) -> SchemaLoadResponse:
        branch = branch or self.client.default_branch
        url = f"{self.client.address}/api/schema/load?branch={branch}"
        response = await self.client._post(
            url=url, timeout=max(120, self.client.default_timeout), payload={"schemas": schemas}
        )

        if wait_until_converged:
            await self.wait_until_converged(branch=branch)

        return self._validate_load_schema_response(response=response)

    async def wait_until_converged(self, branch: str | None = None) -> None:
        """Wait until the schema has converged on the selected branch or the timeout has been reached"""
        waited = 0
        while True:
            if await self.in_sync(branch=branch):
                self.client.log.info(f"Schema successfully converged after {waited} seconds")
                return

            if waited >= self.client.config.schema_converge_timeout:
                self.client.log.warning(f"Schema not converged after {waited} seconds, proceeding regardless")
                return

            waited += 1
            await asyncio.sleep(delay=1)

    async def in_sync(self, branch: str | None = None) -> bool:
        """Indicate if the schema is in sync across all workers for the provided branch"""
        response = await self.client.execute_graphql(query=SCHEMA_HASH_SYNC_STATUS, branch_name=branch)
        return response["InfrahubStatus"]["summary"]["schema_hash_synced"]

    async def check(self, schemas: list[dict], branch: str | None = None) -> tuple[bool, dict | None]:
        branch = branch or self.client.default_branch
        url = f"{self.client.address}/api/schema/check?branch={branch}"
        response = await self.client._post(
            url=url, timeout=max(120, self.client.default_timeout), payload={"schemas": schemas}
        )

        if response.status_code == httpx.codes.ACCEPTED:
            return True, response.json()

        if response.status_code == httpx.codes.UNPROCESSABLE_ENTITY:
            return False, response.json()

        response.raise_for_status()
        return False, None

    async def _get_kind_and_attribute_schema(
        self, kind: str | InfrahubNodeTypes, attribute: str, branch: str | None = None
    ) -> tuple[str, AttributeSchema]:
        node_kind: str = kind._schema.kind if not isinstance(kind, str) else kind
        node_schema = await self.client.schema.get(kind=node_kind, branch=branch)
        schema_attr = node_schema.get_attribute(name=attribute)

        if schema_attr is None:
            raise ValueError(f"Unable to find attribute {attribute}")

        return node_kind, schema_attr

    async def _mutate_enum_attribute(
        self,
        mutation: EnumMutation,
        kind: str | InfrahubNodeTypes,
        attribute: str,
        option: str | int,
        branch: str | None = None,
    ) -> None:
        node_kind, schema_attr = await self._get_kind_and_attribute_schema(
            kind=kind, attribute=attribute, branch=branch
        )

        if schema_attr.enum is None:
            raise ValueError(f"Attribute '{schema_attr.name}' is not of kind Enum")

        input_data = {"data": {"kind": node_kind, "attribute": schema_attr.name, "enum": option}}

        query = Mutation(mutation=mutation.value, input_data=input_data, query={"ok": None})
        await self.client.execute_graphql(
            query=query.render(),
            branch_name=branch,
            tracker=f"mutation-{mutation.name}-add",
            timeout=max(60, self.client.default_timeout),
        )

    async def add_enum_option(
        self, kind: str | InfrahubNodeTypes, attribute: str, option: str | int, branch: str | None = None
    ) -> None:
        await self._mutate_enum_attribute(
            mutation=EnumMutation.add, kind=kind, attribute=attribute, option=option, branch=branch
        )

    async def remove_enum_option(
        self, kind: str | InfrahubNodeTypes, attribute: str, option: str | int, branch: str | None = None
    ) -> None:
        await self._mutate_enum_attribute(
            mutation=EnumMutation.remove, kind=kind, attribute=attribute, option=option, branch=branch
        )

    async def _mutate_dropdown_attribute(
        self,
        mutation: DropdownMutation,
        kind: str | InfrahubNodeTypes,
        attribute: str,
        option: str,
        branch: str | None = None,
        dropdown_optional_args: DropdownMutationOptionalArgs | None = None,
    ) -> None:
        dropdown_optional_args = dropdown_optional_args or DropdownMutationOptionalArgs(
            color="", description="", label=""
        )

        node_kind, schema_attr = await self._get_kind_and_attribute_schema(
            kind=kind, attribute=attribute, branch=branch
        )

        if schema_attr.kind != "Dropdown":
            raise ValueError(f"Attribute '{schema_attr.name}' is not of kind Dropdown")

        input_data: dict[str, Any] = {
            "data": {
                "kind": node_kind,
                "attribute": schema_attr.name,
                "dropdown": option,
            }
        }
        if mutation == DropdownMutation.add:
            input_data["data"].update(dropdown_optional_args)

        query = Mutation(mutation=mutation.value, input_data=input_data, query={"ok": None})
        await self.client.execute_graphql(
            query=query.render(),
            branch_name=branch,
            tracker=f"mutation-{mutation.name}-remove",
            timeout=max(60, self.client.default_timeout),
        )

    async def remove_dropdown_option(
        self, kind: str | InfrahubNodeTypes, attribute: str, option: str, branch: str | None = None
    ) -> None:
        await self._mutate_dropdown_attribute(
            mutation=DropdownMutation.remove, kind=kind, attribute=attribute, option=option, branch=branch
        )

    async def add_dropdown_option(
        self,
        kind: str | InfrahubNodeTypes,
        attribute: str,
        option: str,
        color: str | None = "",
        description: str | None = "",
        label: str | None = "",
        branch: str | None = None,
    ) -> None:
        dropdown_optional_args = DropdownMutationOptionalArgs(color=color, description=description, label=label)
        await self._mutate_dropdown_attribute(
            mutation=DropdownMutation.add,
            kind=kind,
            attribute=attribute,
            option=option,
            branch=branch,
            dropdown_optional_args=dropdown_optional_args,
        )

    async def fetch(
        self, branch: str, namespaces: list[str] | None = None, timeout: int | None = None, populate_cache: bool = True
    ) -> MutableMapping[str, MainSchemaTypesAPI]:
        """Fetch the schema from the server for a given branch.

        Args:
            branch: Name of the branch to fetch the schema for.
            timeout: Overrides default timeout used when querying the schema. deprecated.
            populate_cache: Whether to populate the cache with the fetched schema. Defaults to True.

        Returns:
            dict[str, MainSchemaTypes]: Dictionary of all schema organized by kind
        """

        if timeout:
            self._deprecated_schema_timeout()

        branch_schema = await self._fetch(branch=branch, namespaces=namespaces)

        if populate_cache:
            self.cache[branch] = branch_schema

        return branch_schema.nodes

    async def export(
        self,
        branch: str | None = None,
        namespaces: list[str] | None = None,
    ) -> SchemaExport:
        """Export user-defined schemas organized by namespace.

        Fetches schemas from the server, filters out system types and
        restricted namespaces (see :data:`RESTRICTED_NAMESPACES`), and returns
        a :class:`SchemaExport` object with per-namespace data.  Restricted
        namespaces such as ``Core`` and ``Builtin`` are always excluded even if
        explicitly listed in *namespaces*; a warning is emitted when this
        happens.

        Args:
            branch: Branch to export from. Defaults to default_branch.
            namespaces: Optional list of namespaces to include. If empty/None,
                all user-defined namespaces are exported.

        Returns:
            A :class:`SchemaExport` containing user-defined schemas by namespace.
        """
        branch = branch or self.client.default_branch
        schema_nodes = await self.fetch(branch=branch, namespaces=namespaces, populate_cache=False)
        return self._build_export_schemas(schema_nodes=schema_nodes, namespaces=namespaces)

    async def get_graphql_schema(self, branch: str | None = None) -> str:
        """Get the GraphQL schema as a string.

        Args:
            branch: The branch to get the schema for. Defaults to default_branch.

        Returns:
            The GraphQL schema as a string.
        """
        branch = branch or self.client.default_branch
        url = f"{self.client.address}/schema.graphql?branch={branch}"

        response = await self.client._get(url=url)

        if response.status_code != 200:
            raise ValueError(f"Failed to fetch GraphQL schema: HTTP {response.status_code} - {response.text}")

        return response.text

    async def _fetch(self, branch: str, namespaces: list[str] | None = None) -> BranchSchema:
        url_parts = [("branch", branch)]
        if namespaces:
            url_parts.extend([("namespaces", ns) for ns in namespaces])
        query_params = urlencode(url_parts)
        url = f"{self.client.address}/api/schema?{query_params}"

        response = await self.client._get(url=url)

        data = self._parse_schema_response(response=response, branch=branch)

        return BranchSchema.from_api_response(data=data)


class InfrahubSchemaSync(InfrahubSchemaBase):
    client: InfrahubClientSync

    def all(
        self,
        branch: str | None = None,
        refresh: bool = False,
        namespaces: list[str] | None = None,
        schema_hash: str | None = None,
    ) -> MutableMapping[str, MainSchemaTypesAPI]:
        """Retrieve the entire schema for a given branch.

        if present in cache, the schema will be served from the cache, unless refresh is set to True
        if the schema is not present in the cache, it will be fetched automatically from the server

        Args:
            branch (str, optional): Name of the branch to query. Defaults to default_branch.
            refresh (bool, optional): Force a refresh of the schema. Defaults to False.
            schema_hash (str, optional): Only refresh if the current schema doesn't match this hash.

        Returns:
            dict[str, MainSchemaTypes]: Dictionary of all schema organized by kind
        """
        branch = branch or self.client.default_branch
        if refresh and branch in self.cache and schema_hash and self.cache[branch].hash == schema_hash:
            refresh = False

        if refresh or branch not in self.cache:
            self.cache[branch] = self._fetch(branch=branch, namespaces=namespaces)

        return self.cache[branch].nodes

    def get(
        self,
        kind: type[SchemaType | SchemaTypeSync] | str,
        branch: str | None = None,
        refresh: bool = False,
        timeout: int | None = None,
    ) -> MainSchemaTypesAPI:
        """
        Retrieve a specific schema object from the server.

        Args:
            kind: The kind of schema object to retrieve.
            branch: The branch to retrieve the schema from.
            refresh: Whether to refresh the schema.
            timeout: Overrides default timeout used when querying the GraphQL API. Specified in seconds (deprecated).

        Returns:
            MainSchemaTypes: The schema object.
        """
        branch = branch or self.client.default_branch

        kind_str = self._get_schema_name(schema=kind)

        if timeout:
            self._deprecated_schema_timeout()

        if refresh:
            self.cache[branch] = self._fetch(branch=branch)

        if branch in self.cache and kind_str in self.cache[branch].nodes:
            return self.cache[branch].nodes[kind_str]

        # Fetching the latest schema from the server if we didn't fetch it earlier
        #   because we coulnd't find the object on the local cache
        if not refresh:
            self.cache[branch] = self._fetch(branch=branch)

        if branch in self.cache and kind_str in self.cache[branch].nodes:
            return self.cache[branch].nodes[kind_str]

        raise SchemaNotFoundError(identifier=kind_str)

    def _get_kind_and_attribute_schema(
        self, kind: str | InfrahubNodeTypes, attribute: str, branch: str | None = None
    ) -> tuple[str, AttributeSchemaAPI]:
        node_kind: str = kind._schema.kind if not isinstance(kind, str) else kind
        node_schema = self.client.schema.get(kind=node_kind, branch=branch)
        schema_attr = node_schema.get_attribute(name=attribute)

        if schema_attr is None:
            raise ValueError(f"Unable to find attribute {attribute}")

        return node_kind, schema_attr

    def _mutate_enum_attribute(
        self,
        mutation: EnumMutation,
        kind: str | InfrahubNodeTypes,
        attribute: str,
        option: str | int,
        branch: str | None = None,
    ) -> None:
        node_kind, schema_attr = self._get_kind_and_attribute_schema(kind=kind, attribute=attribute, branch=branch)

        if schema_attr.enum is None:
            raise ValueError(f"Attribute '{schema_attr.name}' is not of kind Enum")

        input_data = {"data": {"kind": node_kind, "attribute": schema_attr.name, "enum": option}}

        query = Mutation(mutation=mutation.value, input_data=input_data, query={"ok": None})
        self.client.execute_graphql(
            query=query.render(),
            branch_name=branch,
            tracker=f"mutation-{mutation.name}-add",
            timeout=max(60, self.client.default_timeout),
        )

    def add_enum_option(
        self, kind: str | InfrahubNodeTypes, attribute: str, option: str | int, branch: str | None = None
    ) -> None:
        self._mutate_enum_attribute(
            mutation=EnumMutation.add, kind=kind, attribute=attribute, option=option, branch=branch
        )

    def remove_enum_option(
        self, kind: str | InfrahubNodeTypes, attribute: str, option: str | int, branch: str | None = None
    ) -> None:
        self._mutate_enum_attribute(
            mutation=EnumMutation.remove, kind=kind, attribute=attribute, option=option, branch=branch
        )

    def _mutate_dropdown_attribute(
        self,
        mutation: DropdownMutation,
        kind: str | InfrahubNodeTypes,
        attribute: str,
        option: str,
        branch: str | None = None,
        dropdown_optional_args: DropdownMutationOptionalArgs | None = None,
    ) -> None:
        dropdown_optional_args = dropdown_optional_args or DropdownMutationOptionalArgs(
            color="", description="", label=""
        )
        node_kind, schema_attr = self._get_kind_and_attribute_schema(kind=kind, attribute=attribute, branch=branch)

        if schema_attr.kind != "Dropdown":
            raise ValueError(f"Attribute '{schema_attr.name}' is not of kind Dropdown")

        input_data: dict[str, Any] = {
            "data": {
                "kind": node_kind,
                "attribute": schema_attr.name,
                "dropdown": option,
            }
        }

        if mutation == DropdownMutation.add:
            input_data["data"].update(dropdown_optional_args)

        query = Mutation(mutation=mutation.value, input_data=input_data, query={"ok": None})
        self.client.execute_graphql(
            query=query.render(),
            branch_name=branch,
            tracker=f"mutation-{mutation.name}-remove",
            timeout=max(60, self.client.default_timeout),
        )

    def remove_dropdown_option(
        self, kind: str | InfrahubNodeTypes, attribute: str, option: str, branch: str | None = None
    ) -> None:
        self._mutate_dropdown_attribute(
            mutation=DropdownMutation.remove, kind=kind, attribute=attribute, option=option, branch=branch
        )

    def add_dropdown_option(
        self,
        kind: str | InfrahubNodeTypes,
        attribute: str,
        option: str,
        color: str | None = "",
        description: str | None = "",
        label: str | None = "",
        branch: str | None = None,
    ) -> None:
        dropdown_optional_args = DropdownMutationOptionalArgs(color=color, description=description, label=label)
        self._mutate_dropdown_attribute(
            mutation=DropdownMutation.add,
            kind=kind,
            attribute=attribute,
            option=option,
            branch=branch,
            dropdown_optional_args=dropdown_optional_args,
        )

    def fetch(
        self, branch: str, namespaces: list[str] | None = None, timeout: int | None = None, populate_cache: bool = True
    ) -> MutableMapping[str, MainSchemaTypesAPI]:
        """Fetch the schema from the server for a given branch.

        Args:
            branch: Name of the branch to fetch the schema for.
            timeout: Overrides default timeout used when querying the GraphQL API. Specified in seconds (deprecated).
            populate_cache: Whether to populate the cache with the fetched schema. Defaults to True.

        Returns:
            dict[str, MainSchemaTypes]: Dictionary of all schema organized by kind
        """
        if timeout:
            self._deprecated_schema_timeout()

        branch_schema = self._fetch(branch=branch, namespaces=namespaces)

        if populate_cache:
            self.cache[branch] = branch_schema

        return branch_schema.nodes

    def export(
        self,
        branch: str | None = None,
        namespaces: list[str] | None = None,
    ) -> SchemaExport:
        """Export user-defined schemas organized by namespace.

        Fetches schemas from the server, filters out system types and
        restricted namespaces (see :data:`RESTRICTED_NAMESPACES`), and returns
        a :class:`SchemaExport` object with per-namespace data.  Restricted
        namespaces such as ``Core`` and ``Builtin`` are always excluded even if
        explicitly listed in *namespaces*; a warning is emitted when this
        happens.

        Args:
            branch: Branch to export from. Defaults to default_branch.
            namespaces: Optional list of namespaces to include. If empty/None,
                all user-defined namespaces are exported.

        Returns:
            A :class:`SchemaExport` containing user-defined schemas by namespace.
        """
        branch = branch or self.client.default_branch
        schema_nodes = self.fetch(branch=branch, namespaces=namespaces, populate_cache=False)
        return self._build_export_schemas(schema_nodes=schema_nodes, namespaces=namespaces)

    def get_graphql_schema(self, branch: str | None = None) -> str:
        """Get the GraphQL schema as a string.

        Args:
            branch: The branch to get the schema for. Defaults to default_branch.

        Returns:
            The GraphQL schema as a string.
        """
        branch = branch or self.client.default_branch
        url = f"{self.client.address}/schema.graphql?branch={branch}"

        response = self.client._get(url=url)

        if response.status_code != 200:
            raise ValueError(f"Failed to fetch GraphQL schema: HTTP {response.status_code} - {response.text}")

        return response.text

    def _fetch(self, branch: str, namespaces: list[str] | None = None) -> BranchSchema:
        url_parts = [("branch", branch)]
        if namespaces:
            url_parts.extend([("namespaces", ns) for ns in namespaces])
        query_params = urlencode(url_parts)
        url = f"{self.client.address}/api/schema?{query_params}"
        response = self.client._get(url=url)

        data = self._parse_schema_response(response=response, branch=branch)

        return BranchSchema.from_api_response(data=data)

    def load(
        self, schemas: list[dict], branch: str | None = None, wait_until_converged: bool = False
    ) -> SchemaLoadResponse:
        branch = branch or self.client.default_branch
        url = f"{self.client.address}/api/schema/load?branch={branch}"
        response = self.client._post(
            url=url, timeout=max(120, self.client.default_timeout), payload={"schemas": schemas}
        )

        if wait_until_converged:
            self.wait_until_converged(branch=branch)

        return self._validate_load_schema_response(response=response)

    def wait_until_converged(self, branch: str | None = None) -> None:
        """Wait until the schema has converged on the selected branch or the timeout has been reached"""
        waited = 0
        while True:
            if self.in_sync(branch=branch):
                self.client.log.info(f"Schema successfully converged after {waited} seconds")
                return

            if waited >= self.client.config.schema_converge_timeout:
                self.client.log.warning(f"Schema not converged after {waited} seconds, proceeding regardless")
                return

            waited += 1
            sleep(1)

    def in_sync(self, branch: str | None = None) -> bool:
        """Indicate if the schema is in sync across all workers for the provided branch"""
        response = self.client.execute_graphql(query=SCHEMA_HASH_SYNC_STATUS, branch_name=branch)
        return response["InfrahubStatus"]["summary"]["schema_hash_synced"]

    def check(self, schemas: list[dict], branch: str | None = None) -> tuple[bool, dict | None]:
        branch = branch or self.client.default_branch
        url = f"{self.client.address}/api/schema/check?branch={branch}"
        response = self.client._post(
            url=url, timeout=max(120, self.client.default_timeout), payload={"schemas": schemas}
        )

        if response.status_code == httpx.codes.ACCEPTED:
            return True, response.json()

        if response.status_code == httpx.codes.UNPROCESSABLE_ENTITY:
            return False, response.json()

        response.raise_for_status()
        return False, None


class SchemaLoadResponse(BaseModel):
    hash: str = Field(default="", description="The new hash for the entire schema")
    previous_hash: str = Field(default="", description="The previous hash for the entire schema")
    errors: dict = Field(default_factory=dict, description="Errors reported by the server")
    warnings: list[SchemaWarning] = Field(default_factory=list, description="Warnings reported by the server")

    @property
    def schema_updated(self) -> bool:
        return bool(self.hash and self.previous_hash and self.hash != self.previous_hash)
