# generated by datamodel-codegen:
#   filename:  http://0.0.0.0:8080/execution/openapi.json
#   version:   0.33.0

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from datetime import timedelta
from enum import Enum
from typing import Annotated, Any, Final, Literal
from uuid import UUID

from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel

API_VERSION: Final[str] = "2026-04-06"


class AssetAliasReferenceAssetEventDagRun(BaseModel):
    """
    Schema for AssetAliasModel used in AssetEventDagRunReference.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    name: Annotated[str, Field(title="Name")]


class AssetProfile(BaseModel):
    """
    Profile of an asset-like object.

    Asset will have name, uri defined, with type set to 'Asset'.
    AssetNameRef will have name defined, type set to 'AssetNameRef'.
    AssetUriRef will have uri defined, type set to 'AssetUriRef'.
    AssetAlias will have name defined, type set to 'AssetAlias'.

    Note that 'type' here is distinct from 'asset_type' the user declares on an
    Asset (or subclass). This field is for distinguishing between different
    asset-related types (Asset, AssetRef, or AssetAlias).
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    name: Annotated[str | None, Field(title="Name")] = None
    uri: Annotated[str | None, Field(title="Uri")] = None
    type: Annotated[str, Field(title="Type")]


class ConnectionResponse(BaseModel):
    """
    Connection schema for responses with fields that are needed for Runtime.
    """

    conn_id: Annotated[str, Field(title="Conn Id")]
    conn_type: Annotated[str, Field(title="Conn Type")]
    host: Annotated[str | None, Field(title="Host")] = None
    schema_: Annotated[str | None, Field(alias="schema", title="Schema")] = None
    login: Annotated[str | None, Field(title="Login")] = None
    password: Annotated[str | None, Field(title="Password")] = None
    port: Annotated[int | None, Field(title="Port")] = None
    extra: Annotated[str | None, Field(title="Extra")] = None


class DagResponse(BaseModel):
    """
    Schema for DAG response.
    """

    dag_id: Annotated[str, Field(title="Dag Id")]
    is_paused: Annotated[bool, Field(title="Is Paused")]
    bundle_name: Annotated[str | None, Field(title="Bundle Name")] = None
    bundle_version: Annotated[str | None, Field(title="Bundle Version")] = None
    relative_fileloc: Annotated[str | None, Field(title="Relative Fileloc")] = None
    owners: Annotated[str | None, Field(title="Owners")] = None
    tags: Annotated[list[str], Field(title="Tags")]
    next_dagrun: Annotated[AwareDatetime | None, Field(title="Next Dagrun")] = None


class DagRunAssetReference(BaseModel):
    """
    DagRun serializer for asset responses.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    run_id: Annotated[str, Field(title="Run Id")]
    dag_id: Annotated[str, Field(title="Dag Id")]
    logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None
    start_date: Annotated[AwareDatetime, Field(title="Start Date")]
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None
    state: Annotated[str, Field(title="State")]
    data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None
    data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None
    partition_key: Annotated[str | None, Field(title="Partition Key")] = None


class DagRunState(str, Enum):
    """
    All possible states that a DagRun can be in.

    These are "shared" with TaskInstanceState in some parts of the code,
    so please ensure that their values always match the ones with the
    same name in TaskInstanceState.
    """

    QUEUED = "queued"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"


class DagRunStateResponse(BaseModel):
    """
    Schema for DAG Run State response.
    """

    state: DagRunState


class DagRunType(str, Enum):
    """
    Class with DagRun types.
    """

    BACKFILL = "backfill"
    SCHEDULED = "scheduled"
    MANUAL = "manual"
    ASSET_TRIGGERED = "asset_triggered"
    ASSET_MATERIALIZATION = "asset_materialization"


class HITLUser(BaseModel):
    """
    Schema for a Human-in-the-loop users.
    """

    id: Annotated[str, Field(title="Id")]
    name: Annotated[str, Field(title="Name")]


class InactiveAssetsResponse(BaseModel):
    """
    Response for inactive assets.
    """

    inactive_assets: Annotated[list[AssetProfile] | None, Field(title="Inactive Assets")] = None


class IntermediateTIState(str, Enum):
    """
    States that a Task Instance can be in that indicate it is not yet in a terminal or running state.
    """

    SCHEDULED = "scheduled"
    QUEUED = "queued"
    RESTARTING = "restarting"
    UP_FOR_RETRY = "up_for_retry"
    UP_FOR_RESCHEDULE = "up_for_reschedule"
    DEFERRED = "deferred"


class PrevSuccessfulDagRunResponse(BaseModel):
    """
    Schema for response with previous successful DagRun information for Task Template Context.
    """

    data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None
    data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None
    start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None


class PreviousTIResponse(BaseModel):
    """
    Schema for response with previous TaskInstance information.
    """

    task_id: Annotated[str, Field(title="Task Id")]
    dag_id: Annotated[str, Field(title="Dag Id")]
    run_id: Annotated[str, Field(title="Run Id")]
    logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None
    start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None
    state: Annotated[str | None, Field(title="State")] = None
    try_number: Annotated[int, Field(title="Try Number")]
    map_index: Annotated[int | None, Field(title="Map Index")] = -1
    duration: Annotated[float | None, Field(title="Duration")] = None


class TIDeferredStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to a deferred state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["deferred"] | None, Field(title="State")] = "deferred"
    classpath: Annotated[str, Field(title="Classpath")]
    trigger_kwargs: Annotated[dict[str, JsonValue] | str | None, Field(title="Trigger Kwargs")] = None
    trigger_timeout: Annotated[timedelta | None, Field(title="Trigger Timeout")] = None
    queue: Annotated[str | None, Field(title="Queue")] = None
    next_method: Annotated[str, Field(title="Next Method")]
    next_kwargs: Annotated[dict[str, JsonValue] | None, Field(title="Next Kwargs")] = None
    rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None


class TIEnterRunningPayload(BaseModel):
    """
    Schema for updating TaskInstance to 'RUNNING' state with minimal required fields.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["running"] | None, Field(title="State")] = "running"
    hostname: Annotated[str, Field(title="Hostname")]
    unixname: Annotated[str, Field(title="Unixname")]
    pid: Annotated[int, Field(title="Pid")]
    start_date: Annotated[AwareDatetime, Field(title="Start Date")]


class TIHeartbeatInfo(BaseModel):
    """
    Schema for TaskInstance heartbeat endpoint.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    hostname: Annotated[str, Field(title="Hostname")]
    pid: Annotated[int, Field(title="Pid")]


class TIRescheduleStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to a up_for_reschedule state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["up_for_reschedule"] | None, Field(title="State")] = "up_for_reschedule"
    reschedule_date: Annotated[AwareDatetime, Field(title="Reschedule Date")]
    end_date: Annotated[AwareDatetime, Field(title="End Date")]


class TIRetryStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to up_for_retry.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["up_for_retry"] | None, Field(title="State")] = "up_for_retry"
    end_date: Annotated[AwareDatetime, Field(title="End Date")]
    rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None


class TISkippedDownstreamTasksStatePayload(BaseModel):
    """
    Schema for updating downstream tasks to a skipped state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    tasks: Annotated[list[str | tuple[str, int]], Field(title="Tasks")]


class TISuccessStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to success state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: Annotated[Literal["success"] | None, Field(title="State")] = "success"
    end_date: Annotated[AwareDatetime, Field(title="End Date")]
    task_outlets: Annotated[list[AssetProfile] | None, Field(title="Task Outlets")] = None
    outlet_events: Annotated[list[dict[str, Any]] | None, Field(title="Outlet Events")] = None
    rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None


class TITargetStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to a target state, excluding terminal and running states.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: IntermediateTIState


class TaskBreadcrumbsResponse(BaseModel):
    """
    Response for task breadcrumbs.
    """

    breadcrumbs: Annotated[list[dict[str, Any]], Field(title="Breadcrumbs")]


class TaskInstanceState(str, Enum):
    """
    All possible states that a Task Instance can be in.

    Note that None is also allowed, so always use this in a type hint with Optional.
    """

    REMOVED = "removed"
    SCHEDULED = "scheduled"
    QUEUED = "queued"
    RUNNING = "running"
    SUCCESS = "success"
    RESTARTING = "restarting"
    FAILED = "failed"
    UP_FOR_RETRY = "up_for_retry"
    UP_FOR_RESCHEDULE = "up_for_reschedule"
    UPSTREAM_FAILED = "upstream_failed"
    SKIPPED = "skipped"
    DEFERRED = "deferred"


class TaskStatesResponse(BaseModel):
    """
    Response for task states with run_id, task and state.
    """

    task_states: Annotated[dict[str, Any], Field(title="Task States")]


class TerminalStateNonSuccess(str, Enum):
    """
    TaskInstance states that can be reported without extra information.
    """

    FAILED = "failed"
    SKIPPED = "skipped"
    REMOVED = "removed"
    UPSTREAM_FAILED = "upstream_failed"


class TriggerDAGRunPayload(BaseModel):
    """
    Schema for Trigger DAG Run API request.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None
    conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None
    reset_dag_run: Annotated[bool | None, Field(title="Reset Dag Run")] = False
    partition_key: Annotated[str | None, Field(title="Partition Key")] = None
    note: Annotated[str | None, Field(title="Note")] = None


class UpdateHITLDetailPayload(BaseModel):
    """
    Schema for writing the response part of a Human-in-the-loop detail for a specific task instance.
    """

    ti_id: Annotated[UUID, Field(title="Ti Id")]
    chosen_options: Annotated[list[str], Field(min_length=1, title="Chosen Options")]
    params_input: Annotated[dict[str, Any] | None, Field(title="Params Input")] = None


class ValidationError(BaseModel):
    loc: Annotated[list[str | int], Field(title="Location")]
    msg: Annotated[str, Field(title="Message")]
    type: Annotated[str, Field(title="Error Type")]
    input: Annotated[Any | None, Field(title="Input")] = None
    ctx: Annotated[dict[str, Any] | None, Field(title="Context")] = None


class VariablePostBody(BaseModel):
    """
    Request body schema for creating variables.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    val: Annotated[str | None, Field(title="Val")] = None
    description: Annotated[str | None, Field(title="Description")] = None


class VariableResponse(BaseModel):
    """
    Variable schema for responses with fields that are needed for Runtime.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    key: Annotated[str, Field(title="Key")]
    value: Annotated[str | None, Field(title="Value")] = None


class XComResponse(BaseModel):
    """
    XCom schema for responses with fields that are needed for Runtime.
    """

    key: Annotated[str, Field(title="Key")]
    value: JsonValue


class XComSequenceIndexResponse(RootModel[JsonValue]):
    root: Annotated[
        JsonValue,
        Field(
            description="XCom schema with minimal structure for index-based access.",
            title="XComSequenceIndexResponse",
        ),
    ]


class XComSequenceSliceResponse(RootModel[list[JsonValue]]):
    """
    XCom schema with minimal structure for slice-based access.
    """

    root: Annotated[
        list[JsonValue],
        Field(
            description="XCom schema with minimal structure for slice-based access.",
            title="XComSequenceSliceResponse",
        ),
    ]


class TaskInstance(BaseModel):
    """
    Schema for TaskInstance model with minimal required fields needed for Runtime.
    """

    id: Annotated[UUID, Field(title="Id")]
    task_id: Annotated[str, Field(title="Task Id")]
    dag_id: Annotated[str, Field(title="Dag Id")]
    run_id: Annotated[str, Field(title="Run Id")]
    try_number: Annotated[int, Field(title="Try Number")]
    dag_version_id: Annotated[UUID, Field(title="Dag Version Id")]
    map_index: Annotated[int | None, Field(title="Map Index")] = -1
    hostname: Annotated[str | None, Field(title="Hostname")] = None
    context_carrier: Annotated[dict[str, Any] | None, Field(title="Context Carrier")] = None


class BundleInfo(BaseModel):
    """
    Schema for telling task which bundle to run with.
    """

    name: Annotated[str, Field(title="Name")]
    version: Annotated[str | None, Field(title="Version")] = None


class TerminalTIState(str, Enum):
    SUCCESS = "success"
    FAILED = "failed"
    SKIPPED = "skipped"
    UPSTREAM_FAILED = "upstream_failed"
    REMOVED = "removed"


class WeightRule(str, Enum):
    DOWNSTREAM = "downstream"
    UPSTREAM = "upstream"
    ABSOLUTE = "absolute"


class TriggerRule(str, Enum):
    ALL_SUCCESS = "all_success"
    ALL_FAILED = "all_failed"
    ALL_DONE = "all_done"
    ALL_DONE_MIN_ONE_SUCCESS = "all_done_min_one_success"
    ALL_DONE_SETUP_SUCCESS = "all_done_setup_success"
    ONE_SUCCESS = "one_success"
    ONE_FAILED = "one_failed"
    ONE_DONE = "one_done"
    NONE_FAILED = "none_failed"
    NONE_SKIPPED = "none_skipped"
    ALWAYS = "always"
    NONE_FAILED_MIN_ONE_SUCCESS = "none_failed_min_one_success"
    ALL_SKIPPED = "all_skipped"


class DagAttributeTypes(str, Enum):
    OP = "operator"
    TASK_GROUP = "taskgroup"


class AssetReferenceAssetEventDagRun(BaseModel):
    """
    Schema for AssetModel used in AssetEventDagRunReference.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    name: Annotated[str, Field(title="Name")]
    uri: Annotated[str, Field(title="Uri")]
    extra: Annotated[dict[str, JsonValue], Field(title="Extra")]


class AssetResponse(BaseModel):
    """
    Asset schema for responses with fields that are needed for Runtime.
    """

    name: Annotated[str, Field(title="Name")]
    uri: Annotated[str, Field(title="Uri")]
    group: Annotated[str, Field(title="Group")]
    extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None


class HITLDetailRequest(BaseModel):
    """
    Schema for the request part of a Human-in-the-loop detail for a specific task instance.
    """

    ti_id: Annotated[UUID, Field(title="Ti Id")]
    options: Annotated[list[str], Field(min_length=1, title="Options")]
    subject: Annotated[str, Field(title="Subject")]
    body: Annotated[str | None, Field(title="Body")] = None
    defaults: Annotated[list[str] | None, Field(title="Defaults")] = None
    multiple: Annotated[bool | None, Field(title="Multiple")] = False
    params: Annotated[dict[str, Any] | None, Field(title="Params")] = None
    assigned_users: Annotated[list[HITLUser] | None, Field(title="Assigned Users")] = None


class HITLDetailResponse(BaseModel):
    """
    Schema for the response part of a Human-in-the-loop detail for a specific task instance.
    """

    response_received: Annotated[bool, Field(title="Response Received")]
    responded_by_user: HITLUser | None = None
    responded_at: Annotated[AwareDatetime | None, Field(title="Responded At")] = None
    chosen_options: Annotated[list[str] | None, Field(title="Chosen Options")] = None
    params_input: Annotated[dict[str, Any] | None, Field(title="Params Input")] = None


class HTTPValidationError(BaseModel):
    detail: Annotated[list[ValidationError] | None, Field(title="Detail")] = None


class TITerminalStatePayload(BaseModel):
    """
    Schema for updating TaskInstance to a terminal state except SUCCESS state.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    state: TerminalStateNonSuccess
    end_date: Annotated[AwareDatetime, Field(title="End Date")]
    rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None


class AssetEventDagRunReference(BaseModel):
    """
    Schema for AssetEvent model used in DagRun.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    asset: AssetReferenceAssetEventDagRun
    extra: Annotated[dict[str, JsonValue], Field(title="Extra")]
    source_task_id: Annotated[str | None, Field(title="Source Task Id")] = None
    source_dag_id: Annotated[str | None, Field(title="Source Dag Id")] = None
    source_run_id: Annotated[str | None, Field(title="Source Run Id")] = None
    source_map_index: Annotated[int | None, Field(title="Source Map Index")] = None
    source_aliases: Annotated[list[AssetAliasReferenceAssetEventDagRun], Field(title="Source Aliases")]
    timestamp: Annotated[AwareDatetime, Field(title="Timestamp")]


class AssetEventResponse(BaseModel):
    """
    Asset event schema with fields that are needed for Runtime.
    """

    id: Annotated[int, Field(title="Id")]
    timestamp: Annotated[AwareDatetime, Field(title="Timestamp")]
    extra: Annotated[dict[str, JsonValue] | None, Field(title="Extra")] = None
    asset: AssetResponse
    created_dagruns: Annotated[list[DagRunAssetReference], Field(title="Created Dagruns")]
    source_task_id: Annotated[str | None, Field(title="Source Task Id")] = None
    source_dag_id: Annotated[str | None, Field(title="Source Dag Id")] = None
    source_run_id: Annotated[str | None, Field(title="Source Run Id")] = None
    source_map_index: Annotated[int | None, Field(title="Source Map Index")] = None
    partition_key: Annotated[str | None, Field(title="Partition Key")] = None


class AssetEventsResponse(BaseModel):
    """
    Collection of AssetEventResponse.
    """

    asset_events: Annotated[list[AssetEventResponse], Field(title="Asset Events")]


class DagRun(BaseModel):
    """
    Schema for DagRun model with minimal required fields needed for Runtime.
    """

    model_config = ConfigDict(
        extra="forbid",
    )
    dag_id: Annotated[str, Field(title="Dag Id")]
    run_id: Annotated[str, Field(title="Run Id")]
    logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")] = None
    data_interval_start: Annotated[AwareDatetime | None, Field(title="Data Interval Start")] = None
    data_interval_end: Annotated[AwareDatetime | None, Field(title="Data Interval End")] = None
    run_after: Annotated[AwareDatetime, Field(title="Run After")]
    start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
    end_date: Annotated[AwareDatetime | None, Field(title="End Date")] = None
    clear_number: Annotated[int | None, Field(title="Clear Number")] = 0
    run_type: DagRunType
    state: DagRunState
    conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None
    triggering_user_name: Annotated[str | None, Field(title="Triggering User Name")] = None
    consumed_asset_events: Annotated[list[AssetEventDagRunReference], Field(title="Consumed Asset Events")]
    partition_key: Annotated[str | None, Field(title="Partition Key")] = None
    note: Annotated[str | None, Field(title="Note")] = None


class TIRunContext(BaseModel):
    """
    Response schema for TaskInstance run context.
    """

    dag_run: DagRun
    task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule Count")] = 0
    max_tries: Annotated[int, Field(title="Max Tries")]
    variables: Annotated[list[VariableResponse] | None, Field(title="Variables")] = None
    connections: Annotated[list[ConnectionResponse] | None, Field(title="Connections")] = None
    next_method: Annotated[str | None, Field(title="Next Method")] = None
    next_kwargs: Annotated[dict[str, Any] | str | None, Field(title="Next Kwargs")] = None
    xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None
    should_retry: Annotated[bool | None, Field(title="Should Retry")] = False
