"""
GraphQL query mixin for annotation, relationship, label, labelset, and note queries.
"""

import logging
import re
from typing import Any

import graphene
from django.db.models import Q
from graphene import relay
from graphene_django.fields import DjangoConnectionField
from graphene_django.filter import DjangoFilterConnectionField
from graphql import GraphQLError
from graphql_jwt.decorators import login_required
from graphql_relay import from_global_id, to_global_id

from config.graphql.filters import (
    AuthorityFrontierFilter,
    AuthorityKeyEquivalenceFilter,
    AuthorityNamespaceFilter,
    LabelFilter,
    LabelsetFilter,
    RelationshipFilter,
)
from config.graphql.graphene_types import (
    AnnotationLabelType,
    AnnotationType,
    AuthorityDetailType,
    AuthorityFrontierNode,
    AuthorityFrontierStatsType,
    AuthorityKeyEquivalenceNode,
    AuthorityMappingStatsType,
    AuthorityNamespaceNode,
    AuthorityNamespaceStatsType,
    AuthoritySourceProviderType,
    CorpusReferenceType,
    GovernanceGraphCorpusType,
    GovernanceGraphEdgeType,
    GovernanceGraphNodeType,
    GovernanceGraphType,
    LabelSetType,
    NoteType,
    PageAwareAnnotationType,
    PdfPageInfoType,
    RelationshipType,
    WantedAuthorityType,
)
from config.graphql.ratelimits import get_user_tier_rate, graphql_ratelimit_dynamic
from opencontractserver.annotations.models import (
    Annotation,
    AnnotationLabel,
    LabelSet,
    Note,
    Relationship,
)
from opencontractserver.constants.annotations import (
    DOCUMENT_ANNOTATION_INDEX_LIMIT,
    MANUAL_ANNOTATION_SENTINEL,
)
from opencontractserver.constants.stats import GOVERNANCE_GRAPH_MAX_NODES
from opencontractserver.documents.models import Document
from opencontractserver.enrichment import constants as enrichment_constants
from opencontractserver.shared.services.base import BaseService
from opencontractserver.types.enums import LabelType

logger = logging.getLogger(__name__)


class AnnotationQueryMixin:
    """Query fields and resolvers for annotation, relationship, label, labelset, and note queries."""

    # CORPUS REFERENCE RESOLVERS ###############################
    corpus_references = DjangoConnectionField(
        CorpusReferenceType,
        corpus_id=graphene.ID(required=True),
        reference_type=graphene.String(),
        canonical_key=graphene.String(),
        document_id=graphene.ID(
            description=(
                "Restrict to references touching this document on EITHER side "
                "(source mention's document or resolved target document) — the "
                "single-fetch shape the document References panel needs."
            )
        ),
    )

    def resolve_corpus_references(self, info, corpus_id, **kwargs) -> Any:
        """List enrichment cross-references for a corpus the user can read.

        Visibility is enforced by ``CorpusReferenceService`` (corpus-derived);
        no inline Tier-0 permission fusion here.
        """
        from opencontractserver.annotations.models import CorpusReference
        from opencontractserver.documents.models import Document
        from opencontractserver.enrichment.services import CorpusReferenceService
        from opencontractserver.shared.services.base import BaseService

        pk_str = from_global_id(corpus_id)[1]
        if not str(pk_str).isdigit():
            return CorpusReference.objects.none()
        pk = int(pk_str)
        qs = CorpusReferenceService.for_corpus(info.context.user, pk)
        if kwargs.get("reference_type"):
            qs = qs.filter(reference_type=kwargs["reference_type"])
        if kwargs.get("canonical_key"):
            qs = qs.filter(canonical_key=kwargs["canonical_key"])
        if kwargs.get("document_id"):
            doc_pk_str = from_global_id(kwargs["document_id"])[1]
            if not str(doc_pk_str).isdigit():
                return CorpusReference.objects.none()
            doc_pk = int(doc_pk_str)
            # IDOR: validate the document is READ-visible to the caller before
            # filtering by it. Without this a corpus reader could probe whether
            # an arbitrary (possibly invisible) document has references in this
            # corpus. An invisible document yields the same empty result as one
            # with no references.
            if (
                not BaseService.filter_visible(
                    Document, info.context.user, request=info.context
                )
                .filter(id=doc_pk)
                .exists()
            ):
                return CorpusReference.objects.none()
            qs = qs.filter(
                Q(source_annotation__document_id=doc_pk) | Q(target_document_id=doc_pk)
            )
        # Pull the FK targets the type resolves in one pass — without this each
        # CorpusReferenceType row fires a separate query per FK (N+1).
        return qs.select_related(
            "source_annotation",
            "corpus",
            "target_document",
            "target_annotation",
            "target_corpus",
        )

    governance_graph = graphene.Field(
        GovernanceGraphType,
        corpus_id=graphene.ID(required=True),
        limit=graphene.Int(required=False),
        description=(
            "The corpus-scoped reference web in node-link form: documents, "
            "statute sections, and external-citation ghost nodes, with "
            "mention-weighted LAW / LAW_EXTERNAL / DOCUMENT edges. Powers the "
            "Governance Graph panel on the Corpus Intelligence home."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_governance_graph(self, info, corpus_id, limit=None) -> Any:
        """Build the governance graph through ``GovernanceGraphService``.

        All visibility decisions (corpus READ gate, per-document READ checks,
        ghost degradation for invisible targets) live in the service; this
        resolver only translates raw PKs / canonical keys into relay ids.
        """
        # Function-local service import (services import GraphQL types
        # transitively — module-level import would cycle).
        from opencontractserver.enrichment.services import GovernanceGraphService

        empty = GovernanceGraphType(
            corpora=[],
            nodes=[],
            edges=[],
            document_count=0,
            external_key_count=0,
            edge_count=0,
            mention_count=0,
            truncated=False,
        )

        corpus_pk = from_global_id(corpus_id)[1]
        # Malformed/empty global ids decode to a non-numeric pk; treat as
        # not-found (empty graph) rather than erroring.
        if not str(corpus_pk).isdigit():
            return empty

        node_cap = GOVERNANCE_GRAPH_MAX_NODES
        if limit is not None and 0 < limit < node_cap:
            node_cap = limit

        data = GovernanceGraphService.build(
            info.context.user, int(corpus_pk), node_cap, request=info.context
        )
        if data is None:
            return empty

        def _node_id(endpoint) -> str:
            kind, val = endpoint
            if kind == "doc":
                return to_global_id("DocumentType", val)
            return f"key:{val}"

        nodes = [
            GovernanceGraphNodeType(
                id=to_global_id("DocumentType", n["doc_pk"]),
                document_id=to_global_id("DocumentType", n["doc_pk"]),
                title=n["title"],
                kind=n["kind"],
                corpus_id=(
                    to_global_id("CorpusType", n["corpus_pk"])
                    if n["corpus_pk"]
                    else None
                ),
                authority=n["authority"],
                jurisdiction=n.get("jurisdiction"),
                authority_type=n.get("authority_type"),
                discovery_state=None,
                degree=n["degree"],
            )
            for n in data["doc_nodes"]
        ] + [
            GovernanceGraphNodeType(
                id=f"key:{g['key']}",
                document_id=None,
                title=g["key"],
                kind=enrichment_constants.GRAPH_NODE_EXTERNAL,
                corpus_id=None,
                authority=g["authority"],
                jurisdiction=g.get("jurisdiction"),
                authority_type=g.get("authority_type"),
                discovery_state=g.get("discovery_state"),
                degree=g["degree"],
            )
            for g in data["ghost_nodes"]
        ]

        return GovernanceGraphType(
            corpora=[
                GovernanceGraphCorpusType(
                    id=to_global_id("CorpusType", c["corpus_pk"]),
                    title=c["title"],
                    kind=c["kind"],
                )
                for c in data["corpora"]
            ],
            nodes=nodes,
            edges=[
                GovernanceGraphEdgeType(
                    source=_node_id(e["source"]),
                    target=_node_id(e["target"]),
                    edge_type=e["edge_type"],
                    weight=e["weight"],
                )
                for e in data["edges"]
            ],
            document_count=data["document_count"],
            external_key_count=data["external_key_count"],
            edge_count=data["edge_count"],
            mention_count=data["mention_count"],
            truncated=data["truncated"],
        )

    wanted_authorities = graphene.List(
        graphene.NonNull(WantedAuthorityType),
        corpus_id=graphene.ID(
            required=False,
            description="Restrict the backlog to one corpus; omit for all visible.",
        ),
        required=True,
        description=(
            "The missing-authority backlog: EXTERNAL law citations visible to "
            "the user, aggregated by authority prefix and ranked by mention "
            "volume — what to bootstrap next to resolve the most references."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_wanted_authorities(self, info, corpus_id=None) -> Any:
        """Aggregate through ``CorpusReferenceService`` (visibility-scoped);
        the service returns plain dicts that graphene's default resolver maps
        onto ``WantedAuthorityType`` fields."""
        from opencontractserver.enrichment.services import CorpusReferenceService

        pk: int | None = None
        if corpus_id:
            pk_str = from_global_id(corpus_id)[1]
            if not str(pk_str).isdigit():
                return []
            pk = int(pk_str)
        return CorpusReferenceService.wanted_authorities(
            info.context.user, corpus_id=pk
        )

    # AUTHORITY FRONTIER (global, superuser-only) ##############
    authority_frontier = DjangoFilterConnectionField(
        AuthorityFrontierNode,
        filterset_class=AuthorityFrontierFilter,
        description=(
            "Global authority-source discovery queue (AuthorityFrontier): the "
            "crawl/ingestion state of every wanted section-root key across all "
            "corpora, ranked by citation demand. SUPERUSER-ONLY (empty "
            "otherwise) — gating + default order live on the node's get_queryset."
        ),
    )

    authority_frontier_stats = graphene.Field(
        AuthorityFrontierStatsType,
        jurisdiction=graphene.String(required=False),
        authority_type=graphene.String(required=False),
        provider=graphene.String(required=False),
        authority=graphene.String(required=False),
        search=graphene.String(required=False),
        required=True,
        description=(
            "Facet-aware per-discovery_state row counts for the authority-"
            "sources monitor's summary chips. Honours the non-state facets but "
            "not a state filter. SUPERUSER-ONLY (empty otherwise)."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_LIGHT"))
    def resolve_authority_frontier_stats(
        self,
        info,
        jurisdiction=None,
        authority_type=None,
        provider=None,
        authority=None,
        search=None,
    ) -> Any:
        """Delegate the (superuser-gated) aggregation to the service; graphene's
        default resolver maps the returned dict onto ``AuthorityFrontierStatsType``."""
        from opencontractserver.enrichment.services import AuthorityFrontierService

        return AuthorityFrontierService.admin_state_counts(
            info.context.user,
            jurisdiction=jurisdiction,
            authority_type=authority_type,
            provider=provider,
            authority=authority,
            search=search,
        )

    # AUTHORITY MAPPINGS (global, superuser-only) ##############
    authority_key_equivalences = DjangoFilterConnectionField(
        AuthorityKeyEquivalenceNode,
        filterset_class=AuthorityKeyEquivalenceFilter,
        description=(
            "Runtime authority key-equivalence registry (AuthorityKeyEquivalence): "
            "act-section ↔ USC/CFR codification synonyms used to bridge citations "
            "across namespaces. SUPERUSER-ONLY (empty otherwise) — gating + "
            "default order live on the node's get_queryset."
        ),
    )

    authority_mapping_stats = graphene.Field(
        AuthorityMappingStatsType,
        search=graphene.String(required=False),
        required=True,
        description=(
            "Facet-aware per-source row counts for the authority-mappings panel's "
            "summary chips. Honours the search facet but not a source filter. "
            "SUPERUSER-ONLY (empty otherwise)."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_LIGHT"))
    def resolve_authority_mapping_stats(self, info, search=None) -> Any:
        """Delegate the (superuser-gated) aggregation to the service; graphene's
        default resolver maps the returned dict onto ``AuthorityMappingStatsType``."""
        from opencontractserver.enrichment.services import (
            AuthorityKeyEquivalenceService,
        )

        return AuthorityKeyEquivalenceService.stats(info.context.user, search=search)

    # AUTHORITY NAMESPACES (global registry, superuser-only) ###
    authority_namespaces = DjangoFilterConnectionField(
        AuthorityNamespaceNode,
        filterset_class=AuthorityNamespaceFilter,
        description=(
            "The registry of bodies of law (AuthorityNamespace): one row per "
            "canonical-key prefix (e.g. 'usc-15', 'dgcl') whose aliases drive "
            "Tier-1 citation extraction. SUPERUSER-ONLY (empty otherwise) — "
            "gating + default order live on the node's get_queryset."
        ),
    )

    authority_namespace_stats = graphene.Field(
        AuthorityNamespaceStatsType,
        search=graphene.String(required=False),
        required=True,
        description=(
            "Faceted per-jurisdiction / authority_type / scope row counts for "
            "the registry panel's summary chips. Honours the search facet but "
            "not the facet selects. SUPERUSER-ONLY (empty otherwise)."
        ),
    )

    authority_namespace_detail = graphene.Field(
        AuthorityDetailType,
        prefix=graphene.String(required=True),
        description=(
            "Everything about one body of law, string-joined across the "
            "authority models: the namespace + its aliases, in/out key-"
            "equivalences, discovery-queue rows, and reference demand. "
            "SUPERUSER-ONLY (null otherwise or for an unknown prefix)."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_LIGHT"))
    def resolve_authority_namespace_stats(self, info, search=None) -> Any:
        from opencontractserver.enrichment.services import AuthorityNamespaceService

        return AuthorityNamespaceService.stats(info.context.user, search=search)

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_authority_namespace_detail(self, info, prefix) -> Any:
        from opencontractserver.enrichment.services import AuthorityNamespaceService

        return AuthorityNamespaceService.detail(info.context.user, prefix)

    # AUTHORITY SOURCE PROVIDERS (registry, superuser-only) ####
    authority_source_providers = graphene.List(
        graphene.NonNull(AuthoritySourceProviderType),
        required=True,
        description=(
            "The registered authority source providers (scrapers): US Code / "
            "eCFR / Federal Register / agentic web locator, with their supported "
            "prefixes, license, priority, enabled flag and whether the secrets "
            "vault holds credentials. SUPERUSER-ONLY (empty otherwise)."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_LIGHT"))
    def resolve_authority_source_providers(self, info) -> Any:
        from opencontractserver.enrichment.services import (
            AuthoritySourceProviderService,
        )

        return AuthoritySourceProviderService.list_providers(info.context.user)

    # ANNOTATION RESOLVERS #####################################
    annotations = DjangoConnectionField(
        AnnotationType,
        raw_text_contains=graphene.String(),
        annotation_label_id=graphene.ID(),
        annotation_label__text=graphene.String(),
        annotation_label__text_contains=graphene.String(),
        annotation_label__description_contains=graphene.String(),
        annotation_label__label_type=graphene.String(),
        analysis_isnull=graphene.Boolean(),
        corpus_action_isnull=graphene.Boolean(),
        agent_created=graphene.Boolean(),
        document_id=graphene.ID(),
        corpus_id=graphene.ID(),
        structural=graphene.Boolean(),
        uses_label_from_labelset_id=graphene.ID(),
        created_by_analysis_ids=graphene.String(),
        created_with_analyzer_id=graphene.String(),
        order_by=graphene.String(),
        # Higher limit for Document Annotation Index which needs full hierarchy
        max_limit=DOCUMENT_ANNOTATION_INDEX_LIMIT,
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_annotations(
        self,
        info,
        analysis_isnull=None,
        structural=None,
        corpus_action_isnull=None,
        agent_created=None,
        **kwargs,
    ) -> Any:
        # Import the query optimizer
        from opencontractserver.annotations.services import AnnotationService

        document_id = kwargs.get("document_id")
        corpus_id = kwargs.get("corpus_id")

        # Decoded PKs of the requested context, used below to scope the
        # structural-set document prefetch so structural annotations
        # (document_id=NULL) resolve to the context-local document.
        doc_django_pk: int | None = None
        corpus_django_pk: int | None = None

        if document_id:
            # Use document-specific query optimizer
            doc_django_pk = int(from_global_id(document_id)[1])
            corpus_django_pk = int(from_global_id(corpus_id)[1]) if corpus_id else None

            # Use query optimizer which handles permissions properly
            queryset = AnnotationService.get_document_annotations(
                document_id=doc_django_pk,
                user=info.context.user,
                corpus_id=corpus_django_pk,
                analysis_id=None,  # Will be handled below if needed
                extract_id=None,
            )

        elif corpus_id:
            # Use corpus-wide query optimizer (handles structural annotations correctly)
            # This optimizer already applies structural, analysis_isnull, and corpus filters
            corpus_django_pk = int(from_global_id(corpus_id)[1])
            queryset = AnnotationService.get_corpus_annotations(
                corpus_id=corpus_django_pk,
                user=info.context.user,
                structural=structural,
                analysis_isnull=analysis_isnull,
                context=info.context,
            )
            # Mark filters already applied by optimizer to prevent double-filtering
            corpus_id = None
            structural = None
            analysis_isnull = None

        else:
            # Fallback to visible_to_user for queries without document or
            # corpus. This un-scoped "Browse annotations" path uses a cached
            # exact totalCount (scoped paths above keep live counts) — see
            # ``CachedCountQuerySetMixin`` / issue #1908.
            queryset = BaseService.filter_visible(
                Annotation, info.context.user, request=info.context
            ).with_cached_count()

        queryset = queryset.select_related(
            "annotation_label",
            "creator",
            "document",
            "document__creator",
            "corpus",
            "analysis",
            "analysis__analyzer",
            "corpus_action",
            "structural_set",
        ).prefetch_related(
            # Scope the structural-set documents to the requested corpus (or
            # document) so AnnotationType.resolve_document returns the
            # context-local copy rather than an arbitrary member of a
            # content-hash-shared StructuralAnnotationSet. See
            # AnnotationService.structural_document_prefetch.
            AnnotationService.structural_document_prefetch(
                user=info.context.user,
                corpus_id=corpus_django_pk,
                document_id=doc_django_pk,
            ),
        )

        # Filter by uses_label_from_labelset_id
        labelset_id = kwargs.get("uses_label_from_labelset_id")
        if labelset_id:
            django_pk = from_global_id(labelset_id)[1]
            queryset = queryset.filter(annotation_label__included_in_labelset=django_pk)

        # Filter by created_by_analysis_ids
        analysis_ids = kwargs.get("created_by_analysis_ids")
        if analysis_ids:
            analysis_id_list = analysis_ids.split(",")
            if MANUAL_ANNOTATION_SENTINEL in analysis_id_list:
                analysis_id_list = [
                    id for id in analysis_id_list if id != MANUAL_ANNOTATION_SENTINEL
                ]
                analysis_pks = [
                    int(from_global_id(value)[1]) for value in analysis_id_list
                ]
                queryset = queryset.filter(
                    Q(analysis__isnull=True) | Q(analysis_id__in=analysis_pks)
                )
            else:
                analysis_pks = [
                    int(from_global_id(value)[1]) for value in analysis_id_list
                ]
                queryset = queryset.filter(analysis_id__in=analysis_pks)

        # Filter by created_with_analyzer_id
        analyzer_ids = kwargs.get("created_with_analyzer_id")
        if analyzer_ids:
            analyzer_id_list = analyzer_ids.split(",")
            if MANUAL_ANNOTATION_SENTINEL in analyzer_id_list:
                analyzer_id_list = [
                    id for id in analyzer_id_list if id != MANUAL_ANNOTATION_SENTINEL
                ]
                analyzer_pks = [
                    int(from_global_id(id)[1])
                    for id in analyzer_id_list
                    if id != MANUAL_ANNOTATION_SENTINEL
                ]
                queryset = queryset.filter(
                    Q(analysis__isnull=True) | Q(analysis__analyzer_id__in=analyzer_pks)
                )
            elif len(analyzer_id_list) > 0:
                analyzer_pks = [int(from_global_id(id)[1]) for id in analyzer_id_list]
                queryset = queryset.filter(analysis__analyzer_id__in=analyzer_pks)

        # Filter by raw_text
        raw_text = kwargs.get("raw_text_contains")
        if raw_text:
            queryset = queryset.filter(raw_text__contains=raw_text)

        # Filter by annotation_label_id
        annotation_label_id = kwargs.get("annotation_label_id")
        if annotation_label_id:
            django_pk = from_global_id(annotation_label_id)[1]
            queryset = queryset.filter(annotation_label_id=django_pk)

        # Filter by annotation_label__text
        label_text = kwargs.get("annotation_label__text")
        if label_text:
            queryset = queryset.filter(annotation_label__text=label_text)

        label_text_contains = kwargs.get("annotation_label__text_contains")
        if label_text_contains:
            queryset = queryset.filter(
                annotation_label__text__contains=label_text_contains
            )

        # Filter by annotation_label__description
        label_description = kwargs.get("annotation_label__description_contains")
        if label_description:
            queryset = queryset.filter(
                annotation_label__description__contains=label_description
            )

        # Filter by annotation_label__label_type
        label_type = kwargs.get("annotation_label__label_type")
        if label_type:
            queryset = queryset.filter(annotation_label__label_type=label_type)

        # Filter by analysis
        if analysis_isnull is not None:
            queryset = queryset.filter(analysis__isnull=analysis_isnull)

        # Filter by corpus_action
        if corpus_action_isnull is not None:
            queryset = queryset.filter(corpus_action__isnull=corpus_action_isnull)

        # Combined agent filter: annotations created by analysis OR corpus action
        if agent_created is not None:
            agent_q = Q(analysis__isnull=False) | Q(corpus_action__isnull=False)
            if agent_created:
                queryset = queryset.filter(agent_q)
            else:
                queryset = queryset.exclude(agent_q)

        # Skip document_id and corpus_id filtering if already handled by optimizer
        if not document_id:
            # Filter by document_id
            document_id = kwargs.get("document_id")
            if document_id:
                django_pk = from_global_id(document_id)[1]
                queryset = queryset.filter(document_id=django_pk)

            # Filter by corpus_id
            corpus_id = kwargs.get("corpus_id")
            if corpus_id:
                django_pk = from_global_id(corpus_id)[1]
                queryset = queryset.filter(corpus_id=django_pk)

        # Filter by structural
        if structural is not None:
            queryset = queryset.filter(structural=structural)

        # Ordering
        order_by = kwargs.get("order_by")
        if order_by:
            queryset = queryset.order_by(order_by)
        else:
            queryset = queryset.order_by("-modified")

        return queryset

    label_type_enum = graphene.Enum.from_enum(LabelType)

    #############################################################################################
    # For some annotations, it's not clear exactly how to paginate them and, mostllikely        #
    # the total # of such annotations will be pretty minimal (specifically relationships and    #
    # doc types). The bulk_doc_annotations_in_corpus field allows you to request                #
    # full complement of annotations for a given doc in a given corpus as a list                #
    # rather than a Relay-style connection.                                                     #
    #############################################################################################

    bulk_doc_relationships_in_corpus = graphene.Field(
        graphene.List(RelationshipType),
        corpus_id=graphene.ID(required=True),
        document_id=graphene.ID(required=True),
    )

    def resolve_bulk_doc_relationships_in_corpus(
        self, info, corpus_id, document_id
    ) -> Any:
        # Get the base queryset using visible_to_user
        queryset = BaseService.filter_visible(
            Relationship, info.context.user, request=info.context
        )

        doc_django_pk = from_global_id(document_id)[1]
        corpus_django_pk = from_global_id(corpus_id)[1]

        queryset = queryset.filter(
            corpus_id=corpus_django_pk, document_id=doc_django_pk
        )  # Existing filter
        queryset = queryset.select_related(
            "relationship_label",
            "corpus",
            "document",
            "creator",
            "analyzer",  # If needed
            "analysis",  # If needed
        ).prefetch_related(
            "source_annotations",  # If RelationshipType shows source annotations
            "target_annotations",  # If RelationshipType shows target annotations
        )
        return queryset

    bulk_doc_annotations_in_corpus = graphene.Field(
        graphene.List(AnnotationType),
        corpus_id=graphene.ID(required=True),
        document_id=graphene.ID(required=False),
        for_analysis_ids=graphene.String(required=False),
        label_type=graphene.Argument(label_type_enum),
    )

    def resolve_bulk_doc_annotations_in_corpus(self, info, corpus_id, **kwargs) -> Any:

        corpus_django_pk = from_global_id(corpus_id)[1]

        # Get the base queryset using visible_to_user
        queryset = BaseService.filter_visible(
            Annotation, info.context.user, request=info.context
        ).order_by("page")

        # Now build query to stuff they want to see (filter to annotations in this corpus or with NO corpus FK, which
        # travel with document.
        q_objects = Q(corpus_id=corpus_django_pk) | Q(corpus_id__isnull=True)

        # If for_analysis_ids is passed in, only show annotations from those analyses, otherwise only show human
        # annotations.
        for_analysis_ids = kwargs.get("for_analysis_ids", None)
        if for_analysis_ids is not None and len(for_analysis_ids) > 0:
            logger.info(
                f"resolve_bulk_doc_annotations - Split ids: {for_analysis_ids.split(',')}"
            )
            analysis_pks = [
                int(from_global_id(value)[1])
                for value in list(
                    filter(lambda raw_id: len(raw_id) > 0, for_analysis_ids.split(","))
                )
            ]
            logger.info(f"resolve_bulk_doc_annotations - Analysis pks: {analysis_pks}")
            q_objects.add(Q(analysis_id__in=analysis_pks), Q.AND)
        # else:
        #     q_objects.add(Q(analysis__isnull=True), Q.AND)

        label_type = kwargs.get("label_type", None)
        if label_type is not None:
            q_objects.add(Q(annotation_label__label_type=label_type), Q.AND)

        document_id = kwargs.get("document_id", None)
        if document_id is not None:
            doc_pk = from_global_id(document_id)[1]
            q_objects.add(Q(document_id=doc_pk), Q.AND)

        logger.info(f"Filter queryset {queryset} bulk annotations: {q_objects}")

        final_queryset = queryset.filter(q_objects).order_by(
            "created", "page"
        )  # Existing filter/order
        final_queryset = final_queryset.select_related(
            "annotation_label",
            "creator",
            "document",
            "corpus",
            "analysis",
            "analysis__analyzer",
            # 'embeddings' # If needed
        )
        return final_queryset

    page_annotations = graphene.Field(
        PageAwareAnnotationType,
        current_page=graphene.Int(required=False),
        page_number_list=graphene.String(required=False),
        page_containing_annotation_with_id=graphene.ID(required=False),
        corpus_id=graphene.ID(required=False),
        document_id=graphene.ID(required=True),
        for_analysis_ids=graphene.String(required=False),
        label_type=graphene.Argument(label_type_enum),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_page_annotations(
        self, info, document_id, corpus_id=None, **kwargs
    ) -> Any:

        doc_django_pk = int(from_global_id(document_id)[1])

        # Fetch the document (consider select_related if creator/etc. are used elsewhere)
        # Using get_object_or_404 for better error handling if document not found/accessible
        # For simplicity, assuming simple get for now based on original code.
        try:
            # Add select_related if document creator/etc. needed later
            document = Document.objects.get(id=doc_django_pk)
        except Document.DoesNotExist:
            # Handle error appropriately, maybe return null or raise GraphQL error
            logger.error(f"Document with pk {doc_django_pk} not found.")
            return None  # Or raise appropriate GraphQL error

        # Get the base queryset using visible_to_user
        queryset = BaseService.filter_visible(
            Annotation, info.context.user, request=info.context
        )

        # Apply select_related EARLY to the base queryset
        queryset = queryset.select_related(
            "annotation_label",
            "creator",
            "document",  # Document already fetched, but good practice if base queryset reused
            "corpus",
            "analysis",
            "analysis__analyzer",
        )

        # Now build query filters
        q_objects = Q(document_id=doc_django_pk)
        if corpus_id is not None:
            corpus_pk = from_global_id(corpus_id)[
                1
            ]  # Get corpus_pk only if corpus_id is present
            q_objects.add(Q(corpus_id=corpus_pk), Q.AND)

        # If for_analysis_ids is passed in, only show annotations from those analyses
        for_analysis_ids = kwargs.get("for_analysis_ids", None)
        if for_analysis_ids is not None:
            analysis_pks = [
                int(from_global_id(value)[1])
                for value in list(
                    filter(lambda raw_id: len(raw_id) > 0, for_analysis_ids.split(","))
                )
            ]
            if analysis_pks:  # Only add filter if there are valid PKs
                logger.info(
                    f"resolve_page_annotations - Filtering by Analysis pks: {analysis_pks}"
                )
                q_objects.add(Q(analysis_id__in=analysis_pks), Q.AND)
            else:
                # Handle case maybe? Or assume UI prevents empty string if filter applied
                logger.warning(
                    "resolve_page_annotations - for_analysis_ids provided but resulted in empty PK list."
                )
        else:
            logger.info(
                "resolve_page_annotations - for_analysis_ids is None, filtering for analysis__isnull=True"
            )
            q_objects.add(Q(analysis__isnull=True), Q.AND)

        label_type = kwargs.get("label_type", None)
        if label_type is not None:
            logger.info(
                f"resolve_page_annotations - Filtering by label_type: {label_type}"
            )
            q_objects.add(Q(annotation_label__label_type=label_type), Q.AND)

        # Apply filters to the optimized base queryset
        # Order by page first for potential pagination logic, then created
        all_pages_annotations = queryset.filter(q_objects).order_by("page", "created")

        # --- Determine the current page ---
        page_containing_annotation_with_id = kwargs.get(
            "page_containing_annotation_with_id", None
        )
        page_number_list = kwargs.get("page_number_list", None)
        current_page = 1  # Default to page 1 (1-indexed)
        pages: list[int] = []  # Parsed page list from page_number_list (1-indexed)

        # Always parse page_number_list when provided so `pages` is available
        # for the filtering step below, regardless of which branch sets current_page.
        if page_number_list is not None:
            if re.search(r"^(?:\d+,)*\d+$", page_number_list):
                pages = [int(page) for page in page_number_list.split(",")]
            else:
                logger.warning(
                    f"Invalid format for page_number_list: {page_number_list}"
                )

        if kwargs.get("current_page", None) is not None:
            current_page = int(kwargs["current_page"])
            logger.info(
                f"resolve_page_annotations - Using provided current_page: {current_page}"
            )
        elif pages:
            current_page = pages[-1]
            logger.info(
                f"resolve_page_annotations - Using last page from page_number_list: {current_page}"
            )
        elif page_containing_annotation_with_id:
            try:
                annotation_pk = int(
                    from_global_id(page_containing_annotation_with_id)[1]
                )
                # Optimized fetch for just the page number
                annotation_page_zero_indexed = (
                    Annotation.objects.filter(pk=annotation_pk)
                    .values_list("page", flat=True)
                    .first()
                )  # Use first() to avoid DoesNotExist

                if annotation_page_zero_indexed is not None:
                    current_page = (
                        annotation_page_zero_indexed + 1
                    )  # Convert 0-indexed DB value to 1-indexed page number
                    logger.info(
                        f"resolve_page_annotations - Found page {current_page} for annotation pk {annotation_pk}"
                    )
                else:
                    logger.warning(
                        f"resolve_page_annotations - Annotation pk {annotation_pk} not found for page lookup."
                    )
                    # Keep default current_page = 1
            except (ValueError, TypeError) as e:
                logger.error(
                    f"Error parsing annotation ID {page_containing_annotation_with_id}: {e}"
                )
                # Keep default current_page = 1

        # Convert 1-indexed current page to 0-indexed for DB filtering
        current_page_zero_indexed = max(0, current_page - 1)  # Ensure it's not negative

        # --- Filter annotations for the specific page(s) ---
        if page_number_list is not None and re.search(
            r"^(?:\d+,)*\d+$", page_number_list
        ):
            # Use validated page list from earlier
            pages_zero_indexed = [max(0, page - 1) for page in pages]
            page_annotations = all_pages_annotations.filter(
                page__in=pages_zero_indexed
            )  # Order already applied
        else:
            page_annotations = all_pages_annotations.filter(
                page=current_page_zero_indexed
            )  # Order already applied

        logger.info(
            f"resolve_page_annotations - final page annotations count: {page_annotations.count()}"
        )  # Use .count() carefully if queryset is large

        pdf_page_info = PdfPageInfoType(
            page_count=document.page_count,
            current_page=current_page_zero_indexed,  # Return 0-indexed as per original logic
            has_next_page=current_page_zero_indexed < document.page_count - 1,
            has_previous_page=current_page_zero_indexed > 0,
            corpus_id=corpus_id,
            document_id=document_id,
            for_analysis_ids=for_analysis_ids,
            label_type=label_type,
        )

        return PageAwareAnnotationType(
            page_annotations=page_annotations, pdf_page_info=pdf_page_info
        )

    annotation = relay.Node.Field(AnnotationType)

    def resolve_annotation(self, info, **kwargs) -> Any:
        django_pk = from_global_id(kwargs["id"])[1]
        queryset = BaseService.filter_visible(
            Annotation, info.context.user, request=info.context
        )
        queryset = queryset.select_related(
            "annotation_label",
            "creator",
            "document",
            "corpus",
            "analysis",
            "analysis__analyzer",  # 'embeddings'
        )
        return queryset.get(id=django_pk)

    # RELATIONSHIP RESOLVERS #####################################
    relationships = DjangoFilterConnectionField(
        RelationshipType, filterset_class=RelationshipFilter
    )

    def resolve_relationships(self, info, **kwargs) -> Any:
        queryset = BaseService.filter_visible(
            Relationship, info.context.user, request=info.context
        )
        queryset = queryset.select_related(
            "relationship_label",
            "corpus",
            "document",
            "creator",
            "analyzer",
            "analysis",
        ).prefetch_related("source_annotations", "target_annotations")
        return queryset

    relationship = relay.Node.Field(RelationshipType)

    def resolve_relationship(self, info, **kwargs) -> Any:
        django_pk = from_global_id(kwargs["id"])[1]
        queryset = BaseService.filter_visible(
            Relationship, info.context.user, request=info.context
        )
        queryset = queryset.select_related(
            "relationship_label",
            "corpus",
            "document",
            "creator",
            "analyzer",
            "analysis",
        ).prefetch_related(  # Prefetch might be overkill for a single object, but harmless
            "source_annotations", "target_annotations"
        )
        return queryset.get(id=django_pk)

    # LABEL RESOLVERS #####################################

    annotation_labels = DjangoFilterConnectionField(
        AnnotationLabelType, filterset_class=LabelFilter
    )

    def resolve_annotation_labels(self, info, **kwargs) -> Any:
        return BaseService.filter_visible(
            AnnotationLabel, info.context.user, request=info.context
        )

    annotation_label = relay.Node.Field(AnnotationLabelType)

    def resolve_annotation_label(self, info, **kwargs) -> Any:
        django_pk = from_global_id(kwargs["id"])[1]
        return BaseService.filter_visible(
            AnnotationLabel, info.context.user, request=info.context
        ).get(id=django_pk)

    # LABEL SET RESOLVERS #####################################

    labelsets = DjangoFilterConnectionField(
        LabelSetType, filterset_class=LabelsetFilter
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_LIGHT"))
    def resolve_labelsets(self, info, **kwargs) -> Any:
        return BaseService.filter_visible(
            LabelSet, info.context.user, request=info.context
        )

    labelset = relay.Node.Field(LabelSetType)

    def resolve_labelset(self, info, **kwargs) -> Any:
        django_pk = from_global_id(kwargs["id"])[1]
        return BaseService.filter_visible(
            LabelSet, info.context.user, request=info.context
        ).get(id=django_pk)

    default_labelset = graphene.Field(
        LabelSetType,
        description=(
            "The install-wide default LabelSet (is_default=True), or null if "
            "none has been seeded yet or the current user cannot see it. "
            "Used by the new-corpus modal to pre-fill the label set field."
        ),
    )

    @login_required
    def resolve_default_labelset(self, info, **kwargs) -> Any:
        return (
            BaseService.filter_visible(
                LabelSet, info.context.user, request=info.context
            )
            .filter(is_default=True)
            .first()
        )

    # NOTE RESOLVERS #####################################
    notes = DjangoConnectionField(
        NoteType,
        title_contains=graphene.String(),
        content_contains=graphene.String(),
        document_id=graphene.ID(),
        annotation_id=graphene.ID(),
        order_by=graphene.String(),
    )

    @login_required
    def resolve_notes(self, info, **kwargs) -> Any:
        # Base filtering for user permissions
        queryset = BaseService.filter_visible(
            Note, info.context.user, request=info.context
        )

        # Filter by title
        title_contains = kwargs.get("title_contains")
        if title_contains:
            logger.info(f"Filtering by title containing: {title_contains}")
            queryset = queryset.filter(title__contains=title_contains)

        # Filter by content
        content_contains = kwargs.get("content_contains")
        if content_contains:
            logger.info(f"Filtering by content containing: {content_contains}")
            queryset = queryset.filter(content__contains=content_contains)

        # Filter by document_id
        document_id = kwargs.get("document_id")
        if document_id:
            logger.info(f"Filtering by document_id: {document_id}")
            django_pk = from_global_id(document_id)[1]
            queryset = queryset.filter(document_id=django_pk)

        # Filter by annotation_id
        annotation_id = kwargs.get("annotation_id")
        if annotation_id:
            logger.info(f"Filtering by annotation_id: {annotation_id}")
            django_pk = from_global_id(annotation_id)[1]
            queryset = queryset.filter(annotation_id=django_pk)

        # Ordering
        order_by = kwargs.get("order_by")
        if order_by:
            logger.info(f"Ordering by: {order_by}")
            queryset = queryset.order_by(order_by)
        else:
            logger.info("Ordering by default: -modified")
            queryset = queryset.order_by("-modified")

        logger.info(f"Final queryset: {queryset}")
        return queryset

    note = relay.Node.Field(NoteType)

    @login_required
    def resolve_note(self, info, **kwargs) -> Any:
        django_pk = from_global_id(kwargs["id"])[1]
        return BaseService.filter_visible(
            Note, info.context.user, request=info.context
        ).get(id=django_pk)

    # GEOGRAPHIC ANNOTATION AGGREGATION ########################
    # Issue #1819 — surface aggregated map pins without leaking the
    # underlying annotation rows. The query mixin holds both the corpus
    # and global flavours so frontend (#1820 / #1821) calls one schema.
    geographic_annotations_for_corpus = graphene.List(
        lambda: GeographicAnnotationPinType,
        corpus_id=graphene.ID(required=True),
        bbox=graphene.Argument(lambda: BBoxInputType),
        zoom=graphene.Float(
            description=(
                "Optional map zoom level used by the consumer to pick a label "
                "type. Not currently consumed server-side — the resolver "
                "returns every label type and lets the client decide which "
                "to render at the current zoom. ``Float`` accommodates the "
                "fractional zoom levels (e.g. 12.5) that Mapbox / MapLibre "
                "use natively."
            )
        ),
        label_types=graphene.List(
            graphene.String,
            description=(
                "Optional subset of label types to include: 'country', "
                "'state', 'city'. Defaults to all three."
            ),
        ),
        description=(
            "Aggregated geographic pins for a single corpus. Pins are "
            "deduplicated by ``(label_type, canonical_name, lat, lng)`` and "
            "ship a bounded ``sample_document_ids`` preview rather than the "
            "full annotation row set. Document visibility uses MIN(document, "
            "corpus) so private documents inside a public corpus stay hidden."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_geographic_annotations_for_corpus(
        self, info, corpus_id, bbox=None, zoom=None, label_types=None
    ) -> list:
        """Resolve corpus-scoped pins via :class:`GeographicAnnotationService`.

        ``zoom`` is accepted for forward compatibility with the map UI but
        is not consumed today; the server returns all label types and the
        frontend picks the right one for the current cluster threshold.
        """
        from opencontractserver.annotations.services import (
            BBox,
            GeographicAnnotationService,
        )
        from opencontractserver.corpuses.models import Corpus

        django_pk = from_global_id(corpus_id)[1]
        corpus = BaseService.get_or_none(
            Corpus, django_pk, info.context.user, request=info.context
        )
        # IDOR-safe: same empty response whether the corpus doesn't exist or
        # is invisible — never leaks existence.
        if corpus is None:
            return []

        # ``BBox`` raises ``ValueError`` on a degenerate ``south > north``
        # box, and ``aggregate_for_corpus`` raises on an unknown
        # ``label_types`` entry. Surface both as clean ``GraphQLError`` so
        # the client gets an actionable, sanitised message instead of an
        # unhandled-exception 500.
        try:
            bbox_obj = (
                BBox(
                    south=bbox.south,
                    west=bbox.west,
                    north=bbox.north,
                    east=bbox.east,
                )
                if bbox is not None
                else None
            )
            return GeographicAnnotationService.aggregate_for_corpus(
                user=info.context.user,
                corpus=corpus,
                bbox=bbox_obj,
                label_types=label_types,
                request=info.context,
            )
        except ValueError as exc:
            raise GraphQLError(str(exc)) from exc

    global_geographic_annotations = graphene.List(
        lambda: GeographicAnnotationPinType,
        bbox=graphene.Argument(lambda: BBoxInputType),
        zoom=graphene.Float(),
        label_types=graphene.List(graphene.String),
        description=(
            "Aggregated geographic pins across every annotation visible to "
            "the requesting user (the Discover map surface). Same shape as "
            "``geographicAnnotationsForCorpus``."
        ),
    )

    @graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_MEDIUM"))
    def resolve_global_geographic_annotations(
        self, info, bbox=None, zoom=None, label_types=None
    ) -> list:
        """Resolve global pins via :class:`GeographicAnnotationService`.

        No ``@login_required`` — the Discover page must work for anonymous
        visitors. ``Annotation.objects.visible_to_user`` enforces the
        anonymous-friendly visibility rules (public corpus + public
        document) inside the service.
        """
        from opencontractserver.annotations.services import (
            BBox,
            GeographicAnnotationService,
        )

        # Symmetric with ``resolve_geographic_annotations_for_corpus``:
        # convert ``ValueError`` from either ``BBox`` construction (degenerate
        # south > north box) or the service's label-type validation into a
        # ``GraphQLError`` rather than letting it escape as a generic 500.
        try:
            bbox_obj = (
                BBox(
                    south=bbox.south,
                    west=bbox.west,
                    north=bbox.north,
                    east=bbox.east,
                )
                if bbox is not None
                else None
            )
            return GeographicAnnotationService.aggregate_global(
                user=info.context.user,
                bbox=bbox_obj,
                label_types=label_types,
                request=info.context,
            )
        except ValueError as exc:
            raise GraphQLError(str(exc)) from exc


class BBoxInputType(graphene.InputObjectType):
    """Map bounding-box input shared by both geographic queries.

    Fields use standard map conventions: ``south <= north`` (degenerate
    ``south > north`` boxes are rejected with a ``GraphQLError``); ``west``
    may exceed ``east`` for boxes that cross the antimeridian (180°/-180°
    longitude seam) and the resolver handles the wrap-around explicitly.
    """

    south = graphene.Float(required=True)
    west = graphene.Float(required=True)
    north = graphene.Float(required=True)
    east = graphene.Float(required=True)


class GeographicAnnotationPinType(graphene.ObjectType):
    """A single aggregated geographic pin returned to the map UI.

    Mirrors :class:`GeographicPin` from the service layer one-to-one — the
    resolver projects the dataclass directly into this type via field
    resolvers below. ``label_type`` is a literal string ("country" /
    "state" / "city") rather than an enum so a future label-type expansion
    doesn't break the schema.
    """

    canonical_name = graphene.String(required=True)
    label_type = graphene.String(required=True)
    lat = graphene.Float(required=True)
    lng = graphene.Float(required=True)
    document_count = graphene.Int(required=True)
    sample_document_ids = graphene.List(graphene.ID)

    def resolve_sample_document_ids(self, info) -> list:
        """Wrap raw integer PKs as Relay global IDs.

        The service layer carries integer PKs for cheap lookup; the GraphQL
        contract is ``[ID]`` so we encode here. Done in the resolver
        rather than the service so the service stays decoupled from the
        Relay encoding scheme.
        """
        from graphql_relay import to_global_id

        return [to_global_id("DocumentType", pk) for pk in self.sample_document_ids]
