from __future__ import annotations

import contextlib
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from graphene import Argument, Boolean, DateTime, Field, InputObjectType, Int, List, NonNull, ObjectType, String
from graphene import Enum as GrapheneEnum
from opentelemetry import trace

from infrahub.core import registry
from infrahub.core.constants import DiffAction, RelationshipCardinality, RelationshipDirection
from infrahub.core.constants.database import DatabaseEdgeType
from infrahub.core.diff.diff_locker import DiffLocker
from infrahub.core.diff.model.path import BranchTrackingId, NameTrackingId, TrackingId
from infrahub.core.diff.query.filters import EnrichedDiffQueryFilters
from infrahub.core.diff.repository.repository import DiffRepository
from infrahub.core.query.diff import DiffCountChanges
from infrahub.core.timestamp import Timestamp
from infrahub.dependencies.registry import get_component_registry
from infrahub.exceptions import BranchNotFoundError, SchemaNotFoundError, ValidationError
from infrahub.graphql.enums import ConflictSelection as GraphQLConflictSelection
from infrahub.graphql.field_extractor import extract_graphql_fields

if TYPE_CHECKING:
    from datetime import datetime

    from graphql import GraphQLResolveInfo

    from infrahub.core.diff.model.path import (
        EnrichedDiffAttribute,
        EnrichedDiffConflict,
        EnrichedDiffNode,
        EnrichedDiffProperty,
        EnrichedDiffRelationship,
        EnrichedDiffRoot,
        EnrichedDiffSingleRelationship,
    )
    from infrahub.core.schema import MainSchemaTypes
    from infrahub.database import InfrahubDatabase
    from infrahub.graphql.initialization import GraphqlContext

GrapheneDiffActionEnum = GrapheneEnum.from_enum(DiffAction)
GrapheneCardinalityEnum = GrapheneEnum.from_enum(RelationshipCardinality)


@dataclass
class DiffBranchInfo:
    name: str
    branch_start_timestamp: Timestamp | None
    is_terminal: bool


@dataclass
class DiffQueryParams:
    from_time: Timestamp | None
    to_time: Timestamp | None
    tracking_id: TrackingId | None
    exclude_merged: bool


@dataclass
class ParentNodeInfo:
    node: EnrichedDiffNode
    relationship_name: str


class ConflictDetails(ObjectType):
    uuid = String(required=True)
    base_branch_action = Field(GrapheneDiffActionEnum, required=True)
    base_branch_value = String()
    base_branch_changed_at = DateTime(required=True)
    base_branch_label = String()
    diff_branch_action = Field(GrapheneDiffActionEnum, required=True)
    diff_branch_value = String()
    diff_branch_changed_at = DateTime(required=True)
    diff_branch_label = String()
    selected_branch = Field(GraphQLConflictSelection)
    resolvable = Boolean()


class DiffSummaryCounts(ObjectType):
    num_added = Int(required=True)
    num_updated = Int(required=True)
    num_removed = Int(required=True)
    num_conflicts = Int(required=True)


class DiffProperty(ObjectType):
    property_type = String(required=True)
    last_changed_at = DateTime(required=True)
    previous_value = String(required=False)
    new_value = String(required=False)
    previous_label = String(required=False)
    new_label = String(required=False)
    status = Field(GrapheneDiffActionEnum, required=True)
    path_identifier = String(required=True)
    conflict = Field(ConflictDetails, required=False)


class DiffAttribute(DiffSummaryCounts):
    name = String(required=True)
    last_changed_at = DateTime(required=True)
    status = Field(GrapheneDiffActionEnum, required=True)
    path_identifier = String(required=True)
    properties = List(NonNull(DiffProperty))
    contains_conflict = Boolean(required=True)
    conflict = Field(ConflictDetails, required=False)


class DiffSingleRelationship(DiffSummaryCounts):
    last_changed_at = DateTime(required=False)
    status = Field(GrapheneDiffActionEnum, required=True)
    peer_id = String(required=True)
    peer_label = String(required=False)
    path_identifier = String(required=True)
    contains_conflict = Boolean(required=True)
    conflict = Field(ConflictDetails, required=False)
    properties = List(NonNull(DiffProperty))


class DiffRelationship(DiffSummaryCounts):
    name = String(required=True)
    label = String(required=False)
    last_changed_at = DateTime(required=False)
    cardinality = Field(GrapheneCardinalityEnum, required=True)
    status = Field(GrapheneDiffActionEnum, required=True)
    path_identifier = String(required=True)
    elements = List(NonNull(DiffSingleRelationship), required=True)
    contains_conflict = Boolean(required=True)


class DiffNodeParent(ObjectType):
    uuid = String(required=True)
    kind = String(required=False)
    relationship_name = String(required=False)


class DiffNode(DiffSummaryCounts):
    uuid = String(required=True)
    kind = String(required=True)
    label = String(required=True)
    status = Field(GrapheneDiffActionEnum, required=True)
    path_identifier = String(required=True)
    conflict = Field(ConflictDetails, required=False)
    contains_conflict = Boolean(required=True)
    last_changed_at = DateTime(required=False)
    parent = Field(DiffNodeParent, required=False)
    attributes = List(NonNull(DiffAttribute), required=True)
    relationships = List(NonNull(DiffRelationship), required=True)


class DiffTree(DiffSummaryCounts):
    base_branch = String(required=True)
    diff_branch = String(required=True)
    from_time = DateTime(required=True)
    to_time = DateTime(required=True)
    num_untracked_base_changes = Int(required=False)
    num_untracked_diff_changes = Int(required=False)
    name = String(required=False)
    nodes = List(NonNull(DiffNode))


class DiffTreeSummary(DiffSummaryCounts):
    base_branch = String(required=True)
    diff_branch = String(required=True)
    from_time = DateTime(required=True)
    to_time = DateTime(required=True)
    num_unchanged = Int(required=True)
    num_untracked_base_changes = Int(required=False)
    num_untracked_diff_changes = Int(required=False)


class DiffTreeResolver:
    def __init__(self) -> None:
        self.source_branch_name: str | None = None

    def initialize(self, enriched_diff_root: EnrichedDiffRoot) -> None:
        self.source_branch_name = enriched_diff_root.diff_branch_name

    async def to_diff_tree(
        self, enriched_diff_root: EnrichedDiffRoot, graphql_context: GraphqlContext | None = None
    ) -> DiffTree:
        self.initialize(enriched_diff_root=enriched_diff_root)
        all_nodes = list(enriched_diff_root.nodes)
        tree_nodes = [self.to_diff_node(enriched_node=e_node, graphql_context=graphql_context) for e_node in all_nodes]
        name = None
        if enriched_diff_root.tracking_id and isinstance(enriched_diff_root.tracking_id, NameTrackingId):
            name = enriched_diff_root.tracking_id.name
        return DiffTree(
            base_branch=enriched_diff_root.base_branch_name,
            diff_branch=enriched_diff_root.diff_branch_name,
            from_time=await enriched_diff_root.from_time.to_graphql(),
            to_time=await enriched_diff_root.to_time.to_graphql(),
            name=name,
            nodes=tree_nodes,
            num_added=enriched_diff_root.num_added,
            num_updated=enriched_diff_root.num_updated,
            num_removed=enriched_diff_root.num_removed,
            num_conflicts=enriched_diff_root.num_conflicts,
        )

    def _get_parent_info(
        self, diff_node: EnrichedDiffNode, graphql_context: GraphqlContext | None = None
    ) -> ParentNodeInfo | None:
        for r in diff_node.relationships:
            for n in r.nodes:
                relationship_name: str = "undefined"

                if not graphql_context or not self.source_branch_name:
                    return ParentNodeInfo(node=n, relationship_name=relationship_name)

                node_schema: MainSchemaTypes | None = None
                parent_schema: MainSchemaTypes | None = None
                for branch_name in [self.source_branch_name, registry.default_branch]:
                    with contextlib.suppress(BranchNotFoundError, SchemaNotFoundError):
                        node_schema = graphql_context.db.schema.get(
                            name=diff_node.kind, branch=branch_name, duplicate=False
                        )
                        parent_schema = graphql_context.db.schema.get(name=n.kind, branch=branch_name, duplicate=False)
                        break

                if not node_schema or not parent_schema:
                    return ParentNodeInfo(node=n, relationship_name=relationship_name)

                rel_schema = node_schema.get_relationship(name=r.name)
                rels_parent = parent_schema.get_relationships_by_identifier(id=rel_schema.get_identifier())

                if rels_parent and len(rels_parent) == 1:
                    relationship_name = rels_parent[0].name
                elif rels_parent and len(rels_parent) > 1:
                    for rel_parent in rels_parent:
                        if (
                            rel_schema.direction == RelationshipDirection.INBOUND
                            and rel_parent.direction == RelationshipDirection.OUTBOUND
                        ) or (
                            rel_schema.direction == RelationshipDirection.OUTBOUND
                            and rel_parent.direction == RelationshipDirection.INBOUND
                        ):
                            relationship_name = rel_parent.name
                            break

                return ParentNodeInfo(node=n, relationship_name=relationship_name)
        return None

    def to_diff_node(self, enriched_node: EnrichedDiffNode, graphql_context: GraphqlContext | None = None) -> DiffNode:
        diff_attributes = [
            self.to_diff_attribute(enriched_attribute=e_attr, graphql_context=graphql_context)
            for e_attr in enriched_node.attributes
        ]
        diff_relationships = [
            self.to_diff_relationship(enriched_relationship=e_rel, graphql_context=graphql_context)
            for e_rel in enriched_node.relationships
            if e_rel.include_in_response
        ]
        conflict = None
        if enriched_node.conflict:
            conflict = self.to_diff_conflict(enriched_conflict=enriched_node.conflict, graphql_context=graphql_context)

        parent = None
        if parent_info := self._get_parent_info(diff_node=enriched_node, graphql_context=graphql_context):
            parent = DiffNodeParent(
                uuid=parent_info.node.uuid,
                kind=parent_info.node.kind,
                relationship_name=parent_info.relationship_name,
            )

        return DiffNode(
            uuid=enriched_node.uuid,
            kind=enriched_node.kind,
            label=enriched_node.label,
            status=enriched_node.action,
            parent=parent,
            last_changed_at=enriched_node.changed_at.to_datetime() if enriched_node.changed_at else None,
            path_identifier=enriched_node.path_identifier,
            attributes=diff_attributes,
            relationships=diff_relationships,
            contains_conflict=enriched_node.contains_conflict,
            conflict=conflict,
            num_added=enriched_node.num_added,
            num_updated=enriched_node.num_updated,
            num_removed=enriched_node.num_removed,
            num_conflicts=enriched_node.num_conflicts,
        )

    def to_diff_attribute(
        self, enriched_attribute: EnrichedDiffAttribute, graphql_context: GraphqlContext | None = None
    ) -> DiffAttribute:
        diff_properties = [
            self.to_diff_property(enriched_property=e_prop, graphql_context=graphql_context)
            for e_prop in enriched_attribute.properties
        ]
        conflict = None
        for diff_prop in diff_properties:
            if diff_prop.property_type == DatabaseEdgeType.HAS_VALUE.value and diff_prop.conflict:
                conflict = diff_prop.conflict
                diff_prop.conflict = None
        return DiffAttribute(
            name=enriched_attribute.name,
            last_changed_at=enriched_attribute.changed_at.to_datetime(),
            status=enriched_attribute.action,
            path_identifier=enriched_attribute.path_identifier,
            properties=diff_properties,
            contains_conflict=enriched_attribute.contains_conflict,
            conflict=conflict,
            num_added=enriched_attribute.num_added,
            num_updated=enriched_attribute.num_updated,
            num_removed=enriched_attribute.num_removed,
            num_conflicts=enriched_attribute.num_conflicts,
        )

    def to_diff_relationship(
        self, enriched_relationship: EnrichedDiffRelationship, graphql_context: GraphqlContext | None = None
    ) -> DiffRelationship:
        diff_elements = [
            self.to_diff_relationship_element(enriched_element=element, graphql_context=graphql_context)
            for element in enriched_relationship.relationships
        ]
        return DiffRelationship(
            name=enriched_relationship.name,
            label=enriched_relationship.label,
            last_changed_at=enriched_relationship.changed_at.to_datetime()
            if enriched_relationship.changed_at
            else None,
            status=enriched_relationship.action,
            cardinality=enriched_relationship.cardinality,
            path_identifier=enriched_relationship.path_identifier,
            elements=diff_elements,
            contains_conflict=enriched_relationship.contains_conflict,
            num_added=enriched_relationship.num_added,
            num_updated=enriched_relationship.num_updated,
            num_removed=enriched_relationship.num_removed,
            num_conflicts=enriched_relationship.num_conflicts,
        )

    def to_diff_relationship_element(
        self, enriched_element: EnrichedDiffSingleRelationship, graphql_context: GraphqlContext | None = None
    ) -> DiffSingleRelationship:
        diff_properties = [self.to_diff_property(e_prop) for e_prop in enriched_element.properties]
        conflict = None
        if enriched_element.conflict:
            conflict = self.to_diff_conflict(
                enriched_conflict=enriched_element.conflict, graphql_context=graphql_context
            )
        return DiffSingleRelationship(
            last_changed_at=enriched_element.changed_at.to_datetime(),
            status=enriched_element.action,
            peer_id=enriched_element.peer_id,
            peer_label=enriched_element.peer_label,
            path_identifier=enriched_element.path_identifier,
            conflict=conflict,
            properties=diff_properties,
            contains_conflict=enriched_element.contains_conflict,
            num_added=enriched_element.num_added,
            num_updated=enriched_element.num_updated,
            num_removed=enriched_element.num_removed,
            num_conflicts=enriched_element.num_conflicts,
        )

    def to_diff_property(
        self, enriched_property: EnrichedDiffProperty, graphql_context: GraphqlContext | None = None
    ) -> DiffProperty:
        conflict = None
        if enriched_property.conflict:
            conflict = self.to_diff_conflict(
                enriched_conflict=enriched_property.conflict, graphql_context=graphql_context
            )
        return DiffProperty(
            property_type=enriched_property.property_type.value,
            last_changed_at=enriched_property.changed_at.to_datetime(),
            previous_value=enriched_property.previous_value,
            new_value=enriched_property.new_value,
            previous_label=enriched_property.previous_label,
            new_label=enriched_property.new_label,
            status=enriched_property.action,
            path_identifier=enriched_property.path_identifier,
            conflict=conflict,
        )

    def to_diff_conflict(
        self,
        enriched_conflict: EnrichedDiffConflict,
        graphql_context: GraphqlContext | None = None,  # noqa: ARG002
    ) -> ConflictDetails:
        return ConflictDetails(
            uuid=enriched_conflict.uuid,
            base_branch_action=enriched_conflict.base_branch_action,
            base_branch_value=enriched_conflict.base_branch_value,
            base_branch_changed_at=enriched_conflict.base_branch_changed_at.to_datetime()
            if enriched_conflict.base_branch_changed_at
            else None,
            base_branch_label=enriched_conflict.base_branch_label,
            diff_branch_action=enriched_conflict.diff_branch_action,
            diff_branch_value=enriched_conflict.diff_branch_value,
            diff_branch_changed_at=enriched_conflict.diff_branch_changed_at.to_datetime()
            if enriched_conflict.diff_branch_changed_at
            else None,
            diff_branch_label=enriched_conflict.diff_branch_label,
            selected_branch=enriched_conflict.selected_branch.value if enriched_conflict.selected_branch else None,
            resolvable=enriched_conflict.resolvable,
        )

    async def to_graphql(
        self, fields: dict[str, dict], diff_object: Any | None
    ) -> list[dict[str, Any]] | dict[str, Any] | None:
        if diff_object is None:
            return None
        if isinstance(diff_object, list):
            list_response = True
            diff_elements = diff_object
        else:
            list_response = False
            diff_elements = [diff_object]

        response_list = []
        for diff_object_element in diff_elements:
            element_response = {}
            for field_name, sub_fields in fields.items():
                if sub_fields is None:
                    element_response[field_name] = getattr(diff_object_element, field_name, None)
                elif hasattr(diff_object_element, field_name):
                    element_response[field_name] = await self.to_graphql(
                        sub_fields, getattr(diff_object_element, field_name)
                    )
                else:
                    continue
            response_list.append(element_response)
        if list_response:
            return response_list
        return response_list[0]

    async def _add_untracked_fields(
        self,
        db: InfrahubDatabase,
        diff_response: DiffTreeSummary | DiffTree,
        from_time: Timestamp,
        base_branch_name: str | None = None,
        diff_branch_name: str | None = None,
    ) -> None:
        if not (base_branch_name or diff_branch_name):
            return
        branch_names = []
        if base_branch_name:
            branch_names.append(base_branch_name)
        if diff_branch_name:
            branch_names.append(diff_branch_name)
        query = await DiffCountChanges.init(db=db, branch_names=branch_names, diff_from=from_time, diff_to=Timestamp())
        await query.execute(db=db)
        branch_change_map = query.get_num_changes_by_branch()
        if base_branch_name:
            diff_response.num_untracked_base_changes = branch_change_map.get(base_branch_name, 0)
        if diff_branch_name:
            diff_response.num_untracked_diff_changes = branch_change_map.get(diff_branch_name, 0)

    def _get_timestamp(
        self,
        from_time: datetime | None,
        to_time: datetime | None,
        branch_start_timestamp: Timestamp | None,
        at: Timestamp,
        proposed_change_id: str | None,
    ) -> tuple[Timestamp | None, Timestamp | None]:
        if from_time:
            from_timestamp = Timestamp(from_time.isoformat())
        # don't set default time if filtering on proposed change
        elif proposed_change_id:
            from_timestamp = None
        else:
            from_timestamp = branch_start_timestamp
        if to_time:
            to_timestamp = Timestamp(to_time.isoformat())
        # don't set default time if filtering on proposed change
        elif proposed_change_id:
            to_timestamp = None
        else:
            to_timestamp = at
        return from_timestamp, to_timestamp

    async def _get_branch_info(
        self,
        db: InfrahubDatabase,
        branch: str | None,
        proposed_change_id: str | None,
    ) -> DiffBranchInfo:
        try:
            diff_branch = await registry.get_branch(db=db, branch=branch)
            return DiffBranchInfo(
                name=diff_branch.name,
                branch_start_timestamp=Timestamp(diff_branch.get_branched_from()),
                is_terminal=diff_branch.is_terminal,
            )
        except BranchNotFoundError:
            if not branch and not proposed_change_id:
                raise ValidationError("Must include the branch or proposed_change_id argument") from None
            # case for deleted branch with an input proposed_change_id
            if branch and proposed_change_id:
                return DiffBranchInfo(name=branch, branch_start_timestamp=None, is_terminal=False)
            raise

    def _get_diff_query_params(
        self,
        branch_info: DiffBranchInfo,
        at: Timestamp,
        from_time: datetime | None = None,
        to_time: datetime | None = None,
        name: str | None = None,
        proposed_change_id: str | None = None,
    ) -> DiffQueryParams:
        # For terminal branches with no explicit time filters,
        # return the latest stored diff for the branch regardless of time
        if branch_info.is_terminal and not from_time and not to_time:
            return DiffQueryParams(
                from_time=None,
                to_time=None,
                tracking_id=BranchTrackingId(branch_info.name),
                exclude_merged=False,
            )

        from_timestamp, to_timestamp = self._get_timestamp(
            from_time=from_time,
            to_time=to_time,
            branch_start_timestamp=branch_info.branch_start_timestamp,
            at=at,
            proposed_change_id=proposed_change_id,
        )
        return DiffQueryParams(
            from_time=from_timestamp,
            to_time=to_timestamp,
            tracking_id=NameTrackingId(name) if name else None,
            exclude_merged=branch_info.is_terminal or not proposed_change_id,
        )

    async def resolve(
        self,
        root: dict,  # noqa: ARG002
        info: GraphQLResolveInfo,
        branch: str | None = None,
        name: str | None = None,
        from_time: datetime | None = None,
        to_time: datetime | None = None,
        filters: dict | None = None,
        include_parents: bool = True,
        root_node_uuids: list[str] | None = None,
        limit: int | None = None,
        offset: int | None = None,
        proposed_change_id: str | None = None,
    ) -> list[dict[str, Any]] | dict[str, Any] | None:
        component_registry = get_component_registry()
        graphql_context: GraphqlContext = info.context
        base_branch = await registry.get_branch(db=graphql_context.db, branch=registry.default_branch)

        branch_info = await self._get_branch_info(
            db=graphql_context.db, branch=branch, proposed_change_id=proposed_change_id
        )
        query_params = self._get_diff_query_params(
            branch_info=branch_info,
            at=graphql_context.at or Timestamp(),
            from_time=from_time,
            to_time=to_time,
            name=name,
            proposed_change_id=proposed_change_id,
        )

        diff_repo = await component_registry.get_component(DiffRepository, db=graphql_context.db, branch=base_branch)

        # Convert filters to dict and merge root_node_uuids for compatibility
        filters_dict = dict(filters or {})
        if root_node_uuids and "ids" in filters_dict and isinstance(filters_dict["ids"], list):
            filters_dict["ids"].extend(root_node_uuids)
        elif root_node_uuids:
            filters_dict["ids"] = root_node_uuids

        # ensure any ongoing diff updates complete before we try to retrieve them
        lock_request_time = Timestamp()
        diff_locker = DiffLocker()
        async with (
            diff_locker.acquire_lock(
                target_branch_name=base_branch.name, source_branch_name=branch_info.name, is_incremental=True
            ),
            diff_locker.acquire_lock(
                target_branch_name=base_branch.name, source_branch_name=branch_info.name, is_incremental=False
            ),
        ):
            lock_acquired_time = Timestamp()
            with trace.get_tracer(__name__).start_as_current_span("diff_tree_request") as span:
                get_start_time = Timestamp()
                enriched_diffs = await diff_repo.get(
                    base_branch_name=base_branch.name,
                    diff_branch_names=[branch_info.name],
                    from_time=query_params.from_time,
                    to_time=query_params.to_time,
                    filters=EnrichedDiffQueryFilters(**filters_dict),
                    include_parents=include_parents,
                    limit=limit,
                    offset=offset,
                    tracking_id=query_params.tracking_id,
                    include_empty=True,
                    proposed_change_id=proposed_change_id,
                    exclude_merged=query_params.exclude_merged,
                )
                get_end_time = Timestamp()

                span.set_attribute("base_branch_name", base_branch.name)
                span.set_attribute("diff_branch_name", branch_info.name)
                span.set_attribute(
                    "from_time", query_params.from_time.to_string() if query_params.from_time else "null"
                )
                span.set_attribute("to_time", query_params.to_time.to_string() if query_params.to_time else "null")
                span.set_attribute("proposed_change_id", proposed_change_id or "null")
                span.set_attribute("lock_request_time", lock_request_time.to_string())
                span.set_attribute("lock_acquired_time", lock_acquired_time.to_string())
                span.set_attribute("get_start_time", get_start_time.to_string())
                span.set_attribute("get_end_time", get_end_time.to_string())
                span.set_attribute("num_diffs", str(len(enriched_diffs)))

        if not enriched_diffs:
            return None
        if len(enriched_diffs) > 0:
            # take the one with the longest duration that covers multiple branches
            enriched_diff = sorted(
                enriched_diffs,
                key=lambda d: (
                    d.base_branch_name != d.diff_branch_name,
                    d.to_time.to_datetime() - d.from_time.to_datetime(),
                ),
                reverse=True,
            )[0]
        else:
            enriched_diff = enriched_diffs[0]

        full_fields = extract_graphql_fields(info=info)
        diff_tree = await self.to_diff_tree(enriched_diff_root=enriched_diff, graphql_context=graphql_context)
        need_base_changes = "num_untracked_base_changes" in full_fields
        need_branch_changes = "num_untracked_diff_changes" in full_fields
        if need_base_changes or need_branch_changes:
            await self._add_untracked_fields(
                db=graphql_context.db,
                diff_response=diff_tree,
                from_time=enriched_diff.to_time,
                base_branch_name=base_branch.name if need_base_changes else None,
                diff_branch_name=branch_info.name if need_branch_changes else None,
            )
        return await self.to_graphql(fields=full_fields, diff_object=diff_tree)

    async def summary(
        self,
        root: dict,  # noqa: ARG002
        info: GraphQLResolveInfo,
        branch: str | None = None,
        from_time: datetime | None = None,
        to_time: datetime | None = None,
        filters: dict | None = None,
        proposed_change_id: str | None = None,
    ) -> list[dict[str, Any]] | dict[str, Any] | None:
        component_registry = get_component_registry()
        graphql_context: GraphqlContext = info.context
        base_branch = await registry.get_branch(db=graphql_context.db, branch=registry.default_branch)

        branch_info = await self._get_branch_info(
            db=graphql_context.db, branch=branch, proposed_change_id=proposed_change_id
        )
        query_params = self._get_diff_query_params(
            branch_info=branch_info,
            at=graphql_context.at or Timestamp(),
            from_time=from_time,
            to_time=to_time,
            proposed_change_id=proposed_change_id,
        )

        diff_repo = await component_registry.get_component(DiffRepository, db=graphql_context.db, branch=base_branch)

        filters_dict = dict(filters or {})

        # ensure any ongoing diff updates complete before we try to retrieve them
        lock_request_time = Timestamp()
        diff_locker = DiffLocker()
        async with (
            diff_locker.acquire_lock(
                target_branch_name=base_branch.name, source_branch_name=branch_info.name, is_incremental=True
            ),
            diff_locker.acquire_lock(
                target_branch_name=base_branch.name, source_branch_name=branch_info.name, is_incremental=False
            ),
        ):
            lock_acquired_time = Timestamp()
            with trace.get_tracer(__name__).start_as_current_span("diff_tree_summary_request") as span:
                get_start_time = Timestamp()
                summary = await diff_repo.summary(
                    base_branch_name=base_branch.name,
                    diff_branch_names=[branch_info.name],
                    from_time=query_params.from_time,
                    to_time=query_params.to_time,
                    tracking_id=query_params.tracking_id,
                    filters=filters_dict,
                    proposed_change_id=proposed_change_id,
                    exclude_merged=query_params.exclude_merged,
                )
                get_end_time = Timestamp()

                span.set_attribute("base_branch_name", base_branch.name)
                span.set_attribute("diff_branch_name", branch_info.name)
                span.set_attribute(
                    "from_time", query_params.from_time.to_string() if query_params.from_time else "null"
                )
                span.set_attribute("to_time", query_params.to_time.to_string() if query_params.to_time else "null")
                span.set_attribute("proposed_change_id", proposed_change_id or "null")
                span.set_attribute("lock_request_time", lock_request_time.to_string())
                span.set_attribute("lock_acquired_time", lock_acquired_time.to_string())
                span.set_attribute("get_start_time", get_start_time.to_string())
                span.set_attribute("get_end_time", get_end_time.to_string())
                span.set_attribute("summary_retrieved", str(summary is not None))

        if summary is None:
            return None

        diff_tree_summary = DiffTreeSummary(
            base_branch=base_branch.name,
            diff_branch=branch_info.name,
            from_time=summary.from_time.to_datetime(),
            to_time=summary.to_time.to_datetime(),
            **summary.model_dump(exclude={"from_time", "to_time"}),
        )
        full_fields = extract_graphql_fields(info=info)
        need_base_changes = "num_untracked_base_changes" in full_fields
        need_branch_changes = "num_untracked_diff_changes" in full_fields
        if need_base_changes or need_branch_changes:
            await self._add_untracked_fields(
                db=graphql_context.db,
                diff_response=diff_tree_summary,
                from_time=summary.to_time,
                base_branch_name=base_branch.name if need_base_changes else None,
                diff_branch_name=branch_info.name if need_branch_changes else None,
            )
        return diff_tree_summary


class IncExclFilterOptions(InputObjectType):
    includes = List(String)
    excludes = List(String)


class IncExclFilterStatusOptions(InputObjectType):
    includes = List(GrapheneDiffActionEnum)
    excludes = List(GrapheneDiffActionEnum)


class DiffTreeQueryFilters(InputObjectType):
    ids = List(String)
    status = IncExclFilterStatusOptions()
    kind = IncExclFilterOptions()
    namespace = IncExclFilterOptions()


DiffTreeQuery = Field(
    DiffTree,
    name=String(),
    branch=String(),
    from_time=DateTime(),
    to_time=DateTime(),
    root_node_uuids=Argument(List(String), deprecation_reason="replaced by filters"),
    include_parents=Boolean(),
    filters=DiffTreeQueryFilters(),
    limit=Int(),
    offset=Int(),
    proposed_change_id=String(),
    resolver=DiffTreeResolver().resolve,
    required=False,
)

DiffTreeSummaryQuery = Field(
    DiffTreeSummary,
    name=String(),
    branch=String(),
    from_time=DateTime(),
    to_time=DateTime(),
    filters=DiffTreeQueryFilters(),
    proposed_change_id=String(),
    resolver=DiffTreeResolver().summary,
    required=False,
)
