from typing import Callable, Coroutine, TypedDict

from latch_o11y.o11y import dict_to_attrs, trace_function_with_span
from opentelemetry.trace import Span

from graphql_generate.support import (
    GqlContext,
    GqlSubscriptionData,
    GqlSubscriptionErrors,
    GqlWebSocketContext,
    tracer,
)
from queries_lib.schema import (
    ID,
    JSON,
    AccountAccessPoliciesOrderBy,
    AccountInfosOrderBy,
    AccountSecretsOrderBy,
    AutomationInfosOrderBy,
    AutomationRunIntervalEventsOrderBy,
    AutomationRunLdataEventsOrderBy,
    AutomationTriggerIntervalEventsOrderBy,
    AutomationTriggerLdataEventsOrderBy,
    AutomationTriggerType,
    AutomationVersionsOrderBy,
    Base64EncodedBinary,
    BasespaceApiTokensOrderBy,
    BenchlingApiTokensOrderBy,
    BenchlingImportRunsOrderBy,
    BenchlingImportsOrderBy,
    BigFloat,
    BigInt,
    Boolean,
    CatalogEventsOrderBy,
    CatalogEventType,
    CatalogExperimentColumnDefinitionsOrderBy,
    CatalogExperimentsOrderBy,
    CatalogExperimentViewsOrderBy,
    CatalogProjectsOrderBy,
    CatalogSampleColumnDataOrderBy,
    CatalogSamplesOrderBy,
    ContainerInfosOrderBy,
    CreditsOrderBy,
    Cursor,
    DatasetAgreementActivatedsOrderBy,
    DatasetAvailability,
    DatasetDownloadInfosOrderBy,
    DatasetDownloadStatus,
    DatasetInfosOrderBy,
    DatasetInfoTagsOrderBy,
    DatasetPurchasesOrderBy,
    DatasetRequestsOrderBy,
    Datetime,
    EbsSessionsOrderBy,
    ExecutionBatchInfosOrderBy,
    ExecutionBatchMembersOrderBy,
    ExecutionCreatorsOrderBy,
    ExecutionGraphNodesOrderBy,
    ExecutionInfosOrderBy,
    ExecutionNodeStatus,
    ExecutionNotificationInfosOrderBy,
    ExecutionRefLpInfosOrderBy,
    ExecutionStatus,
    Float,
    GithubPersonalAccessTokensOrderBy,
    IgvGenomesOrderBy,
    IgvTracksOrderBy,
    IncidentInfosOrderBy,
    InstanceSize,
    Int,
    IntegrationsGoogleAuthsOrderBy,
    JwtToken,
    LdataDataCreator,
    LdataEventType,
    LdataGcpMountAccessProvensOrderBy,
    LdataGcpMountProjectsOrderBy,
    LdataGoogleVerifiedBucketsOrderBy,
    LdataGoogleVerifiedProjectsOrderBy,
    LdataLinkDataOrderBy,
    LdataNodeEventChildRemovesOrderBy,
    LdataNodeEventIngressesOrderBy,
    LdataNodeEventMovesOrderBy,
    LdataNodeEventRemovesOrderBy,
    LdataNodeEventsOrderBy,
    LdataNodesOrderBy,
    LdataNodeTagsOrderBy,
    LdataNodeType,
    LdataObjectMetasOrderBy,
    LdataS3MountAccessProvensOrderBy,
    LdataS3MountConfiguratorRolesOrderBy,
    LdataS3MountRolesOrderBy,
    LdataShareInvitesOrderBy,
    LdataShareLinkNodesOrderBy,
    LdataSharePermissionsOrderBy,
    LdataTagType,
    LdataTreeEdgesOrderBy,
    LdataWfExesOrderBy,
    LdataWfExStatus,
    LpInfosOrderBy,
    NextflowNodeInfosOrderBy,
    NextflowTaskType,
    NotificationStatus,
    NotificationType,
    PermLevel,
    PlotDataSourceInfosOrderBy,
    PlotInfosOrderBy,
    PlotLayoutInfosOrderBy,
    PlotNotebookCrdtUpdatesOrderBy,
    PlotNotebookInfosOrderBy,
    PlotTransformExecutionsOrderBy,
    PlotTransformInfosOrderBy,
    PlotTransformSourceCodesOrderBy,
    PodBackupsOrderBy,
    PodBackupStatus,
    PodGroupInfosOrderBy,
    PodInfosOrderBy,
    PodSessionsOrderBy,
    PodStatus,
    PodTemplatesOrderBy,
    PodTemplateVersionsOrderBy,
    PublicKeyInfosOrderBy,
    ReportGroupsOrderBy,
    ReportsOrderBy,
    SharingReceiverSharedDataTypesOrderBy,
    SnakemakeExecutionInfosOrderBy,
    String,
    StripeSubscriptionInfosOrderBy,
    TaskExecutionInfosOrderBy,
    TaskExecutionMessagesOrderBy,
    TaskExecutionRuntimeInfosOrderBy,
    TaskExecutionStatus,
    TaskInfosOrderBy,
    TeamInfosOrderBy,
    TeamInvitesOrderBy,
    TeamMemberSharedDataTypesOrderBy,
    TeamMembersOrderBy,
    TeamRole,
    UserInfosOrderBy,
    VerifiedEmailsOrderBy,
    WorkflowFamilyInfosOrderBy,
    WorkflowGraphBranchesOrderBy,
    WorkflowGraphBranchType,
    WorkflowGraphEdgesOrderBy,
    WorkflowGraphNodesOrderBy,
    WorkflowGraphNodeType,
    WorkflowGraphSubWfsOrderBy,
    WorkflowInfosOrderBy,
    WorkflowPersonalizationsOrderBy,
    WorkflowPinsOrderBy,
    WorkflowPreviewsOrderBy,
    WorkflowSubscriptionsOrderBy,
    WorkflowTagsOrderBy,
)


# >>> NodeInfo
class NodeInfoFragment_LdataObjectMeta(TypedDict):
    """
    Part of NodeInfoFragment_LdataObjectMeta
    """
    contentSize: str | None
    modifyTime: str | None
    accessTime: str | None
    birthTime: str | None

class NodeInfoFragment(TypedDict):
    """
    Part of NodeInfoFragment
    """
    id: str
    name: str
    type: LdataNodeType
    removed: bool | None
    ldataObjectMeta: NodeInfoFragment_LdataObjectMeta | None
    """
    Reads a single `LdataObjectMeta` that is related to this `LdataNode`.
    """

query_str_node_info_fragment = "fragment NodeInfo on LdataNode{id name type removed ldataObjectMeta{contentSize modifyTime accessTime birthTime}}"

# >>> LDataSubtreeByPath
class LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes_Child_FinalLinkTarget(NodeInfoFragment, TypedDict):
    """
    Part of LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes_Child_FinalLinkTarget
    """
    ...

class LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes_Child(NodeInfoFragment, TypedDict):
    """
    Part of LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes_Child
    """
    finalLinkTarget: LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes_Child_FinalLinkTarget | None

class LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes(TypedDict):
    """
    Part of LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes
    """
    child: LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes_Child
    """
    Reads a single `LdataNode` that is related to this `LdataTreeEdge`.
    """

class LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges(TypedDict):
    """
    Part of LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges
    """
    nodes: list[LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges_Nodes]
    """
    A list of `LdataTreeEdge` objects.
    """

class LDataSubtreeByPathQueryResult_LdataResolvePathData(TypedDict):
    """
    Part of LDataSubtreeByPathQueryResult_LdataResolvePathData
    """
    id: str
    childLdataTreeEdges: LDataSubtreeByPathQueryResult_LdataResolvePathData_ChildLdataTreeEdges
    """
    Reads and enables pagination through a set of `LdataTreeEdge`.
    """

class LDataSubtreeByPathQueryResult(TypedDict):
    """
    Part of LDataSubtreeByPathQueryResult
    """
    ldataResolvePathData: LDataSubtreeByPathQueryResult_LdataResolvePathData | None

class LDataSubtreeByPathQueryVariables(TypedDict):
    path: str

query_str_l_data_subtree_by_path_query = query_str_node_info_fragment + "query LDataSubtreeByPath($path:String!){ldataResolvePathData(argPath:$path){id childLdataTreeEdges{nodes{child{...NodeInfo finalLinkTarget{...NodeInfo}}}}}}"

@trace_function_with_span(tracer)
async def query_l_data_subtree_by_path_query(span: Span, ctx: GqlContext, variables: LDataSubtreeByPathQueryVariables):
    span.set_attributes({**dict_to_attrs(variables, "variables")})
    return await ctx.query(query_str=query_str_l_data_subtree_by_path_query, variables=variables, result_type=LDataSubtreeByPathQueryResult)

# >>> Unlink
class UnlinkMutationResult_LdataUnlink(TypedDict):
    """
    Part of UnlinkMutationResult_LdataUnlink
    """
    clientMutationId: str | None
    """
    The exact same `clientMutationId` that was provided in the mutation input,
    unchanged and unused. May be used by a client to track mutations.
    """

class UnlinkMutationResult(TypedDict):
    """
    Part of UnlinkMutationResult
    """
    ldataUnlink: UnlinkMutationResult_LdataUnlink | None

class UnlinkMutationVariables(TypedDict):
    path: str

query_str_unlink_mutation = "mutation Unlink($path:String!){ldataUnlink(input:{argPath: $path}){clientMutationId}}"

@trace_function_with_span(tracer)
async def query_unlink_mutation(span: Span, ctx: GqlContext, variables: UnlinkMutationVariables):
    span.set_attributes({**dict_to_attrs(variables, "variables")})
    return await ctx.query(query_str=query_str_unlink_mutation, variables=variables, result_type=UnlinkMutationResult)

# >>> Mkdir
class MkdirMutationResult_LdataMkdirNoParents(TypedDict):
    """
    Part of MkdirMutationResult_LdataMkdirNoParents
    """
    clientMutationId: str | None
    """
    The exact same `clientMutationId` that was provided in the mutation input,
    unchanged and unused. May be used by a client to track mutations.
    """

class MkdirMutationResult(TypedDict):
    """
    Part of MkdirMutationResult
    """
    ldataMkdirNoParents: MkdirMutationResult_LdataMkdirNoParents | None

class MkdirMutationVariables(TypedDict):
    path: str

query_str_mkdir_mutation = "mutation Mkdir($path:String!){ldataMkdirNoParents(input:{argPath: $path}){clientMutationId}}"

@trace_function_with_span(tracer)
async def query_mkdir_mutation(span: Span, ctx: GqlContext, variables: MkdirMutationVariables):
    span.set_attributes({**dict_to_attrs(variables, "variables")})
    return await ctx.query(query_str=query_str_mkdir_mutation, variables=variables, result_type=MkdirMutationResult)

# >>> Rmdir
class RmdirMutationResult_LdataRmdir(TypedDict):
    """
    Part of RmdirMutationResult_LdataRmdir
    """
    clientMutationId: str | None
    """
    The exact same `clientMutationId` that was provided in the mutation input,
    unchanged and unused. May be used by a client to track mutations.
    """

class RmdirMutationResult(TypedDict):
    """
    Part of RmdirMutationResult
    """
    ldataRmdir: RmdirMutationResult_LdataRmdir | None

class RmdirMutationVariables(TypedDict):
    path: str

query_str_rmdir_mutation = "mutation Rmdir($path:String!){ldataRmdir(input:{argPath: $path}){clientMutationId}}"

@trace_function_with_span(tracer)
async def query_rmdir_mutation(span: Span, ctx: GqlContext, variables: RmdirMutationVariables):
    span.set_attributes({**dict_to_attrs(variables, "variables")})
    return await ctx.query(query_str=query_str_rmdir_mutation, variables=variables, result_type=RmdirMutationResult)

# >>> LDataRenameFuse
class LDataRenameFuseMutationResult_LdataRenameFuse(TypedDict):
    """
    Part of LDataRenameFuseMutationResult_LdataRenameFuse
    """
    clientMutationId: str | None
    """
    The exact same `clientMutationId` that was provided in the mutation input,
    unchanged and unused. May be used by a client to track mutations.
    """

class LDataRenameFuseMutationResult(TypedDict):
    """
    Part of LDataRenameFuseMutationResult
    """
    ldataRenameFuse: LDataRenameFuseMutationResult_LdataRenameFuse | None

class LDataRenameFuseMutationVariables(TypedDict):
    srcPath: str
    destPath: str

query_str_l_data_rename_fuse_mutation = "mutation LDataRenameFuse($srcPath:String!,$destPath:String!){ldataRenameFuse(input:{argSrcPath: $srcPath,argDestPath: $destPath}){clientMutationId}}"

@trace_function_with_span(tracer)
async def query_l_data_rename_fuse_mutation(span: Span, ctx: GqlContext, variables: LDataRenameFuseMutationVariables):
    span.set_attributes({**dict_to_attrs(variables, "variables")})
    return await ctx.query(query_str=query_str_l_data_rename_fuse_mutation, variables=variables, result_type=LDataRenameFuseMutationResult)

# >>> BasicInfo
class BasicInfoQueryResult_AccountInfoCurrent(TypedDict):
    """
    Part of BasicInfoQueryResult_AccountInfoCurrent
    """
    id: str

class BasicInfoQueryResult(TypedDict):
    """
    Part of BasicInfoQueryResult
    """
    accountInfoCurrent: BasicInfoQueryResult_AccountInfoCurrent | None

query_str_basic_info_query = "query BasicInfo{accountInfoCurrent{id}}"

@trace_function_with_span(tracer)
async def query_basic_info_query(span: Span, ctx: GqlContext):
    return await ctx.query(query_str=query_str_basic_info_query, result_type=BasicInfoQueryResult)

# >>> LatchData
class LatchDataSubscriptionResult_ConsoleLdataNode(TypedDict):
    """
    Part of LatchDataSubscriptionResult_ConsoleLdataNode
    """
    table: str
    rowIds: list[str] | None

class LatchDataSubscriptionResult(TypedDict):
    """
    Part of LatchDataSubscriptionResult
    """
    consoleLdataNode: LatchDataSubscriptionResult_ConsoleLdataNode | None

class LatchDataSubscriptionVariables(TypedDict):
    workspaceId: str

query_str_latch_data_subscription = "subscription LatchData($workspaceId:String!){consoleLdataNode(workspaceId:$workspaceId){table rowIds}}"

@trace_function_with_span(tracer)
async def subscribe_latch_data_subscription(span: Span, ctx: GqlWebSocketContext, *, callback: Callable[[GqlSubscriptionData[LatchDataSubscriptionResult] | GqlSubscriptionErrors], Coroutine[object, object, object]], operation_id: str, variables: LatchDataSubscriptionVariables):
    span.set_attributes({"operation_id": operation_id, "callback": callback.__name__, **dict_to_attrs(variables, "variables")})
    return await ctx.subscribe(query_str=query_str_latch_data_subscription, operation_id=operation_id, callback=callback, variables=variables, result_type=LatchDataSubscriptionResult)

