from __future__ import annotations

import ipaddress
from typing import TYPE_CHECKING

from graphene import Field, Int, ObjectType, String
from netaddr import IPSet

from infrahub.core.constants import InfrahubKind
from infrahub.core.ipam.resource_allocator import IPAMResourceAllocator
from infrahub.core.manager import NodeManager
from infrahub.core.protocols import BuiltinIPPrefix
from infrahub.exceptions import NodeNotFoundError, ValidationError
from infrahub.pools.prefix import get_next_available_prefix

if TYPE_CHECKING:
    from graphql import GraphQLResolveInfo

    from infrahub.graphql.initialization import GraphqlContext


class IPAddressGetNextAvailable(ObjectType):
    address = String(required=True)

    @staticmethod
    async def resolve(
        root: dict,  # noqa: ARG004
        info: GraphQLResolveInfo,
        prefix_id: str,
        prefix_length: int | None = None,
    ) -> dict[str, str]:
        graphql_context: GraphqlContext = info.context

        prefix = await NodeManager.get_one(
            id=prefix_id, kind=BuiltinIPPrefix, db=graphql_context.db, branch=graphql_context.branch
        )

        if not prefix:
            raise NodeNotFoundError(
                branch_name=graphql_context.branch.name, node_type=InfrahubKind.IPPREFIX, identifier=prefix_id
            )

        ip_prefix = ipaddress.ip_network(prefix.prefix.value)  # type: ignore[attr-defined]
        prefix_length = prefix_length or ip_prefix.prefixlen

        if not ip_prefix.prefixlen <= prefix_length <= ip_prefix.max_prefixlen:
            raise ValidationError(input_value="Invalid prefix length for current selected prefix")

        namespace = await prefix.ip_namespace.get_peer(db=graphql_context.db)  # type: ignore[attr-defined]
        allocator = IPAMResourceAllocator(db=graphql_context.db, namespace=namespace, branch=graphql_context.branch)
        next_address = await allocator.get_next_address(
            ip_prefix=ip_prefix,
            is_pool=prefix.is_pool.value,  # type: ignore[attr-defined]
        )

        if not next_address:
            raise IndexError("No addresses available in prefix")

        return {"address": f"{next_address.ip}/{prefix_length}"}


class IPPrefixGetNextAvailable(ObjectType):
    prefix = String(required=True)

    @staticmethod
    async def resolve(
        root: dict,  # noqa: ARG004
        info: GraphQLResolveInfo,
        prefix_id: str,
        prefix_length: int | None = None,
    ) -> dict[str, str]:
        graphql_context: GraphqlContext = info.context

        prefix = await NodeManager.get_one(
            id=prefix_id, db=graphql_context.db, branch=graphql_context.branch, kind=BuiltinIPPrefix
        )

        if not prefix:
            raise NodeNotFoundError(
                branch_name=graphql_context.branch.name, node_type=InfrahubKind.IPPREFIX, identifier=prefix_id
            )

        ip_prefix = prefix.prefix.obj
        namespace = await prefix.ip_namespace.get_peer(db=graphql_context.db)
        allocator = IPAMResourceAllocator(db=graphql_context.db, namespace=namespace, branch=graphql_context.branch)

        # Build available pool by removing existing subnets from parent prefix
        subnets = await allocator.get_subnets(ip_prefix=ip_prefix)
        available_pool = IPSet([str(ip_prefix)])
        for subnet in subnets:
            available_pool.remove(str(subnet.prefix))

        if prefix_length is not None and not ip_prefix.prefixlen < prefix_length <= ip_prefix.max_prefixlen:
            raise ValidationError(input_value="Invalid prefix length for current selected prefix")

        try:
            next_prefix = get_next_available_prefix(
                pool=available_pool, prefix_length=prefix_length, prefix_ver=ip_prefix.version
            )
        except ValueError as exc:
            raise IndexError("No prefixes available in prefix") from exc

        return {"prefix": str(next_prefix)}


InfrahubIPAddressGetNextAvailable = Field(
    IPAddressGetNextAvailable,
    prefix_id=String(required=True),
    prefix_length=Int(required=False),
    resolver=IPAddressGetNextAvailable.resolve,
    required=True,
)


InfrahubIPPrefixGetNextAvailable = Field(
    IPPrefixGetNextAvailable,
    prefix_id=String(required=True),
    prefix_length=Int(required=False),
    resolver=IPPrefixGetNextAvailable.resolve,
    required=True,
)

# The following two query fields must be removed once we are sure that people are not using the old queries anymore. Those fields only exist to
# expose a deprecation message.

DeprecatedIPAddressGetNextAvailable = Field(
    IPAddressGetNextAvailable,
    prefix_id=String(required=True),
    prefix_length=Int(required=False),
    resolver=IPAddressGetNextAvailable.resolve,
    required=True,
    deprecation_reason="This query has been renamed to 'InfrahubIPAddressGetNextAvailable'. It will be removed in the next version of Infrahub.",
)


DeprecatedIPPrefixGetNextAvailable = Field(
    IPPrefixGetNextAvailable,
    prefix_id=String(required=True),
    prefix_length=Int(required=False),
    resolver=IPPrefixGetNextAvailable.resolve,
    required=True,
    deprecation_reason="This query has been renamed to 'InfrahubIPPrefixGetNextAvailable'. It will be removed in the next version of Infrahub.",
)
