from typing import Annotated

from fastapi import APIRouter, Depends, status
from fastapi.responses import Response
from psycopg import Connection
from src.__types.index import Tenant
from src.routers.client.authorizer import get_client, get_tenant
from src.routers.client.connection import get_connection, set_session_variables
from src.usecases.quota_usecase import QuotaUseCase

router = APIRouter(
    prefix="/quota",
    tags=["クォータ"],
)

CHECK_DESCRIPTION = """
AIユニットの残量を確認し、生成AI機能の実行可否を判定します。

AIユニットの残量が`0より大きい`場合はHTTPステータスコード200を返します。

ステータスコードの確認の後、生成AIへのリクエストを実行してください。

### エラーハンドリング

| エラー内容 | 対応 |
| --- | --- |
| 4xxエラー （429以外） | 生成AI機能の処理を中断 |
| 429エラー, 5xxエラー, ネットワークエラー | リトライ戦略に従ってリトライ |

### リトライ戦略


| 観点 | 推奨される対応 |
| --- | --- |
| リトライ回数 | `3回` |
| リトライ間隔 | Exponential Backoff |
| リクエストタイムアウト | `10秒` |
| エラー時の対応 | リトライを経ても失敗する場合は、AIユニットの障害とみなし、機能を続行 |

### サンプルコード
```python

# --- 変数定義 ---
endpoint = "https://ai-unit.lab.works-hi.co.jp/v1/quota/check"
headers = {
    "X-Api-Key": "your-api-key",
    "X-Tenant-Name": "your-tenant-name",
}
max_retries = 3
wait_times = [1, 2, 4]

# --- リトライ処理 ---
try:
    for attempt in range(max_retries + 1):
        try:
            response = requests.head(endpoint, headers=headers, timeout=10)

            if response.status_code in [200]:
                logger.info("AIユニットの残量があります。生成AI機能を開始します。")
                break
            elif response.status_code in [400, 401, 403, 422]:
                raise AIUnitErrorException(f"AIユニットがクライアントエラーを返しました。")
            elif response.status_code in [404]:
                raise AIUnitErrorException(f"テナントが存在しません。")
            elif response.status_code in [402]:
                raise AIUnitQuotaExceededException("AIユニットの残量が0です。生成AI機能の処理を中断します。")
            elif response.status_code in [429]:
                logger.warning(f"AIユニットでレートリミットエラーが発生、リトライします。")
            elif 500 <= response.status_code <= 599:
                logger.warning(f"AIユニットでサーバーエラーが発生、リトライします。")
            else:
                raise AIUnitQuotaErrorException(f"AIユニットが予期せぬステータスコード {response.status_code} を返しました。")
        except requests.exceptions.RequestException as e:
            logger.warning("AIユニットへのリクエストでネットワークエラーが発生、リトライします")

        if attempt < max_retries:
            time.sleep(wait_times[attempt])
    else:
        logger.warning("AIユニットへのリトライが全て失敗しました。障害とみなし、処理を続行します。")

    client = boto3.client("bedrock-runtime")
    client.converse(
        model="your-model",
        messages=[{"role": "user", "content": "Hello, world!"}],
        inference_config=inference_config,
    )

except (AIUnitQuotaExceededException, AIUnitErrorException) as e:
    logger.error("生成AI機能の処理を中断します。")
```
"""


@router.head(
    "/check",
    summary="残ユニット確認",
    description=CHECK_DESCRIPTION,
    status_code=status.HTTP_200_OK,
    responses={
        status.HTTP_200_OK: {
            "description": "AIユニットの残量が0より大きい",
        },
        status.HTTP_401_UNAUTHORIZED: {
            "description": "認証失敗",
        },
        status.HTTP_402_PAYMENT_REQUIRED: {
            "description": "AIユニットの残量が0",
        },
        status.HTTP_403_FORBIDDEN: {
            "description": "権限の不足",
        },
        status.HTTP_404_NOT_FOUND: {
            "description": "テナントが存在しない",
        },
        status.HTTP_422_UNPROCESSABLE_CONTENT: {
            "description": "リクエストボディのバリデーションエラー",
        },
        status.HTTP_429_TOO_MANY_REQUESTS: {
            "description": "レートリミットエラー ※ リトライを推奨します。",
        },
        status.HTTP_500_INTERNAL_SERVER_ERROR: {
            "description": "内部エラー ※ AIユニットの障害発生時には機能の利用を続行してください。",
        },
    },
    response_class=Response,
    dependencies=[Depends(get_client)],
)
def head(
    tenant: Annotated[Tenant, Depends(get_tenant)],
    connection: Annotated[Connection, Depends(get_connection)],
) -> Response:
    set_session_variables(tenant=tenant, connection=connection)
    usecase = QuotaUseCase.create(
        tenant=tenant,
        connection=connection,
    )
    if not usecase.check():
        return Response(status_code=status.HTTP_402_PAYMENT_REQUIRED)
    return Response(status_code=status.HTTP_200_OK)
