# 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 typing import Annotated

from fastapi import Depends, HTTPException, status
from sqlalchemy import select

from airflow._shared.state import AssetScope
from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.api_fastapi.core_api.datamodels.asset_state import (
    AssetStateBody,
    AssetStateCollectionResponse,
    AssetStateResponse,
)
from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
from airflow.api_fastapi.core_api.security import requires_access_asset
from airflow.models.asset import AssetModel
from airflow.models.asset_state import AssetStateModel
from airflow.state import get_state_backend

asset_state_router = AirflowRouter(
    tags=["Asset State"],
    prefix="/assets/{asset_id}/states",
)


def _get_asset_or_404(asset_id: int, session: SessionDep) -> int:
    exists = session.scalar(select(AssetModel.id).where(AssetModel.id == asset_id))
    if exists is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Asset with id {asset_id!r} not found",
        )
    return asset_id


AssetIdDep = Annotated[int, Depends(_get_asset_or_404)]


@asset_state_router.get(
    "",
    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
    dependencies=[Depends(requires_access_asset(method="GET"))],
)
def list_asset_states(
    asset_id: AssetIdDep,
    limit: QueryLimit,
    offset: QueryOffset,
    session: SessionDep,
) -> AssetStateCollectionResponse:
    """List all state entries for an asset."""
    base = (
        select(
            AssetStateModel.key,
            AssetStateModel.value,
            AssetStateModel.updated_at,
        )
        .where(AssetStateModel.asset_id == asset_id)
        .order_by(AssetStateModel.key.asc())
    )
    paginated, total_entries = paginated_select(
        statement=base,
        filters=None,
        order_by=None,
        offset=offset,
        limit=limit,
        session=session,
    )
    rows = session.execute(paginated).all()
    entries = [AssetStateResponse(key=r.key, value=r.value, updated_at=r.updated_at) for r in rows]
    return AssetStateCollectionResponse(asset_states=entries, total_entries=total_entries)


@asset_state_router.get(
    "/{key:path}",
    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
    dependencies=[Depends(requires_access_asset(method="GET"))],
)
def get_asset_state(
    asset_id: AssetIdDep,
    key: str,
    session: SessionDep,
) -> AssetStateResponse:
    """Get a single asset state entry."""
    row = session.execute(
        select(
            AssetStateModel.key,
            AssetStateModel.value,
            AssetStateModel.updated_at,
        ).where(
            AssetStateModel.asset_id == asset_id,
            AssetStateModel.key == key,
        )
    ).one_or_none()
    if row is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Asset state key {key!r} not found",
        )
    return AssetStateResponse(key=row.key, value=row.value, updated_at=row.updated_at)


@asset_state_router.put(
    "/{key:path}",
    status_code=status.HTTP_204_NO_CONTENT,
    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
    dependencies=[Depends(requires_access_asset(method="PUT"))],
)
def set_asset_state(
    asset_id: AssetIdDep,
    key: str,
    body: AssetStateBody,
    session: SessionDep,
) -> None:
    """Set an asset state value. Creates or overwrites the key."""
    get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value, session=session)


@asset_state_router.delete(
    "/{key:path}",
    status_code=status.HTTP_204_NO_CONTENT,
    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
    dependencies=[Depends(requires_access_asset(method="DELETE"))],
)
def delete_asset_state(
    asset_id: AssetIdDep,
    key: str,
    session: SessionDep,
) -> None:
    """Delete a single asset state key. No-op if the key does not exist."""
    get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)


@asset_state_router.delete(
    "",
    status_code=status.HTTP_204_NO_CONTENT,
    responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
    dependencies=[Depends(requires_access_asset(method="DELETE"))],
)
def clear_asset_state(
    asset_id: AssetIdDep,
    session: SessionDep,
) -> None:
    """Delete all state keys for an asset."""
    get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)
