# 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 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.execution_api.datamodels.taskinstance import (
    InactiveAssetsResponse,
    PreviousTIResponse,
    PrevSuccessfulDagRunResponse,
    TaskBreadcrumbsResponse,
    TaskStatesResponse,
    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, require_auth
from airflow.configuration import conf
from airflow.exceptions import TaskNotFound
from airflow.models.asset import AssetActive
from airflow.models.dag import DagModel
from airflow.models.dagrun import DagRun as DR
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
from airflow.models.xcom import XComModel
from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey
from airflow.state import get_state_backend
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={
        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
        status.HTTP_409_CONFLICT: {"description": "The TI is already in the requested state"},
        HTTP_422_UNPROCESSABLE_CONTENT: {"description": "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
        )

        from airflow.api_fastapi.execution_api.security import get_team_name_for_ti

        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 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_404_NOT_FOUND: {"description": "Task Instance not found"},
        status.HTTP_409_CONFLICT: {"description": "The TI is already in the requested state"},
        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",
            },
        )

    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,
            },
        )

    # 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 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 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)

    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 _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:
                ti.prepare_db_for_next_try(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,
                )
        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,
        )
        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, 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={
        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
        HTTP_422_UNPROCESSABLE_CONTENT: {"description": "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={
        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
        status.HTTP_409_CONFLICT: {
            "description": "The TI attempting to heartbeat should be terminated for the given reason"
        },
        status.HTTP_410_GONE: {
            "description": "Task Instance not found in the TI table but exists in the Task Instance History table"
        },
        HTTP_422_UNPROCESSABLE_CONTENT: {"description": "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={
        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
        HTTP_422_UNPROCESSABLE_CONTENT: {
            "description": "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)
    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={
        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
        HTTP_422_UNPROCESSABLE_CONTENT: {"description": "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={
        status.HTTP_404_NOT_FOUND: {"description": "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={
        status.HTTP_404_NOT_FOUND: {"description": "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)
