from __future__ import annotations

from typing import Any

from fastapi import HTTPException
from sqlalchemy import text
from sqlalchemy.exc import DBAPIError, OperationalError
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings


class DatabaseRowLockBusy(HTTPException, RuntimeError):
    """A short-lived PostgreSQL row/table lock could not be acquired in time."""

    def __init__(self, message: str = "当前任务正在被其他流程处理，请稍后重试") -> None:
        super().__init__(status_code=409, detail=message)


_LOCK_NOT_AVAILABLE_SQLSTATE = "55P03"


def _sqlstate_from_exception(exc: BaseException | None) -> str | None:
    current: BaseException | None = exc
    seen: set[int] = set()
    while current is not None and id(current) not in seen:
        seen.add(id(current))
        for attr in ("sqlstate", "pgcode"):
            value = getattr(current, attr, None)
            if value:
                return str(value)
        current = getattr(current, "orig", None) or getattr(current, "__cause__", None)
    return None


def is_postgres_lock_timeout(exc: BaseException) -> bool:
    if _sqlstate_from_exception(exc) == _LOCK_NOT_AVAILABLE_SQLSTATE:
        return True
    message = str(exc).lower()
    return "lock timeout" in message or "could not obtain lock" in message


def raise_if_database_lock_busy(exc: BaseException) -> None:
    if is_postgres_lock_timeout(exc):
        raise DatabaseRowLockBusy("数据库任务行正在被其他事务处理，请稍后重试") from exc


async def apply_short_lock_timeout(
    db: AsyncSession,
    *,
    seconds: int | None = None,
) -> None:
    """Apply a transaction-local PostgreSQL lock wait limit.

    It deliberately does not change the global database configuration and is a
    no-op on non-PostgreSQL test/development databases.
    """
    bind = db.get_bind()
    if bind is None or bind.dialect.name != "postgresql":
        return
    timeout_seconds = max(
        1,
        int(
            seconds
            if seconds is not None
            else getattr(settings, "GENERATION_DB_LOCK_TIMEOUT_SECONDS", 5)
            or 5
        ),
    )
    await db.execute(text(f"SET LOCAL lock_timeout = '{timeout_seconds}s'"))


async def execute_with_lock_timeout(
    db: AsyncSession,
    statement: Any,
    *,
    seconds: int | None = None,
):
    await apply_short_lock_timeout(db, seconds=seconds)
    try:
        return await db.execute(statement)
    except (OperationalError, DBAPIError) as exc:
        raise_if_database_lock_busy(exc)
        raise