from __future__ import annotations

from collections.abc import Iterator
from contextlib import contextmanager
from hashlib import sha256
from threading import Lock, RLock

from sqlalchemy import func, select
from sqlalchemy.orm import Session

_FALLBACK_LOCKS_GUARD = Lock()
_FALLBACK_LOCKS: dict[str, RLock] = {}


@contextmanager
def organization_memory_operation_locks(
    db: Session,
    *lock_keys: str,
) -> Iterator[None]:
    """串行化组织记忆作用域和幂等键，锁的生命周期覆盖当前事务。"""

    normalized_keys = sorted({str(key) for key in lock_keys if str(key)})
    bind = db.get_bind()
    if bind.dialect.name == "postgresql":
        for lock_key in normalized_keys:
            db.execute(
                select(
                    func.pg_advisory_xact_lock(
                        organization_memory_advisory_lock_id(lock_key)
                    )
                )
            )
        yield
        return

    # SQLite 等方言没有事务级 advisory lock。进程内锁配合数据库唯一索引
    # 提供安全退化，保证测试与单进程部署不会出现空集合竞态。
    fallback_locks = [_fallback_lock(lock_key) for lock_key in normalized_keys]
    for fallback_lock in fallback_locks:
        fallback_lock.acquire()
    try:
        yield
    finally:
        for fallback_lock in reversed(fallback_locks):
            fallback_lock.release()


def organization_memory_scope_lock_key(
    *,
    tenant_id: str,
    scope_type: str,
    scope_id: str,
    scene: str,
    field_key: str,
) -> str:
    return "|".join(
        (
            "organization-memory-scope",
            tenant_id,
            scope_type,
            scope_id,
            scene,
            field_key,
        )
    )


def organization_memory_request_lock_key(tenant_id: str, request_id: str) -> str:
    return f"organization-memory-request|{tenant_id}|{request_id}"


def _fallback_lock(lock_key: str) -> RLock:
    with _FALLBACK_LOCKS_GUARD:
        return _FALLBACK_LOCKS.setdefault(lock_key, RLock())


def organization_memory_advisory_lock_id(lock_key: str) -> int:
    return int.from_bytes(
        sha256(lock_key.encode("utf-8")).digest()[:8],
        byteorder="big",
        signed=True,
    )
