"""写接口审计上下文工具。"""

from __future__ import annotations

from contextvars import ContextVar, Token
from dataclasses import dataclass, replace
from datetime import datetime
from functools import wraps
from inspect import iscoroutinefunction
from typing import Any, Awaitable, Callable, Optional, TypeVar, cast

from .timezone import CN_TZ


T = TypeVar("T")

VALID_CALLERS = {"system", "user", "ai"}


@dataclass(frozen=True)
class AuditRequestContext:
    trace_id: Optional[str] = None
    interface_path: Optional[str] = None
    caller: str = "system"
    method: Optional[str] = None
    tenant_id: Optional[int] = None
    user_id: Optional[int] = None


_audit_context: ContextVar[AuditRequestContext] = ContextVar(
    "pmflow_audit_request_context",
    default=AuditRequestContext(),
)


def normalize_caller(value: Optional[str]) -> str:
    normalized = str(value or "").strip().lower()
    return normalized if normalized in VALID_CALLERS else "system"


def get_audit_context() -> AuditRequestContext:
    return _audit_context.get()


def set_audit_context(context: AuditRequestContext) -> Token[AuditRequestContext]:
    return _audit_context.set(
        AuditRequestContext(
            trace_id=(str(context.trace_id)[:64] if context.trace_id else None),
            interface_path=(str(context.interface_path)[:255] if context.interface_path else None),
            caller=normalize_caller(context.caller),
            method=(str(context.method)[:16] if context.method else None),
            tenant_id=int(context.tenant_id) if context.tenant_id is not None else None,
            user_id=int(context.user_id) if context.user_id is not None else None,
        )
    )


def reset_audit_context(token: Token[AuditRequestContext]) -> None:
    _audit_context.reset(token)


def update_audit_context(
    *,
    trace_id: Optional[str] = None,
    interface_path: Optional[str] = None,
    caller: Optional[str] = None,
    method: Optional[str] = None,
    tenant_id: Optional[int] = None,
    user_id: Optional[int] = None,
) -> AuditRequestContext:
    current = get_audit_context()
    next_ctx = replace(
        current,
        trace_id=current.trace_id if trace_id is None else trace_id,
        interface_path=current.interface_path if interface_path is None else interface_path,
        caller=current.caller if caller is None else normalize_caller(caller),
        method=current.method if method is None else method,
        tenant_id=current.tenant_id if tenant_id is None else int(tenant_id),
        user_id=current.user_id if user_id is None else int(user_id),
    )
    _audit_context.set(next_ctx)
    return next_ctx


def format_beijing_datetime(dt: Optional[datetime]) -> Optional[str]:
    if dt is None:
        return None
    localized = dt.replace(tzinfo=CN_TZ) if dt.tzinfo is None else dt.astimezone(CN_TZ)
    return localized.strftime("%Y-%m-%d %H:%M:%S")


def audited_write(*, caller: Optional[str] = None) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """写接口装饰器：覆写当前请求 caller 审计上下文。"""

    normalized = normalize_caller(caller)

    def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        if not iscoroutinefunction(func):
            return func

        @wraps(func)
        async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
            current = get_audit_context()
            token = set_audit_context(replace(current, caller=normalized))
            try:
                return await cast(Callable[..., Awaitable[Any]], func)(*args, **kwargs)
            finally:
                reset_audit_context(token)

        return _async_wrapper

    return _decorator