# 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

import contextlib
import itertools
import json
from collections import defaultdict
from collections.abc import Iterator
from typing import TYPE_CHECKING, Annotated, Any, cast
from uuid import UUID

import attrs
import structlog
from cadwyn import VersionedAPIRouter
from fastapi import Body, HTTPException, Query, Response, Security, status
from opentelemetry import trace
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import JsonValue
from sqlalchemy import and_, func, or_, tuple_, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
from sqlalchemy.orm import joinedload
from sqlalchemy.sql import select
from structlog.contextvars import bind_contextvars

from airflow._shared.observability.traces import override_ids
from airflow._shared.state import TaskScope
from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTGenerator
from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
from airflow.api_fastapi.common.db.common import SessionDep
from airflow.api_fastapi.common.types import UtcDateTime
from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
    InactiveAssetsResponse,
    PreviousTIResponse,
    PrevSuccessfulDagRunResponse,
    TaskBreadcrumbsResponse,
    TaskStatesResponse,
    TIAwaitingInputStatePayload,
    TIDeferredStatePayload,
    TIEnterRunningPayload,
    TIHeartbeatInfo,
    TIRescheduleStatePayload,
    TIRetryStatePayload,
    TIRunContext,
    TISkippedDownstreamTasksStatePayload,
    TIStateUpdate,
    TISuccessStatePayload,
    TITerminalStatePayload,
)
from airflow.api_fastapi.execution_api.datamodels.token import TIToken
from airflow.api_fastapi.execution_api.deps import DepContainer
from airflow.api_fastapi.execution_api.security import (
    CurrentTIToken,
    ExecutionAPIRoute,
    get_team_name_for_ti,
    require_auth,
)
from airflow.configuration import conf
from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound
from airflow.models.asset import AssetActive
from airflow.models.base import ID_LEN
from airflow.models.dag import DagModel
from airflow.models.dagrun import DagRun as DR
from airflow.models.hitl import HITLDetail
from airflow.models.log import Log
from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH
from airflow.models.taskreschedule import TaskReschedule
from airflow.models.trigger import Trigger, handle_event_submit
from airflow.models.xcom import XComModel
from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey
from airflow.state import get_state_backend
from airflow.triggers.base import TriggerEvent
from airflow.utils.sqlalchemy import get_dialect_name
from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState

if TYPE_CHECKING:
    from sqlalchemy.sql.dml import Update

router = VersionedAPIRouter()

ti_id_router = VersionedAPIRouter(
    route_class=ExecutionAPIRoute,
    dependencies=[
        Security(require_auth, scopes=["ti:self"]),
    ],
)


log = structlog.get_logger(__name__)
tracer = trace.get_tracer(__name__)


@ti_id_router.patch(
    "/{task_instance_id}/run",
    status_code=status.HTTP_200_OK,
    dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])],
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance not found"),
            (status.HTTP_409_CONFLICT, "The TI is already in the requested state"),
            (HTTP_422_UNPROCESSABLE_CONTENT, "Invalid payload for the state transition"),
        ]
    ),
    response_model_exclude_unset=True,
)
def ti_run(
    task_instance_id: UUID,
    ti_run_payload: Annotated[TIEnterRunningPayload, Body()],
    response: Response,
    session: SessionDep,
    dag_bag: DagBagDep,
    services=DepContainer,
    token: TIToken = CurrentTIToken,
) -> TIRunContext:
    """
    Run a TaskInstance.

    This endpoint is used to start a TaskInstance that is in the QUEUED state.
    """
    bind_contextvars(ti_id=str(task_instance_id))
    log.debug(
        "Starting task instance run",
        hostname=ti_run_payload.hostname,
        unixname=ti_run_payload.unixname,
        pid=ti_run_payload.pid,
    )

    from sqlalchemy.sql import column
    from sqlalchemy.types import JSON

    old = (
        select(
            TI.state,
            TI.dag_id,
            TI.run_id,
            TI.task_id,
            TI.map_index,
            TI.try_number,
            TI.max_tries,
            TI.start_date,
            TI.next_method,
            TI.hostname,
            TI.unixname,
            TI.pid,
            # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the
            # client
            column("next_kwargs", JSON),
            DR.logical_date,
            DagModel.owners,
        )
        .select_from(TI)
        .join(DR, and_(TI.dag_id == DR.dag_id, TI.run_id == DR.run_id))
        .join(DagModel, TI.dag_id == DagModel.dag_id)
        .where(TI.id == task_instance_id)
        .with_for_update(of=TI)
    )
    try:
        ti = session.execute(old).one()
        log.debug("Retrieved task instance details", state=ti.state, dag_id=ti.dag_id, task_id=ti.task_id)
    except NoResultFound:
        log.error("Task Instance not found")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={
                "reason": "not_found",
                "message": "Task Instance not found",
            },
        )

    # We exclude_unset to avoid updating fields that are not set in the payload
    data = ti_run_payload.model_dump(exclude_unset=True)

    # don't update start date when resuming from deferral
    if ti.next_kwargs:
        data.pop("start_date")
        log.debug("Removed start_date from update as task is resuming from deferral")

    query = update(TI).where(TI.id == task_instance_id).values(data)

    previous_state = ti.state

    # If we are already running, but this is a duplicate request from the same client return the same OK
    # -- it's possible there was a network glitch and they never got the response
    if previous_state == TaskInstanceState.RUNNING and (ti.hostname, ti.unixname, ti.pid) == (
        ti_run_payload.hostname,
        ti_run_payload.unixname,
        ti_run_payload.pid,
    ):
        log.info("Duplicate start request received", hostname=ti_run_payload.hostname)
    elif previous_state not in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING):
        log.warning(
            "Cannot start Task Instance in invalid state",
            previous_state=previous_state,
        )

        # TODO: Pass a RFC 9457 compliant error message in "detail" field
        # https://datatracker.ietf.org/doc/html/rfc9457
        # to provide more information about the error
        # FastAPI will automatically convert this to a JSON response
        # This might be added in FastAPI in https://github.com/fastapi/fastapi/issues/10370
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail={
                "reason": "invalid_state",
                "message": "TI was not in a state where it could be marked as running",
                "previous_state": previous_state,
            },
        )
    else:
        log.info("Task started", previous_state=previous_state, hostname=ti_run_payload.hostname)
        session.add(
            Log(
                event=TaskInstanceState.RUNNING.value,
                task_id=ti.task_id,
                dag_id=ti.dag_id,
                run_id=ti.run_id,
                map_index=ti.map_index,
                try_number=ti.try_number,
                logical_date=ti.logical_date,
                owner=ti.owners,
                extra=json.dumps({"host_name": ti_run_payload.hostname}) if ti_run_payload.hostname else None,
            )
        )
    # Ensure there is no end date set and clear retry policy overrides from the previous attempt.
    query = query.values(
        end_date=None,
        hostname=ti_run_payload.hostname,
        unixname=ti_run_payload.unixname,
        pid=ti_run_payload.pid,
        state=TaskInstanceState.RUNNING,
        last_heartbeat_at=timezone.utcnow(),
        retry_delay_override=None,
        retry_reason=None,
    )

    try:
        result = session.execute(query)
        log.info("Task instance state updated", rows_affected=getattr(result, "rowcount", 0))

        dr = (
            session.scalars(
                select(DR)
                .filter_by(dag_id=ti.dag_id, run_id=ti.run_id)
                .options(joinedload(DR.consumed_asset_events))
            )
            .unique()
            .one_or_none()
        )

        if not dr:
            log.error("DagRun not found", dag_id=ti.dag_id, run_id=ti.run_id)
            raise ValueError(f"DagRun with dag_id={ti.dag_id} and run_id={ti.run_id} not found.")

        # Send the keys to the SDK so that the client requests to clear those XComs from the server.
        # The reason we cannot do this here in the server is because we need to issue a purge on custom XCom backends
        # too. With the current assumption, the workers ONLY have access to the custom XCom backends directly and they
        # can issue the purge.

        # However, do not clear it for deferral
        xcom_keys = []
        if not ti.next_method:
            map_index = None if ti.map_index < 0 else ti.map_index
            xcom_query = select(XComModel.key).where(
                XComModel.dag_id == ti.dag_id,
                XComModel.task_id == ti.task_id,
                XComModel.run_id == ti.run_id,
            )
            if map_index is not None:
                xcom_query = xcom_query.where(XComModel.map_index == map_index)

            xcom_keys = list(session.scalars(xcom_query))
        task_reschedule_count = (
            session.scalar(
                select(func.count(TaskReschedule.id)).where(TaskReschedule.ti_id == task_instance_id)
            )
            or 0
        )

        dr.team_name = get_team_name_for_ti(task_instance_id, session)

        context = TIRunContext(
            dag_run=dr,
            task_reschedule_count=task_reschedule_count,
            max_tries=ti.max_tries,
            # TODO: Add variables and connections that are needed (and has perms) for the task
            variables=[],
            connections=[],
            xcom_keys_to_clear=xcom_keys,
            should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
        )

        # Only set if they are non-null
        if ti.next_method:
            context.next_method = ti.next_method
            context.next_kwargs = ti.next_kwargs
            context.start_date = ti.start_date
    except DataError:
        # Let the app-level DataErrorHandler return a 422 (not the opaque 500 below).
        raise
    except SQLAlchemyError:
        log.exception("Error marking Task Instance state as running")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Database error occurred"
        )

    # JWTReissueMiddleware also writes Refreshed-API-Token but skips workload tokens, so we set it here for the workload→execution swap.
    if token.claims.scope == "workload":
        generator: JWTGenerator = services.get(JWTGenerator)
        execution_token = generator.generate(extras={"sub": str(task_instance_id), "scope": "execution"})
        response.headers["Refreshed-API-Token"] = execution_token

    return context


@ti_id_router.patch(
    "/{task_instance_id}/state",
    status_code=status.HTTP_204_NO_CONTENT,
    responses={
        status.HTTP_200_OK: {"description": "The TI was already in the requested state"},
        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
        status.HTTP_409_CONFLICT: {"description": "The TI is not in a valid state for this transition"},
        HTTP_422_UNPROCESSABLE_CONTENT: {"description": "Invalid payload for the state transition"},
    },
)
def ti_update_state(
    task_instance_id: UUID,
    ti_patch_payload: Annotated[TIStateUpdate, Body()],
    session: SessionDep,
    dag_bag: DagBagDep,
):
    """
    Update the state of a TaskInstance.

    Not all state transitions are valid, and transitioning to some states requires extra information to be
    passed along. (Check out the datamodels for details, the rendered docs might not reflect this accurately)
    """
    bind_contextvars(ti_id=str(task_instance_id))
    log.debug("Updating task instance state", new_state=ti_patch_payload.state)

    old = (
        select(
            TI.state,
            TI.try_number,
            TI.max_tries,
            TI.dag_id,
            TI.task_id,
            TI.run_id,
            TI.map_index,
            TI.hostname,
            DR.logical_date,
            DagModel.owners,
        )
        .select_from(TI)
        .join(DR, and_(TI.dag_id == DR.dag_id, TI.run_id == DR.run_id))
        .join(DagModel, TI.dag_id == DagModel.dag_id)
        .where(TI.id == task_instance_id)
        .with_for_update(of=TI)
    )
    try:
        (
            previous_state,
            try_number,
            max_tries,
            dag_id,
            task_id,
            run_id,
            map_index,
            hostname,
            logical_date,
            owners,
        ) = session.execute(old).one()
        log.debug(
            "Retrieved current task instance state",
            previous_state=previous_state,
            try_number=try_number,
            max_tries=max_tries,
        )
    except NoResultFound:
        log.error("Task Instance not found")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={
                "reason": "not_found",
                "message": "Task Instance not found",
            },
        )

    # TIStateUpdate can include terminal and intermediate states. This idempotency check handles
    # duplicate updates when the requested state is already persisted (for example SUCCESS ->
    # SUCCESS or DEFERRED -> DEFERRED), including duplicates that would not pass the RUNNING
    # transition check below.
    if ti_patch_payload.state.value == previous_state:
        log.info(
            "Duplicate state update request received; state already set",
            requested_state=ti_patch_payload.state.value,
            previous_state=previous_state,
        )
        return Response(status_code=status.HTTP_200_OK)

    if previous_state != TaskInstanceState.RUNNING:
        log.warning(
            "Cannot update Task Instance in invalid state",
            previous_state=previous_state,
        )
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail={
                "reason": "invalid_state",
                "message": "TI was not in the running state so it cannot be updated",
                "previous_state": previous_state,
            },
        )

    # Validate outlet event partition keys early, before entering the catch-all
    # except block that would otherwise swallow the HTTPException and mark the TI failed.
    if isinstance(ti_patch_payload, TISuccessStatePayload):
        try:
            _validate_outlet_event_partition_keys(ti_patch_payload.outlet_events)
        except InvalidPartitionKeyError as e:
            raise HTTPException(
                status_code=HTTP_422_UNPROCESSABLE_CONTENT,
                detail={"reason": "invalid_partition_key", "message": str(e)},
            ) from e

    # We exclude_unset to avoid updating fields that are not set in the payload
    data = ti_patch_payload.model_dump(
        exclude={"task_outlets", "outlet_events", "retry_delay_seconds", "retry_reason"},
        exclude_unset=True,
    )
    if "rendered_map_index" in data:
        data["_rendered_map_index"] = data.pop("rendered_map_index")
    query = update(TI).where(TI.id == task_instance_id).values(data)

    try:
        query, updated_state = _create_ti_state_update_query_and_update_state(
            ti_patch_payload=ti_patch_payload,
            task_instance_id=task_instance_id,
            session=session,
            query=query,
            dag_id=dag_id,
            dag_bag=dag_bag,
        )
    except DataError:
        # Let DataErrorHandler return a 422 instead of silently marking the TI FAILED below.
        raise
    except Exception:
        # Set a task to failed in case any unexpected exception happened during task state update
        log.exception(
            "Error updating Task Instance state. Setting the task to failed.",
            payload=ti_patch_payload,
        )
        ti = session.get(TI, task_instance_id, with_for_update={"of": TI})
        if session.bind is not None:
            query = TI.duration_expression_update(timezone.utcnow(), query, session.bind)
        query = query.values(state=(updated_state := TaskInstanceState.FAILED))
        if ti is not None:
            _handle_fail_fast_for_dag(ti=ti, dag_id=dag_id, session=session, dag_bag=dag_bag)

    # TODO: Replace this with FastAPI's Custom Exception handling:
    # https://fastapi.tiangolo.com/tutorial/handling-errors/#install-custom-exception-handlers
    try:
        result = session.execute(query)
        log.info(
            "Task instance state updated",
            new_state=updated_state,
            rows_affected=getattr(result, "rowcount", 0),
        )
        session.add(
            Log(
                event=updated_state.value,
                task_id=task_id,
                dag_id=dag_id,
                run_id=run_id,
                map_index=map_index,
                try_number=try_number,
                logical_date=logical_date,
                owner=owners,
                extra=json.dumps({"host_name": hostname}) if hostname else None,
            )
        )
    except DataError:
        # Let DataErrorHandler return a 422 (not the opaque 500 below).
        raise
    except SQLAlchemyError as e:
        log.error("Error updating Task Instance state", error=str(e))
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Database error occurred"
        )

    if updated_state == TaskInstanceState.SUCCESS:
        if conf.getboolean("state_store", "clear_on_success"):
            scope = TaskScope(
                dag_id=dag_id,
                run_id=run_id,
                task_id=task_id,
                map_index=map_index if map_index is not None else -1,
            )
            try:
                get_state_backend().clear(scope, session=session)
                log.info(
                    "Cleared task state on success",
                    dag_id=dag_id,
                    run_id=run_id,
                    task_id=task_id,
                    map_index=map_index,
                )
            except Exception:
                log.warning(
                    "Failed to clear task state on success",
                    dag_id=dag_id,
                    run_id=run_id,
                    task_id=task_id,
                )


def _emit_task_span(ti, state):
    # just to be safe
    if not ti.dag_run:
        return
    if not isinstance(ti.dag_run.context_carrier, dict):
        return
    if not isinstance(ti.context_carrier, dict):
        return
    dr_ctx = TraceContextTextMapPropagator().extract(ti.dag_run.context_carrier)

    # Skip if the run was head-sampled out, so every span in the run agrees with the
    # carrier's decision. A parent-based sampler would already drop this child span,
    # but the explicit check also covers non-parent-based samplers (which ignore the
    # parent and would re-sample it in) and short-circuits before building the span.
    # An invalid/empty carrier (legacy/NULL) recorded no decision, so it falls through
    # and still emits — preserving prior behavior.
    dr_span_context = trace.get_current_span(context=dr_ctx).get_span_context()
    if dr_span_context.is_valid and not dr_span_context.trace_flags.sampled:
        return

    ti_ctx = TraceContextTextMapPropagator().extract(ti.context_carrier)
    ti_span = trace.get_current_span(context=ti_ctx)
    span_context = ti_span.get_span_context()
    start_time_candidates = (x for x in (ti.queued_dttm, ti.start_date, timezone.utcnow()) if x)
    name = f"task_run.{ti.task_id}"
    if ti.map_index >= 0:
        name += f"[{ti.map_index}]"
    with override_ids(span_context.trace_id, span_context.span_id):
        span = tracer.start_span(
            name=name,
            start_time=int(min(start_time_candidates).timestamp() * 1e9),
            context=dr_ctx,
        )

        span.set_attributes(
            {
                "airflow.dag_id": ti.dag_id,
                "airflow.task_id": ti.task_id,
                "airflow.dag_run.run_id": ti.run_id,
                "airflow.task_instance.try_number": ti.try_number,
                "airflow.task_instance.map_index": ti.map_index if ti.map_index is not None else -1,
                "airflow.task_instance.state": state,
                "airflow.task_instance.id": str(ti.id),
            }
        )
        status_code = StatusCode.OK if state == TaskInstanceState.SUCCESS else StatusCode.ERROR
        span.set_status(status_code)
        span.end()


def _handle_fail_fast_for_dag(ti: TI, dag_id: str, session: SessionDep, dag_bag: DagBagDep) -> None:
    dr = ti.dag_run

    # Check fail_fast from DagModel (simple column lookup) - early exit if False
    # This avoids loading 5-50 MB SerializedDAG in 99% of cases
    fail_fast = session.scalar(select(DagModel.fail_fast).where(DagModel.dag_id == dag_id))
    if not fail_fast:
        return

    # Only load SerializedDAG when fail_fast=True (rare case ~1%)
    ser_dag = dag_bag.get_dag_for_run(dag_run=dr, session=session)
    if ser_dag:
        task_dict = getattr(ser_dag, "task_dict")
        task_teardown_map = {k: v.is_teardown for k, v in task_dict.items()}
        _stop_remaining_tasks(task_instance=ti, task_teardown_map=task_teardown_map, session=session)


def _validate_outlet_event_partition_keys(outlet_events: list[dict[str, Any]]) -> None:
    """
    Validate partition_key values embedded in outlet events.

    Raises ``InvalidPartitionKeyError`` (which the caller translates to HTTP 422)
    if any per-emission partition key is empty/whitespace-only or exceeds the
    ``ID_LEN`` column width used in the metadata database.
    """
    for event in outlet_events:
        if (pk := event.get("partition_key")) is None:
            continue
        if not pk.strip():
            raise InvalidPartitionKeyError(
                f"partition_key in outlet event must not be empty or whitespace-only; got {pk!r}."
            )
        if len(pk) > ID_LEN:
            raise InvalidPartitionKeyError(
                f"partition_key in outlet event must be at most {ID_LEN} characters; got {len(pk)}."
            )


def _create_ti_state_update_query_and_update_state(
    *,
    ti_patch_payload: TIStateUpdate,
    task_instance_id: UUID,
    query: Update,
    session: SessionDep,
    dag_bag: DagBagDep,
    dag_id: str,
) -> tuple[Update, TaskInstanceState]:
    if isinstance(ti_patch_payload, (TITerminalStatePayload, TIRetryStatePayload, TISuccessStatePayload)):
        ti = session.get(TI, task_instance_id, with_for_update={"of": TI})
        updated_state = TaskInstanceState(ti_patch_payload.state.value)
        if session.bind is not None:
            query = TI.duration_expression_update(ti_patch_payload.end_date, query, session.bind)
        query = query.values(state=updated_state, next_method=None, next_kwargs=None)

        if updated_state == TaskInstanceState.FAILED:
            # This is the only case needs extra handling for TITerminalStatePayload
            if ti is not None:
                _handle_fail_fast_for_dag(ti=ti, dag_id=dag_id, session=session, dag_bag=dag_bag)
        elif isinstance(ti_patch_payload, TIRetryStatePayload):
            if ti is not None:
                # Set the overrides on the TI *before* archiving so record_ti()
                # snapshots them into task_instance_history (it copies attrs off
                # the ti object). Otherwise the per-try audit trail is always NULL.
                ti.retry_delay_override = ti_patch_payload.retry_delay_seconds
                ti.retry_reason = (
                    ti_patch_payload.retry_reason[:500] if ti_patch_payload.retry_reason else None
                )
                ti.prepare_db_for_next_try(session=session)
            # Store retry policy overrides so next_retry_datetime() can read them.
            # These are cleared when the task enters RUNNING (ti_run).
            query = query.values(
                retry_delay_override=ti_patch_payload.retry_delay_seconds,
                retry_reason=(ti_patch_payload.retry_reason[:500] if ti_patch_payload.retry_reason else None),
            )
        elif isinstance(ti_patch_payload, TISuccessStatePayload):
            if ti is not None:
                TI.register_asset_changes_in_db(
                    ti,
                    ti_patch_payload.task_outlets,
                    ti_patch_payload.outlet_events,
                    session=session,
                )
        try:
            _emit_task_span(ti, state=updated_state)
        except Exception:
            log.warning("Failed to emit task span", exc_info=True)
    elif isinstance(ti_patch_payload, TIDeferredStatePayload):
        # Calculate timeout if it was passed
        timeout = None
        if ti_patch_payload.trigger_timeout is not None:
            timeout = timezone.utcnow() + ti_patch_payload.trigger_timeout

        trigger_kwargs = ti_patch_payload.trigger_kwargs
        if not isinstance(trigger_kwargs, str):
            # If it's passed as a string, assume the client encrypted it, otherwise assume it doesn't need to
            # be. Just JSON serialize it
            trigger_kwargs = json.dumps(trigger_kwargs)

        trigger_row = Trigger(
            classpath=ti_patch_payload.classpath,
            kwargs={},
            queue=ti_patch_payload.queue,
            team_name=get_team_name_for_ti(task_instance_id, session),
        )
        trigger_row.encrypted_kwargs = trigger_kwargs
        session.add(trigger_row)
        session.flush()

        # TODO: HANDLE execution timeout later as it requires a call to the DB
        # either get it from the serialised DAG or get it from the API

        query = update(TI).where(TI.id == task_instance_id)

        # Store next_kwargs directly (already serialized by worker)
        query = query.values(
            state=TaskInstanceState.DEFERRED,
            trigger_id=trigger_row.id,
            next_method=ti_patch_payload.next_method,
            next_kwargs=ti_patch_payload.next_kwargs,
            trigger_timeout=timeout,
        )
        updated_state = TaskInstanceState.DEFERRED
    elif isinstance(ti_patch_payload, TIAwaitingInputStatePayload):
        # Park the task waiting for human input (Human-in-the-loop). No trigger / triggerer is
        # created: the task is resumed by the Core API response handler or the scheduler timeout
        # sweep. The optional response deadline is stored on the existing trigger_timeout column.
        #
        # Fixed lock order (TaskInstance -> HITLDetail), matching the Core API response path, so a
        # human response racing this park transition cannot deadlock.
        ti = session.get(TI, task_instance_id, with_for_update={"of": TI})
        # Lock only the hitl_detail row (of=...): HITLDetail eager-joins task_instance (lazy="joined"),
        # and Postgres rejects FOR UPDATE against the nullable side of that outer join.
        hitl_detail = session.scalar(
            select(HITLDetail).where(HITLDetail.ti_id == task_instance_id).with_for_update(of=HITLDetail)
        )
        if ti is not None and hitl_detail is not None and hitl_detail.response_received:
            # The human responded in the window between the operator writing the HITL request and
            # the worker parking the task. Resume straight to execute_complete instead of parking,
            # which would otherwise strand an already-responded task (no trigger/sweep would fire).
            # Carry next_method/next_kwargs onto the TI first so the resume dispatches correctly;
            # handle_event_submit then injects the response event into next_kwargs.
            ti.next_method = ti_patch_payload.next_method
            ti.next_kwargs = ti_patch_payload.next_kwargs
            handle_event_submit(
                TriggerEvent(hitl_detail.as_resume_event_payload()),
                task_instance=ti,
                session=session,
            )
            query = update(TI).where(TI.id == task_instance_id).values(state=TaskInstanceState.SCHEDULED)
            updated_state = TaskInstanceState.SCHEDULED
        else:
            timeout = None
            if ti_patch_payload.timeout is not None:
                timeout = timezone.utcnow() + ti_patch_payload.timeout

            query = update(TI).where(TI.id == task_instance_id)
            query = query.values(
                state=TaskInstanceState.AWAITING_INPUT,
                trigger_id=None,
                next_method=ti_patch_payload.next_method,
                next_kwargs=ti_patch_payload.next_kwargs,
                trigger_timeout=timeout,
            )
            updated_state = TaskInstanceState.AWAITING_INPUT
    elif isinstance(ti_patch_payload, TIRescheduleStatePayload):
        # Quick check for poke_interval isn't immediately over MySQL's TIMESTAMP limit.
        # This check is only rudimentary to catch trivial user errors, e.g. mistakenly
        # set the value to milliseconds instead of seconds. There's another check when
        # we actually try to reschedule to ensure database coherence.
        if get_dialect_name(session) == "mysql":
            # As documented in https://dev.mysql.com/doc/refman/5.7/en/datetime.html.
            _MYSQL_TIMESTAMP_MAX = timezone.datetime(2038, 1, 19, 3, 14, 7)
            if ti_patch_payload.reschedule_date > _MYSQL_TIMESTAMP_MAX:
                # Set a task to failed in case any unexpected exception happened during task state update
                log.error(
                    "Cannot reschedule task past MySQL limit. Setting the task to failed.",
                    payload=ti_patch_payload,
                    mysql_timestamp_max=_MYSQL_TIMESTAMP_MAX,
                )
                data = ti_patch_payload.model_dump(exclude={"reschedule_date"}, exclude_unset=True)
                query = update(TI).where(TI.id == task_instance_id).values(data)
                if session.bind is not None:
                    query = TI.duration_expression_update(timezone.utcnow(), query, session.bind)
                query = query.values(state=TaskInstanceState.FAILED)
                ti = session.get(TI, task_instance_id, with_for_update={"of": TI})
                if ti is not None:
                    _handle_fail_fast_for_dag(ti=ti, dag_id=dag_id, session=session, dag_bag=dag_bag)
                return query, TaskInstanceState.FAILED

        actual_start_date = timezone.utcnow()
        session.add(
            TaskReschedule(
                task_instance_id,
                actual_start_date,
                ti_patch_payload.end_date,
                ti_patch_payload.reschedule_date,
            )
        )

        query = update(TI).where(TI.id == task_instance_id)
        # calculate the duration for TI table too
        if session.bind is not None:
            query = TI.duration_expression_update(ti_patch_payload.end_date, query, session.bind)
        # clear the next_method and next_kwargs so that none of the retries pick them up
        updated_state = TaskInstanceState.UP_FOR_RESCHEDULE
        query = query.values(state=updated_state, next_method=None, next_kwargs=None)
    else:
        raise ValueError(f"Unexpected Payload Type {type(ti_patch_payload)}")

    return query, updated_state


@ti_id_router.patch(
    "/{task_instance_id}/skip-downstream",
    status_code=status.HTTP_204_NO_CONTENT,
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance not found"),
            (HTTP_422_UNPROCESSABLE_CONTENT, "Invalid payload for the state transition"),
        ]
    ),
)
def ti_skip_downstream(
    task_instance_id: UUID,
    ti_patch_payload: TISkippedDownstreamTasksStatePayload,
    session: SessionDep,
):
    bind_contextvars(ti_id=str(task_instance_id))
    log.info("Skipping downstream tasks", task_count=len(ti_patch_payload.tasks))

    now = timezone.utcnow()
    tasks = ti_patch_payload.tasks

    query_result = session.execute(select(TI.dag_id, TI.run_id).where(TI.id == task_instance_id))
    row_result = query_result.fetchone()
    if row_result is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={"reason": "not_found", "message": "Task Instance not found"},
        )
    dag_id, run_id = row_result
    log.debug("Retrieved DAG and run info", dag_id=dag_id, run_id=run_id)

    task_ids = [task if isinstance(task, tuple) else (task, -1) for task in tasks]
    log.debug("Prepared task IDs for skipping", task_ids=task_ids)

    # Don't overwrite tasks that are already executing or finished.
    # See: https://github.com/apache/airflow/issues/59378
    # Note: SQL NULL NOT IN (...) is falsy, so we need an explicit IS NULL check.
    skippable_state_clause = or_(
        TI.state.is_(None),
        TI.state.not_in(
            [
                TaskInstanceState.RUNNING,
                TaskInstanceState.SUCCESS,
                TaskInstanceState.FAILED,
            ]
        ),
    )
    query = (
        update(TI)
        .where(
            TI.dag_id == dag_id,
            TI.run_id == run_id,
            tuple_(TI.task_id, TI.map_index).in_(task_ids),
            skippable_state_clause,
        )
        .values(state=TaskInstanceState.SKIPPED, start_date=now, end_date=now)
        .execution_options(synchronize_session=False)
    )

    result = session.execute(query)
    log.info("Downstream tasks skipped", tasks_skipped=getattr(result, "rowcount", 0))


@ti_id_router.put(
    "/{task_instance_id}/heartbeat",
    status_code=status.HTTP_204_NO_CONTENT,
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance not found"),
            (
                status.HTTP_409_CONFLICT,
                "The TI attempting to heartbeat should be terminated for the given reason",
            ),
            (
                status.HTTP_410_GONE,
                "Task Instance not found in the TI table but exists in the Task Instance History table",
            ),
            (HTTP_422_UNPROCESSABLE_CONTENT, "Invalid payload for the state transition"),
        ]
    ),
)
def ti_heartbeat(
    task_instance_id: UUID,
    ti_payload: TIHeartbeatInfo,
    session: SessionDep,
):
    """Update the heartbeat of a TaskInstance to mark it as alive & still running."""
    bind_contextvars(ti_id=str(task_instance_id))
    log.debug("Processing heartbeat", hostname=ti_payload.hostname, pid=ti_payload.pid)

    # Hot path: in the common case the TI is still running on the same host and pid,
    # so we can update last_heartbeat_at directly without first taking a row lock.
    fast_path_result = cast(
        "CursorResult[Any]",
        session.execute(
            update(TI)
            .where(
                TI.id == task_instance_id,
                TI.state == TaskInstanceState.RUNNING,
                TI.hostname == ti_payload.hostname,
                TI.pid == ti_payload.pid,
            )
            .values(last_heartbeat_at=timezone.utcnow())
            .execution_options(synchronize_session=False)
        ),
    )
    if fast_path_result.rowcount is not None and fast_path_result.rowcount > 0:
        log.debug("Heartbeat updated via fast path")
        return

    log.debug("Heartbeat fast path missed; falling back to diagnostic checks")

    old = select(TI.state, TI.hostname, TI.pid).where(TI.id == task_instance_id).with_for_update()

    try:
        (previous_state, hostname, pid) = session.execute(old).one()
        log.debug(
            "Retrieved current task state", state=previous_state, current_hostname=hostname, current_pid=pid
        )
    except NoResultFound:
        # Check if the TI exists in the Task Instance History table.
        # If it does, it was likely cleared while running, so return 410 Gone
        # instead of 404 Not Found to give the client a more specific signal.
        tih_exists = session.scalar(
            select(func.count(TIH.task_instance_id)).where(TIH.task_instance_id == task_instance_id)
        )
        if tih_exists:
            log.error(
                "TaskInstance was previously cleared and archived in history, heartbeat skipped",
                ti_id=str(task_instance_id),
            )
            raise HTTPException(
                status_code=status.HTTP_410_GONE,
                detail={
                    "reason": "not_found",
                    "message": "Task Instance not found, it may have been moved to the Task Instance History table",
                },
            )
        log.error("Task Instance not found")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={
                "reason": "not_found",
                "message": "Task Instance not found",
            },
        )

    if hostname != ti_payload.hostname or pid != ti_payload.pid:
        log.warning(
            "Task running elsewhere",
            current_hostname=hostname,
            current_pid=pid,
            requested_hostname=ti_payload.hostname,
            requested_pid=ti_payload.pid,
        )
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail={
                "reason": "running_elsewhere",
                "message": "TI is already running elsewhere",
                "current_hostname": hostname,
                "current_pid": pid,
            },
        )

    if previous_state != TaskInstanceState.RUNNING:
        log.warning("Task not in running state", current_state=previous_state)
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail={
                "reason": "not_running",
                "message": "TI is no longer in the running state and task should terminate",
                "current_state": previous_state,
            },
        )

    # Update the last heartbeat time!
    session.execute(update(TI).where(TI.id == task_instance_id).values(last_heartbeat_at=timezone.utcnow()))
    log.debug("Heartbeat updated", state=previous_state)


@ti_id_router.put(
    "/{task_instance_id}/rtif",
    status_code=status.HTTP_201_CREATED,
    operation_id="put_rtif",
    summary="Set Rendered Task Instance Fields",
    description="Store the rendered task instance fields (RTIF) for a task instance. "
    "These are the template fields after Jinja rendering has been applied. "
    "Called by the worker after task execution begins.",
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance not found"),
            (
                HTTP_422_UNPROCESSABLE_CONTENT,
                "Invalid payload for the setting rendered task instance fields",
            ),
        ]
    ),
)
def ti_put_rtif(
    task_instance_id: UUID,
    put_rtif_payload: Annotated[dict[str, JsonValue], Body()],
    session: SessionDep,
):
    """Add an RTIF entry for a task instance, sent by the worker."""
    bind_contextvars(ti_id=str(task_instance_id))
    log.info("Updating RenderedTaskInstanceFields", field_count=len(put_rtif_payload))

    task_instance = session.scalar(select(TI).where(TI.id == task_instance_id))
    if not task_instance:
        log.error("Task Instance not found")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
        )
    task_instance.update_rtif(put_rtif_payload, session=session)
    log.debug("RenderedTaskInstanceFields updated successfully")

    return {"message": "Rendered task instance fields successfully set"}


@ti_id_router.patch(
    "/{task_instance_id}/rendered-map-index",
    status_code=status.HTTP_204_NO_CONTENT,
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance not found"),
            (
                HTTP_422_UNPROCESSABLE_CONTENT,
                "Invalid rendered_map_index value",
            ),
        ]
    ),
)
def ti_patch_rendered_map_index(
    task_instance_id: UUID,
    rendered_map_index: Annotated[str, Body()],
    session: SessionDep,
):
    """Update rendered_map_index for a task instance, sent by the worker during task execution."""
    bind_contextvars(ti_id=str(task_instance_id))

    if not rendered_map_index:
        log.error("rendered_map_index cannot be empty")
        raise HTTPException(
            status_code=HTTP_422_UNPROCESSABLE_CONTENT,
            detail="rendered_map_index cannot be empty",
        )

    log.debug("Updating rendered_map_index", length=len(rendered_map_index))

    query = update(TI).where(TI.id == task_instance_id).values(_rendered_map_index=rendered_map_index)
    result = session.execute(query)

    result = cast("CursorResult[Any]", result)
    if result.rowcount == 0:
        log.error("Task Instance not found")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Task Instance not found",
        )


@ti_id_router.get(
    "/{task_instance_id}/previous-successful-dagrun",
    status_code=status.HTTP_200_OK,
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance or Dag Run not found"),
        ]
    ),
)
def get_previous_successful_dagrun(
    task_instance_id: UUID, session: SessionDep
) -> PrevSuccessfulDagRunResponse:
    """
    Get the previous successful DagRun for a TaskInstance.

    The data from this endpoint is used to get values for Task Context.
    """
    bind_contextvars(ti_id=str(task_instance_id))
    log.debug("Retrieving previous successful DAG run")

    task_instance = session.scalar(select(TI).where(TI.id == task_instance_id))
    if not task_instance or not task_instance.logical_date:
        log.debug("No task instance or logical date found")
        return PrevSuccessfulDagRunResponse()

    dag_run = session.scalar(
        select(DR)
        .where(
            DR.dag_id == task_instance.dag_id,
            DR.logical_date < task_instance.logical_date,
            DR.state == DagRunState.SUCCESS,
        )
        .order_by(DR.logical_date.desc())
        .limit(1)
    )
    if not dag_run:
        log.debug("No previous successful DAG run found")
        return PrevSuccessfulDagRunResponse()

    log.debug(
        "Found previous successful DAG run",
        dag_id=dag_run.dag_id,
        run_id=dag_run.run_id,
        logical_date=dag_run.logical_date,
    )
    return PrevSuccessfulDagRunResponse.model_validate(dag_run)


@router.get("/count", status_code=status.HTTP_200_OK)
def get_task_instance_count(
    dag_id: str,
    session: SessionDep,
    dag_bag: DagBagDep,
    map_index: Annotated[int | None, Query()] = None,
    task_ids: Annotated[list[str] | None, Query()] = None,
    task_group_id: Annotated[str | None, Query()] = None,
    logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
    run_ids: Annotated[list[str] | None, Query()] = None,
    states: Annotated[list[str] | None, Query()] = None,
) -> int:
    """Get the count of task instances matching the given criteria."""
    query = select(func.count()).select_from(TI).where(TI.dag_id == dag_id)

    if task_ids:
        query = query.where(TI.task_id.in_(task_ids))

    if map_index is not None:
        query = query.where(TI.map_index == map_index)

    if logical_dates:
        query = query.where(TI.logical_date.in_(logical_dates))

    if run_ids:
        query = query.where(TI.run_id.in_(run_ids))

    if task_group_id:
        group_tasks = _get_group_tasks(
            dag_id, task_group_id, session, dag_bag, logical_dates, run_ids, map_index
        )

        # Get unique (task_id, map_index) pairs
        task_map_pairs = [(ti.task_id, ti.map_index) for ti in group_tasks]

        if not task_map_pairs:
            # If no task group tasks found, default to checking the task group ID itself
            # This matches the behavior in _get_external_task_group_task_ids
            task_map_pairs = [(task_group_id, -1)]

        # Update query to use task_id, map_index pairs
        query = query.where(tuple_(TI.task_id, TI.map_index).in_(task_map_pairs))

    if states:
        if "null" in states:
            not_none_states = [s for s in states if s != "null"]
            if not_none_states:
                query = query.where(or_(TI.state.is_(None), TI.state.in_(not_none_states)))
            else:
                query = query.where(TI.state.is_(None))
        else:
            query = query.where(TI.state.in_(states))

    count = session.scalar(query)
    return count or 0


@router.get("/previous/{dag_id}/{task_id}", status_code=status.HTTP_200_OK)
def get_previous_task_instance(
    dag_id: str,
    task_id: str,
    session: SessionDep,
    logical_date: Annotated[UtcDateTime | None, Query()] = None,
    map_index: Annotated[int, Query()] = -1,
    state: Annotated[TaskInstanceState | None, Query()] = None,
) -> PreviousTIResponse | None:
    """
    Get the previous task instance matching the given criteria.

    :param dag_id: DAG ID (from path)
    :param task_id: Task ID (from path)
    :param logical_date: If provided, finds TI with logical_date < this value (before filter)
    :param map_index: Map index to filter by (defaults to -1 for non-mapped tasks)
    :param state: If provided, filters by TaskInstance state
    """
    query = (
        select(TI)
        .join(DR, (TI.dag_id == DR.dag_id) & (TI.run_id == DR.run_id))
        .options(joinedload(TI.dag_run))
        .where(TI.dag_id == dag_id, TI.task_id == task_id, TI.map_index == map_index)
        .order_by(DR.logical_date.desc())
    )

    if logical_date:
        # Find TI with logical_date BEFORE the provided date (previous)
        query = query.where(DR.logical_date < logical_date)

    if state:
        query = query.where(TI.state == state)

    ti = session.scalars(query).first()

    if not ti:
        return None

    return PreviousTIResponse(
        task_id=ti.task_id,
        dag_id=ti.dag_id,
        run_id=ti.run_id,
        logical_date=ti.dag_run.logical_date,
        start_date=ti.start_date,
        end_date=ti.end_date,
        state=ti.state,
        try_number=ti.try_number,
        map_index=ti.map_index,
        duration=ti.duration,
    )


@router.get("/states", status_code=status.HTTP_200_OK)
def get_task_instance_states(
    dag_id: str,
    session: SessionDep,
    dag_bag: DagBagDep,
    map_index: Annotated[int | None, Query()] = None,
    task_ids: Annotated[list[str] | None, Query()] = None,
    task_group_id: Annotated[str | None, Query()] = None,
    logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
    run_ids: Annotated[list[str] | None, Query()] = None,
) -> TaskStatesResponse:
    """Get the states for Task Instances with the given criteria."""
    run_id_task_state_map: dict[str, dict[str, Any]] = defaultdict(dict)

    query = select(TI).where(TI.dag_id == dag_id)

    if task_ids:
        query = query.where(TI.task_id.in_(task_ids))

    if logical_dates:
        query = query.where(TI.logical_date.in_(logical_dates))

    if run_ids:
        query = query.where(TI.run_id.in_(run_ids))

    if map_index is not None:
        query = query.where(TI.map_index == map_index)

    results = session.scalars(query).all()

    if task_group_id:
        group_tasks = _get_group_tasks(
            dag_id, task_group_id, session, dag_bag, logical_dates, run_ids, map_index
        )

        results = results + group_tasks if task_ids else group_tasks

    [
        run_id_task_state_map[task.run_id].update(
            {task.task_id: task.state}
            if task.map_index < 0
            else {f"{task.task_id}_{task.map_index}": task.state}
        )
        for task in results
    ]

    return TaskStatesResponse(task_states=run_id_task_state_map)


@router.get("/breadcrumbs", status_code=status.HTTP_200_OK)
def get_task_instance_breadcrumbs(dag_id: str, run_id: str, session: SessionDep) -> TaskBreadcrumbsResponse:
    result = session.execute(
        select(TI.task_id, TI.map_index, TI.state, TI.operator, TI.duration)
        .where(TI.dag_id == dag_id, TI.run_id == run_id, TI.state.in_(TerminalTIState))
        .order_by(TI.task_id, TI.map_index)
    ).mappings()

    def _iter_breadcrumbs() -> Iterator[dict[str, Any]]:
        for row in result:
            yield {str(k): v for k, v in row.items()}

    return TaskBreadcrumbsResponse(breadcrumbs=_iter_breadcrumbs())


def _is_eligible_to_retry(state: str, try_number: int, max_tries: int) -> bool:
    """Is task instance is eligible for retry."""
    if state == TaskInstanceState.RESTARTING:
        # If a task is cleared when running, it goes into RESTARTING state and is always
        # eligible for retry
        return True

    # max_tries is initialised with the retries defined at task level, we do not need to explicitly ask for
    # retries from the task SDK now, we can handle using max_tries
    return max_tries != 0 and try_number <= max_tries


def _get_group_tasks(
    dag_id: str,
    task_group_id: str,
    session: SessionDep,
    dag_bag: DagBagDep,
    logical_dates=None,
    run_ids=None,
    map_index: int | None = None,
):
    # Get all tasks in the task group
    dag = get_latest_version_of_dag(dag_bag, dag_id, session, include_reason=True)
    task_group = dag.task_group_dict.get(task_group_id)
    if not task_group:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail={
                "reason": "not_found",
                "message": f"Task group {task_group_id} not found in DAG {dag_id}",
            },
        )

    # First get all task instances to get the task_id, map_index pairs
    group_tasks = session.scalars(
        select(TI).where(
            TI.dag_id == dag_id,
            TI.task_id.in_(task.task_id for task in task_group.iter_tasks()),
            *([TI.logical_date.in_(logical_dates)] if logical_dates else []),
            *([TI.run_id.in_(run_ids)] if run_ids else []),
            *([TI.map_index == map_index] if map_index is not None else []),
        )
    ).all()

    return group_tasks


@ti_id_router.get(
    "/{task_instance_id}/validate-inlets-and-outlets",
    status_code=status.HTTP_200_OK,
    responses=create_openapi_http_exception_doc(
        [
            (status.HTTP_404_NOT_FOUND, "Task Instance not found"),
        ]
    ),
)
def validate_inlets_and_outlets(
    task_instance_id: UUID,
    session: SessionDep,
    dag_bag: DagBagDep,
) -> InactiveAssetsResponse:
    """Validate whether there're inactive assets in inlets and outlets of a given task instance."""
    bind_contextvars(ti_id=str(task_instance_id))

    ti = session.scalar(select(TI).where(TI.id == task_instance_id))
    if not ti:
        log.error("Task Instance not found")
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={
                "reason": "not_found",
                "message": "Task Instance not found",
            },
        )

    if not ti.task:
        dr = ti.dag_run
        dag = dag_bag.get_dag_for_run(dag_run=dr, session=session)
        if dag:
            with contextlib.suppress(TaskNotFound):
                ti.task = dag.get_task(ti.task_id)

    inlets = (
        [asset.asprofile() for asset in ti.task.inlets if isinstance(asset, SerializedAsset)]
        if ti.task
        else []
    )
    outlets = (
        [asset.asprofile() for asset in ti.task.outlets if isinstance(asset, SerializedAsset)]
        if ti.task
        else []
    )
    if not (inlets or outlets):
        return InactiveAssetsResponse(inactive_assets=[])

    all_asset_unique_keys: set[SerializedAssetUniqueKey] = {
        SerializedAssetUniqueKey.from_asset(inlet_or_outlet)  # type: ignore
        for inlet_or_outlet in itertools.chain(inlets, outlets)
    }
    active_asset_unique_keys = {
        SerializedAssetUniqueKey(name, uri)
        for name, uri in session.execute(
            select(AssetActive.name, AssetActive.uri).where(
                tuple_(AssetActive.name, AssetActive.uri).in_(
                    attrs.astuple(key) for key in all_asset_unique_keys
                )
            )
        )
    }
    different = all_asset_unique_keys - active_asset_unique_keys

    return InactiveAssetsResponse(
        inactive_assets=[asset_unique_key.asprofile() for asset_unique_key in different],
    )


# This line should be at the end of the file to ensure all routes are registered
router.include_router(ti_id_router)
