"""
Functions to create audit logs for LiteLLM Proxy
"""

import asyncio
import json
from datetime import datetime, timezone
from typing import Dict

import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import (
    AUDIT_ACTIONS,
    LiteLLM_AuditLogs,
    LitellmTableNames,
    Optional,
    UserAPIKeyAuth,
)
from litellm.repositories.table_repositories import AuditLogRepository
from litellm.types.utils import StandardAuditLogPayload

_audit_log_callback_cache: Dict[str, CustomLogger] = {}
ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY = "allow_litellm_changed_by_header"


def _allows_litellm_changed_by_header(user_api_key_dict: UserAPIKeyAuth) -> bool:
    for admin_metadata in (user_api_key_dict.metadata, user_api_key_dict.team_metadata):
        if (
            isinstance(admin_metadata, dict)
            and admin_metadata.get(ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY) is True
        ):
            return True
    return False


def get_audit_log_changed_by(
    *,
    litellm_changed_by: Optional[str],
    user_api_key_dict: UserAPIKeyAuth,
    litellm_proxy_admin_name: Optional[str],
) -> Optional[str]:
    if litellm_changed_by and _allows_litellm_changed_by_header(user_api_key_dict):
        return litellm_changed_by
    return user_api_key_dict.user_id or litellm_proxy_admin_name


def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]:
    """Resolve a string callback name to a CustomLogger instance, with caching.

    For "s3_v2" with `litellm.s3_audit_callback_params` set, constructs a
    dedicated `S3Logger` so audit logs can target a different bucket than the
    normal-log singleton served by `_init_custom_logger_compatible_class`.
    """
    if name in _audit_log_callback_cache:
        return _audit_log_callback_cache[name]

    instance: Optional[CustomLogger]
    if (
        name == "s3_v2"
        and getattr(litellm, "s3_audit_callback_params", None) is not None
    ):
        from litellm.integrations.s3_v2 import S3Logger as S3V2Logger

        instance = S3V2Logger(
            s3_callback_params_override=litellm.s3_audit_callback_params
        )
    else:
        from litellm.litellm_core_utils.litellm_logging import (
            _init_custom_logger_compatible_class,
        )

        instance = _init_custom_logger_compatible_class(
            logging_integration=name,  # type: ignore
            internal_usage_cache=None,
            llm_router=None,
        )

    if instance is not None:
        _audit_log_callback_cache[name] = instance
    return instance


def reset_audit_log_callback_cache() -> None:
    """Clear cached audit-log callback instances. Call on config reload."""
    _audit_log_callback_cache.clear()


def _build_audit_log_payload(
    request_data: LiteLLM_AuditLogs,
) -> StandardAuditLogPayload:
    """Convert LiteLLM_AuditLogs to StandardAuditLogPayload for callback dispatch."""
    updated_at = ""
    if request_data.updated_at is not None:
        updated_at = request_data.updated_at.isoformat()

    table_name_str: str = (
        request_data.table_name.value
        if isinstance(request_data.table_name, LitellmTableNames)
        else str(request_data.table_name)
    )

    return StandardAuditLogPayload(
        id=request_data.id,
        updated_at=updated_at,
        changed_by=request_data.changed_by or "",
        changed_by_api_key=request_data.changed_by_api_key or "",
        action=request_data.action,
        table_name=table_name_str,
        object_id=request_data.object_id,
        before_value=request_data.before_value,
        updated_values=request_data.updated_values,
    )


def _audit_log_task_done_callback(task: asyncio.Task) -> None:
    """Log exceptions from audit log callback tasks so they don't slip through silently."""
    try:
        exc = task.exception()
    except asyncio.CancelledError:
        return
    if exc is not None:
        verbose_proxy_logger.error(
            "Audit log callback task failed: %s", exc, exc_info=exc
        )


async def _dispatch_audit_log_to_callbacks(
    request_data: LiteLLM_AuditLogs,
) -> None:
    """Dispatch audit log to all registered audit_log_callbacks."""
    if not litellm.audit_log_callbacks:
        return

    payload = _build_audit_log_payload(request_data)

    for callback in litellm.audit_log_callbacks:
        try:
            resolved: Optional[CustomLogger] = (
                callback if isinstance(callback, CustomLogger) else None
            )
            if isinstance(callback, str):
                resolved = _resolve_audit_log_callback(callback)
                if resolved is None:
                    verbose_proxy_logger.warning(
                        "Could not resolve audit log callback: %s", callback
                    )
                    continue

            if isinstance(resolved, CustomLogger):
                task = asyncio.create_task(resolved.async_log_audit_log_event(payload))
                task.add_done_callback(_audit_log_task_done_callback)
        except Exception as e:
            verbose_proxy_logger.error(
                "Failed dispatching audit log to callback: %s", e
            )


async def create_object_audit_log(
    object_id: str,
    action: AUDIT_ACTIONS,
    litellm_changed_by: Optional[str],
    user_api_key_dict: UserAPIKeyAuth,
    litellm_proxy_admin_name: Optional[str],
    table_name: LitellmTableNames,
    before_value: Optional[str] = None,
    after_value: Optional[str] = None,
):
    """
    Create an audit log for an internal user.

    Parameters:
    - user_id: str - The id of the user to create the audit log for.
    - action: AUDIT_ACTIONS - The action to create the audit log for.
    - user_row: LiteLLM_UserTable - The user row to create the audit log for.
    - litellm_changed_by: Optional[str] - The user id of the user who is changing the user.
    - user_api_key_dict: UserAPIKeyAuth - The user api key dictionary.
    - litellm_proxy_admin_name: Optional[str] - The name of the proxy admin.
    """
    from litellm.secret_managers.main import get_secret_bool

    _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool(
        "LITELLM_STORE_AUDIT_LOGS"
    )

    if _store_audit_logs is not True:
        return

    _changed_by = get_audit_log_changed_by(
        litellm_changed_by=litellm_changed_by,
        user_api_key_dict=user_api_key_dict,
        litellm_proxy_admin_name=litellm_proxy_admin_name,
    )

    await create_audit_log_for_update(
        request_data=LiteLLM_AuditLogs(
            id=str(uuid.uuid4()),
            updated_at=datetime.now(timezone.utc),
            changed_by=_changed_by,
            changed_by_api_key=user_api_key_dict.api_key,
            table_name=table_name,
            object_id=object_id,
            action=action,
            updated_values=after_value,
            before_value=before_value,
        )
    )


async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
    """
    Create an audit log for an object.
    """
    from litellm.secret_managers.main import get_secret_bool

    _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool(
        "LITELLM_STORE_AUDIT_LOGS"
    )
    if _store_audit_logs is not True:
        return

    from litellm.proxy.proxy_server import premium_user, prisma_client

    if premium_user is not True:
        return

    verbose_proxy_logger.debug("creating audit log for %s", request_data)

    if isinstance(request_data.updated_values, dict):
        request_data.updated_values = json.dumps(request_data.updated_values)

    if isinstance(request_data.before_value, dict):
        request_data.before_value = json.dumps(request_data.before_value)

    # Dispatch to external audit log callbacks regardless of DB availability
    await _dispatch_audit_log_to_callbacks(request_data)

    if prisma_client is None:
        verbose_proxy_logger.error(
            "prisma_client is None, cannot write audit log to DB"
        )
        return

    _request_data = request_data.model_dump(exclude_none=True)

    try:
        await AuditLogRepository(prisma_client).table.create(
            data={
                **_request_data,  # type: ignore
            }
        )
    except Exception as e:
        # [Non-Blocking Exception. Do not allow blocking LLM API call]
        verbose_proxy_logger.error(f"Failed Creating audit log {e}")
