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

from __future__ import annotations

from collections.abc import Sequence
from typing import Literal

import structlog
from fastapi import HTTPException, status
from fastapi.exceptions import RequestValidationError
from pydantic import ValidationError
from sqlalchemy import select, tuple_
from sqlalchemy.orm import joinedload
from sqlalchemy.orm.session import Session

from airflow._shared.state import TaskScope
from airflow.api_fastapi.app import get_auth_manager
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
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.core_api.datamodels.common import (
    BulkActionNotOnExistence,
    BulkActionResponse,
    BulkBody,
    BulkCreateAction,
    BulkDeleteAction,
    BulkUpdateAction,
)
from airflow.api_fastapi.core_api.datamodels.task_instances import (
    BulkTaskInstanceBody,
    ClearTaskInstancesBody,
    PatchTaskInstanceBody,
)
from airflow.api_fastapi.core_api.security import GetUserDep
from airflow.api_fastapi.core_api.services.public.common import BulkService
from airflow.configuration import conf
from airflow.listeners.listener import get_listener_manager
from airflow.models.dag import DagModel
from airflow.models.taskinstance import TaskInstance as TI
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.state.metastore import _get_db_backend
from airflow.utils.state import TaskInstanceState

log = structlog.get_logger(__name__)


def _clear_task_state_store_on_success(tis: Sequence[TI], session: Session) -> None:
    """Clear task state store rows for each TI if clear_on_success is enabled."""
    if not conf.getboolean("state_store", "clear_on_success", fallback=False):
        return
    backend = _get_db_backend()
    for ti in tis:
        scope = TaskScope(
            dag_id=ti.dag_id,
            run_id=ti.run_id,
            task_id=ti.task_id,
            map_index=ti.map_index if ti.map_index is not None else -1,
        )
        try:
            backend.clear(scope=scope, session=session)
            log.info(
                "Cleared task state on success",
                dag_id=ti.dag_id,
                run_id=ti.run_id,
                task_id=ti.task_id,
                map_index=ti.map_index,
            )
        except Exception:
            log.warning(
                "Failed to clear task state on success",
                dag_id=ti.dag_id,
                run_id=ti.run_id,
                task_id=ti.task_id,
            )


def _validate_patch_task_instance_body(
    body: PatchTaskInstanceBody,
    update_mask: list[str] | None,
) -> dict:
    """Validate the patch body and return the fields to update as a dict."""
    fields_to_update = body.model_fields_set
    if update_mask:
        fields_to_update = fields_to_update.intersection(update_mask)
    else:
        try:
            PatchTaskInstanceBody.model_validate(body)
        except ValidationError as e:
            raise RequestValidationError(errors=e.errors())

    return body.model_dump(include=fields_to_update, by_alias=True)


def _emit_state_listener_hooks(updated_tis: list[TI], new_state: str | TaskInstanceState) -> None:
    """Fire listener hooks for the given TIs based on their new state. Listener errors are logged."""
    for ti in updated_tis:
        try:
            if new_state == TaskInstanceState.SUCCESS:
                get_listener_manager().hook.on_task_instance_success(previous_state=None, task_instance=ti)
            elif new_state == TaskInstanceState.FAILED:
                get_listener_manager().hook.on_task_instance_failed(
                    previous_state=None,
                    task_instance=ti,
                    error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.",
                )
            elif new_state == TaskInstanceState.SKIPPED:
                get_listener_manager().hook.on_task_instance_skipped(previous_state=None, task_instance=ti)
        except Exception:
            log.exception("error calling listener")


def _reload_tis_with_rendered_fields(tis: list[TI], session: Session) -> list[TI]:
    """
    Re-load TIs with ``rendered_task_instance_fields`` eagerly loaded.

    ``set_task_instance_state`` / ``set_task_group_state`` return TIs without this relationship
    loaded; we re-query so they can be serialized without lazy loads.
    ``populate_existing=True`` ensures the joinedload updates TIs already in the identity map.
    """
    if not tis:
        return tis
    return list(
        session.scalars(
            select(TI)
            .options(joinedload(TI.rendered_task_instance_fields))
            .where(TI.id.in_([ti.id for ti in tis]))
            .execution_options(populate_existing=True)
        )
        .unique()
        .all()
    )


def _patch_ti_validate_request(
    dag_id: str,
    dag_run_id: str,
    task_id: str,
    dag_bag: DagBagDep,
    body: PatchTaskInstanceBody,
    session: SessionDep,
    map_index: int | None = -1,
    update_mask: list[str] | None = None,
) -> tuple[SerializedDAG, list[TI], dict]:
    dag = get_latest_version_of_dag(dag_bag, dag_id, session)
    if not dag.has_task(task_id):
        raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not found in Dag '{dag_id}'")

    query = (
        select(TI)
        .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
        .options(joinedload(TI.rendered_task_instance_fields))
    )
    if map_index is not None:
        query = query.where(TI.map_index == map_index)
    else:
        query = query.order_by(TI.map_index)

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

    err_msg_404 = (
        f"The Task Instance with dag_id: `{dag_id}`, run_id: `{dag_run_id}`, task_id: `{task_id}` and map_index: `{map_index}` was not found",
    )
    if len(tis) == 0:
        raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)

    data = _validate_patch_task_instance_body(body, update_mask)
    return dag, list(tis), data


def _get_task_group_task_instances(
    dag_id: str,
    dag_run_id: str,
    task_group_id: str,
    dag: SerializedDAG,
    session: Session,
) -> list[TI]:
    """Get all task instances in a task group for a specific DAG run."""
    task_group = dag.task_group_dict.get(task_group_id)
    if not task_group:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND, f"Task group '{task_group_id}' not found in DAG '{dag_id}'"
        )

    task_ids = [task.task_id for task in task_group.iter_tasks()]

    query = (
        select(TI)
        .where(
            TI.dag_id == dag_id,
            TI.run_id == dag_run_id,
            TI.task_id.in_(task_ids),
        )
        .order_by(TI.task_id, TI.map_index)
    )

    group_tis = list(session.scalars(query).all())
    if not group_tis:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            f"No task instances found for task group '{task_group_id}' in dag run '{dag_run_id}'",
        )

    return group_tis


def _patch_ti_group_validate_request(
    dag_id: str,
    dag_run_id: str,
    task_group_id: str,
    dag_bag: DagBagDep,
    body: PatchTaskInstanceBody,
    session: SessionDep,
    update_mask: list[str] | None = None,
) -> tuple[SerializedDAG, list[TI], dict]:
    """Validate and prepare data for task group patch request."""
    dag = get_latest_version_of_dag(dag_bag, dag_id, session)
    tis = _get_task_group_task_instances(dag_id, dag_run_id, task_group_id, dag, session)

    data = _validate_patch_task_instance_body(body, update_mask)
    return dag, tis, data


def _patch_task_instance_state(
    task_id: str,
    dag_run_id: str,
    dag: SerializedDAG,
    task_instance_body: BulkTaskInstanceBody | PatchTaskInstanceBody,
    data: dict,
    session: Session,
) -> list[TI]:
    map_index = getattr(task_instance_body, "map_index", None)
    map_indexes = None if map_index is None else [map_index]

    updated_tis = dag.set_task_instance_state(
        task_id=task_id,
        run_id=dag_run_id,
        map_indexes=map_indexes,
        state=data["new_state"],
        upstream=task_instance_body.include_upstream,
        downstream=task_instance_body.include_downstream,
        future=task_instance_body.include_future,
        past=task_instance_body.include_past,
        commit=True,
        session=session,
    )
    if not updated_tis:
        raise HTTPException(
            status.HTTP_409_CONFLICT,
            f"Task id {task_id} is already in {data['new_state']} state",
        )

    if data["new_state"] == TaskInstanceState.SUCCESS:
        _clear_task_state_store_on_success(updated_tis, session)

    _emit_state_listener_hooks(updated_tis, data["new_state"])

    return updated_tis


def _patch_task_group_state(
    group_id: str,
    dag_run_id: str,
    dag: SerializedDAG,
    body: PatchTaskInstanceBody,
    data: dict,
    *,
    session: Session,
) -> list[TI]:
    """Update the state of all task instances in a task group."""
    updated_tis = dag.set_task_group_state(
        group_id=group_id,
        run_id=dag_run_id,
        state=data["new_state"],
        upstream=body.include_upstream,
        downstream=body.include_downstream,
        future=body.include_future,
        past=body.include_past,
        commit=True,
        session=session,
    )
    if not updated_tis:
        raise HTTPException(
            status.HTTP_409_CONFLICT,
            f"All task instances in the group are already in {data['new_state']} state",
        )

    if data["new_state"] == TaskInstanceState.SUCCESS:
        _clear_task_state_store_on_success(updated_tis, session)

    _emit_state_listener_hooks(updated_tis, data["new_state"])

    return updated_tis


def _patch_task_instance_note(
    task_instance_body: BulkTaskInstanceBody | ClearTaskInstancesBody | PatchTaskInstanceBody,
    tis: list[TI],
    user: GetUserDep,
    update_mask: list[str] | None = None,
) -> None:
    for ti in tis:
        if update_mask or task_instance_body.note is not None:
            if ti.task_instance_note is None:
                ti.note = (task_instance_body.note, user.get_id())
            else:
                ti.task_instance_note.content = task_instance_body.note
                ti.task_instance_note.user_id = user.get_id()


class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
    """Service for handling bulk operations on task instances."""

    def __init__(
        self,
        session: Session,
        request: BulkBody[BulkTaskInstanceBody],
        dag_id: str,
        dag_run_id: str,
        dag_bag: DagBagDep,
        user: GetUserDep,
    ):
        super().__init__(session, request)
        self.dag_id = dag_id
        self.dag_run_id = dag_run_id
        self.dag_bag = dag_bag
        self.user = user

    def _extract_task_identifiers(
        self, entity: str | BulkTaskInstanceBody
    ) -> tuple[str, str, str, int | None]:
        """
        Extract task identifiers from an id or entity object.

        :param entity: Task identifier as string or BulkTaskInstanceBody object
        :return: tuple of (dag_id, dag_run_id, task_id, map_index)
        """
        if isinstance(entity, str):
            dag_id = self.dag_id
            dag_run_id = self.dag_run_id
            task_id = entity
            map_index = None
        else:
            dag_id = entity.dag_id if entity.dag_id else self.dag_id
            dag_run_id = entity.dag_run_id if entity.dag_run_id else self.dag_run_id
            task_id = entity.task_id
            map_index = entity.map_index

        return dag_id, dag_run_id, task_id, map_index

    def _categorize_entities(
        self,
        entities: Sequence[str | BulkTaskInstanceBody],
        results: BulkActionResponse,
        method: Literal["PUT", "DELETE"],
        action_name: str,
    ) -> tuple[set[tuple[str, str, str, int]], set[tuple[str, str, str]]]:
        """
        Validate entities and categorize them into specific and all map index update sets.

        :param entities: Sequence of entities to validate
        :param results: BulkActionResponse object to track errors
        :return: tuple of (specific_map_index_task_keys, all_map_index_task_keys)
        """
        specific_map_index_task_keys = set()
        all_map_index_task_keys = set()
        dag_authorization_cache: dict[str, bool] = {}

        for entity in entities:
            dag_id, dag_run_id, task_id, map_index = self._extract_task_identifiers(entity)

            # Validate that we have specific values, not wildcards
            if dag_id == "~" or dag_run_id == "~":
                if isinstance(entity, str):
                    error_msg = f"When using wildcard in path, dag_id and dag_run_id must be specified in BulkTaskInstanceBody object, not as string for task_id: {entity}"
                else:
                    error_msg = f"When using wildcard in path, dag_id and dag_run_id must be specified in request body for task_id: {entity.task_id}"
                results.errors.append(
                    {
                        "error": error_msg,
                        "status_code": status.HTTP_400_BAD_REQUEST,
                    }
                )
                continue

            if dag_id not in dag_authorization_cache:
                team_name = DagModel.get_team_name(dag_id, session=self.session)
                dag_authorization_cache[dag_id] = get_auth_manager().is_authorized_dag(
                    method=method,
                    access_entity=DagAccessEntity.TASK_INSTANCE,
                    details=DagDetails(id=dag_id, team_name=team_name),
                    user=self.user,
                )
            if not dag_authorization_cache[dag_id]:
                results.errors.append(
                    {
                        "error": f"User is not authorized to {action_name} task instances for DAG '{dag_id}'",
                        "status_code": status.HTTP_403_FORBIDDEN,
                    }
                )
                continue

            # Separate logic for "update all" vs "update specific"
            if map_index is not None:
                specific_map_index_task_keys.add((dag_id, dag_run_id, task_id, map_index))
            else:
                all_map_index_task_keys.add((dag_id, dag_run_id, task_id))

        return specific_map_index_task_keys, all_map_index_task_keys

    def _categorize_task_instances(
        self, task_keys: set[tuple[str, str, str, int]]
    ) -> tuple[
        dict[tuple[str, str, str, int], TI], set[tuple[str, str, str, int]], set[tuple[str, str, str, int]]
    ]:
        """
        Categorize the given task_keys into matched and not_found based on existing task instances.

        :param task_keys: set of task_keys (tuple of dag_id, dag_run_id, task_id, and map_index)
        :return: tuple of (task_instances_map, matched_task_keys, not_found_task_keys)
        """
        # Filter at database level using exact tuple matching instead of fetching all combinations
        # and filtering in Python
        task_keys_list = list(task_keys)
        query = select(TI).where(tuple_(TI.dag_id, TI.run_id, TI.task_id, TI.map_index).in_(task_keys_list))

        task_instances = self.session.scalars(query).all()
        task_instances_map = {
            (ti.dag_id, ti.run_id, ti.task_id, ti.map_index if ti.map_index is not None else -1): ti
            for ti in task_instances
        }
        matched_task_keys = set(task_instances_map.keys())
        not_found_task_keys = task_keys - matched_task_keys
        return task_instances_map, matched_task_keys, not_found_task_keys

    def _perform_update(
        self,
        entity: BulkTaskInstanceBody,
        dag_id: str,
        dag_run_id: str,
        task_id: str,
        map_index: int,
        results: BulkActionResponse,
        update_mask: list[str] | None = None,
    ) -> None:
        dag, tis, data = _patch_ti_validate_request(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            dag_bag=self.dag_bag,
            body=entity,
            session=self.session,
            map_index=map_index,
            update_mask=update_mask,
        )

        for key, _ in data.items():
            if key == "new_state":
                _patch_task_instance_state(
                    task_id=task_id,
                    dag_run_id=dag_run_id,
                    dag=dag,
                    task_instance_body=entity,
                    session=self.session,
                    data=data,
                )
            elif key == "note":
                _patch_task_instance_note(
                    task_instance_body=entity,
                    tis=tis,
                    user=self.user,
                )

        results.success.append(f"{dag_id}.{dag_run_id}.{task_id}[{map_index}]")

    def handle_bulk_create(
        self, action: BulkCreateAction[BulkTaskInstanceBody], results: BulkActionResponse
    ) -> None:
        results.errors.append(
            {
                "error": "Task instances bulk create is not supported",
                "status_code": status.HTTP_405_METHOD_NOT_ALLOWED,
            }
        )

    def handle_bulk_update(
        self, action: BulkUpdateAction[BulkTaskInstanceBody], results: BulkActionResponse
    ) -> None:
        """Bulk Update Task Instances."""
        # Validate and categorize entities into specific and all map index update sets
        update_specific_map_index_task_keys, update_all_map_index_task_keys = self._categorize_entities(
            action.entities, results, method="PUT", action_name=action.action.value
        )

        try:
            specific_entity_map = {
                self._extract_task_identifiers(entity): entity
                for entity in action.entities
                if entity.map_index is not None
            }
            all_map_entity_map = {
                self._extract_task_identifiers(entity)[:3]: entity
                for entity in action.entities
                if entity.map_index is None
            }

            # Handle updates for specific map_index task instances
            if update_specific_map_index_task_keys:
                _, matched_task_keys, not_found_task_keys = self._categorize_task_instances(
                    update_specific_map_index_task_keys
                )

                if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_task_keys:
                    not_found_task_ids = [
                        {"dag_id": dag_id, "dag_run_id": run_id, "task_id": task_id, "map_index": map_index}
                        for dag_id, run_id, task_id, map_index in not_found_task_keys
                    ]
                    raise HTTPException(
                        status_code=status.HTTP_404_NOT_FOUND,
                        detail=f"The task instances with these identifiers: {not_found_task_ids} were not found",
                    )

                for dag_id, dag_run_id, task_id, map_index in matched_task_keys:
                    entity = specific_entity_map.get((dag_id, dag_run_id, task_id, map_index))

                    if entity is not None:
                        self._perform_update(
                            dag_id=dag_id,
                            dag_run_id=dag_run_id,
                            task_id=task_id,
                            map_index=map_index,
                            entity=entity,
                            results=results,
                            update_mask=action.update_mask,
                        )

            # Handle updates for all map indexes
            if update_all_map_index_task_keys:
                all_dag_ids = {dag_id for dag_id, _, _ in update_all_map_index_task_keys}
                all_run_ids = {run_id for _, run_id, _ in update_all_map_index_task_keys}
                all_task_ids = {task_id for _, _, task_id in update_all_map_index_task_keys}

                batch_task_instances = self.session.scalars(
                    select(TI).where(
                        TI.dag_id.in_(all_dag_ids),
                        TI.run_id.in_(all_run_ids),
                        TI.task_id.in_(all_task_ids),
                    )
                ).all()

                # Group task instances by (dag_id, run_id, task_id)
                task_instances_by_key: dict[tuple[str, str, str], list[TI]] = {}
                for ti in batch_task_instances:
                    key = (ti.dag_id, ti.run_id, ti.task_id)
                    task_instances_by_key.setdefault(key, []).append(ti)

                for dag_id, run_id, task_id in update_all_map_index_task_keys:
                    all_task_instances = task_instances_by_key.get((dag_id, run_id, task_id), [])

                    if (
                        not all_task_instances
                        and action.action_on_non_existence == BulkActionNotOnExistence.FAIL
                    ):
                        raise HTTPException(
                            status_code=status.HTTP_404_NOT_FOUND,
                            detail=f"No task instances found for dag_id: {dag_id}, run_id: {run_id}, task_id: {task_id}",
                        )

                    entity = all_map_entity_map.get((dag_id, run_id, task_id))

                    if entity is not None:
                        for ti in all_task_instances:
                            self._perform_update(
                                dag_id=dag_id,
                                dag_run_id=run_id,
                                task_id=task_id,
                                map_index=ti.map_index if ti.map_index is not None else -1,
                                entity=entity,
                                results=results,
                                update_mask=action.update_mask,
                            )

        except ValidationError as e:
            results.errors.append({"error": f"{e.errors()}"})
        except HTTPException as e:
            results.errors.append({"error": f"{e.detail}", "status_code": e.status_code})

    def handle_bulk_delete(
        self, action: BulkDeleteAction[BulkTaskInstanceBody], results: BulkActionResponse
    ) -> None:
        """Bulk delete task instances."""
        # Validate and categorize entities into specific and all map index delete sets
        delete_specific_map_index_task_keys, delete_all_map_index_task_keys = self._categorize_entities(
            action.entities, results, method="DELETE", action_name=action.action.value
        )

        try:
            # Handle deletion of specific (dag_id, dag_run_id, task_id, map_index) tuples
            if delete_specific_map_index_task_keys:
                task_instances_map, matched_task_keys, not_found_task_keys = self._categorize_task_instances(
                    delete_specific_map_index_task_keys
                )
                not_found_task_ids = [
                    {"dag_id": dag_id, "dag_run_id": run_id, "task_id": task_id, "map_index": map_index}
                    for dag_id, run_id, task_id, map_index in not_found_task_keys
                ]

                if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_task_keys:
                    raise HTTPException(
                        status_code=status.HTTP_404_NOT_FOUND,
                        detail=f"The task instances with these identifiers: {not_found_task_ids} were not found",
                    )

                for task_key in matched_task_keys:
                    dag_id, run_id, task_id, map_index = task_key
                    self.session.delete(task_instances_map[task_key])
                    results.success.append(f"{dag_id}.{run_id}.{task_id}[{map_index}]")

            # Handle deletion of all map indexes for certain (dag_id, dag_run_id, task_id) tuples
            if delete_all_map_index_task_keys:
                all_dag_ids = {dag_id for dag_id, _, _ in delete_all_map_index_task_keys}
                all_run_ids = {run_id for _, run_id, _ in delete_all_map_index_task_keys}
                all_task_ids = {task_id for _, _, task_id in delete_all_map_index_task_keys}

                batch_task_instances = self.session.scalars(
                    select(TI).where(
                        TI.dag_id.in_(all_dag_ids),
                        TI.run_id.in_(all_run_ids),
                        TI.task_id.in_(all_task_ids),
                    )
                ).all()

                # Group task instances by (dag_id, run_id, task_id) for efficient lookup
                task_instances_by_key: dict[tuple[str, str, str], list[TI]] = {}
                for ti in batch_task_instances:
                    key = (ti.dag_id, ti.run_id, ti.task_id)
                    task_instances_by_key.setdefault(key, []).append(ti)

                for dag_id, run_id, task_id in delete_all_map_index_task_keys:
                    all_task_instances = task_instances_by_key.get((dag_id, run_id, task_id), [])

                    if (
                        not all_task_instances
                        and action.action_on_non_existence == BulkActionNotOnExistence.FAIL
                    ):
                        raise HTTPException(
                            status_code=status.HTTP_404_NOT_FOUND,
                            detail=f"No task instances found for dag_id: {dag_id}, run_id: {run_id}, task_id: {task_id}",
                        )

                    for ti in all_task_instances:
                        self.session.delete(ti)
                        results.success.append(f"{dag_id}.{run_id}.{task_id}[{ti.map_index}]")

        except HTTPException as e:
            results.errors.append({"error": f"{e.detail}", "status_code": e.status_code})
