"""Context management and decorators for audit logging integration."""

import uuid
import logging
from contextvars import ContextVar
from typing import Optional, Callable
from contextlib import asynccontextmanager

from .models import AuditLogEntry, AuditActor, AuditTarget, AuditChange
from .enums import AuditAction, AuditResource, AuditResult, AuditSeverity


# Context variable for request-scoped audit context
_audit_context: ContextVar[dict] = ContextVar("audit_context", default={})

# Module-level manager reference for automatic auth logging
_audit_manager = None

logger = logging.getLogger(__name__)


def init_audit_context(manager) -> None:
    """
    Initialize the audit context module with an AuditLogManager.

    Call this once at application startup to enable automatic auth logging
    when set_audit_context is called.

    Usage:
        from audit_logger import init_audit_context
        init_audit_context(audit_manager)
    """
    global _audit_manager
    _audit_manager = manager


def set_audit_context(
    org_id: str,
    user_id: Optional[str] = None,
    api_key: Optional[str] = None,
    ip_address: Optional[str] = None,
    request_id: Optional[str] = None,
) -> None:
    """
    Set audit context for the current request (call in middleware).

    If init_audit_context() was called with an AuditLogManager, this will
    automatically log an AUTH event for DAU/MAU tracking.

    Args:
        org_id: Organization ID
        user_id: Optional user ID
        api_key: Optional API key
        ip_address: Optional IP address
        request_id: Optional request ID (auto-generated if not provided)
        auto_log_auth: Whether to automatically log auth event (default True)
    """
    _audit_context.set({
        "org_id": org_id,
        "user_id": user_id,
        "api_key": api_key,
        "ip_address": ip_address,
        "request_id": request_id or str(uuid.uuid4()),
    })

    # Automatically log auth if manager is configured
    if _audit_manager is not None:
        _audit_manager.log_auth(
            org_id=org_id,
            user_id=user_id,
            api_key=api_key,
            ip_address=ip_address,
            resource_type=AuditResource.API,
        )


def get_audit_context() -> dict:
    """Get the current audit context."""
    return _audit_context.get()


def clear_audit_context() -> None:
    """Clear audit context (call at end of request)."""
    _audit_context.set({})


@asynccontextmanager
async def audit_context(
    org_id: str,
    user_id: Optional[str] = None,
    api_key: Optional[str] = None,
    ip_address: Optional[str] = None,
):
    """Context manager for setting audit context."""
    request_id = str(uuid.uuid4())
    set_audit_context(
        org_id=org_id,
        user_id=user_id,
        api_key=api_key,
        ip_address=ip_address,
        request_id=request_id,
    )
    try:
        yield request_id
    finally:
        clear_audit_context()


# def audit_action(
#     action: AuditAction,
#     resource_type: AuditResource,
#     resource_id_param: Optional[str] = None,  # Name of param containing resource_id
#     resource_name_param: Optional[str] = None,  # Name of param containing resource_name
#     get_changes: Optional[Callable] = None,  # Function to compute changes for updates
#     severity: AuditSeverity = AuditSeverity.LOW,
# ):
#     """
#     Decorator to automatically audit async functions.

#     Usage:
#         @audit_action(
#             action=AuditAction.CREATE,
#             resource_type=AuditResource.AGENT,
#             resource_id_param="agent_id",
#         )
#         async def create_agent(agent_id: str, name: str, ...):
#             ...
#     """
#     def decorator(func: Callable) -> Callable:
#         @functools.wraps(func)
#         async def wrapper(*args, **kwargs) -> Any:
#             # Get audit manager from global or kwargs
#             audit_manager = kwargs.pop("_audit_manager", None)
#             if audit_manager is None:
#                 # If no audit manager, just run the function
#                 return await func(*args, **kwargs)

#             ctx = get_audit_context()
#             if not ctx.get("org_id"):
#                 # No audit context set, run without auditing
#                 return await func(*args, **kwargs)

#             # Extract resource identifiers from kwargs
#             resource_id = kwargs.get(resource_id_param) if resource_id_param else None
#             resource_name = kwargs.get(resource_name_param) if resource_name_param else None

#             # Build actor
#             actor = AuditActor(
#                 org_id=ctx["org_id"],
#                 user_id=ctx.get("user_id"),
#                 api_key_hash=audit_manager._hash_api_key(ctx.get("api_key")) if ctx.get("api_key") else None,
#                 ip_address=ctx.get("ip_address"),
#             )

#             # Build target
#             target = AuditTarget(
#                 resource_type=resource_type,
#                 resource_id=resource_id,
#                 resource_name=resource_name,
#             )

#             # Compute changes for updates if provided
#             changes = None
#             if get_changes and action == AuditAction.UPDATE:
#                 try:
#                     changes = await get_changes(*args, **kwargs)
#                 except Exception:
#                     pass

#             # Execute the function
#             result = AuditResult.SUCCESS
#             error_message = None
#             return_value = None

#             try:
#                 return_value = await func(*args, **kwargs)

#                 # Try to extract resource_id from return value if not in kwargs
#                 if not resource_id and isinstance(return_value, dict):
#                     resource_id = return_value.get("id") or return_value.get("agent_id") or return_value.get("_id")
#                     target.resource_id = resource_id

#             except PermissionError as e:
#                 result = AuditResult.BLOCKED
#                 error_message = str(e)
#                 raise
#             except Exception as e:
#                 result = AuditResult.FAILURE
#                 error_message = str(e)
#                 raise
#             finally:
#                 # Log the audit entry (non-blocking)
#                 entry = AuditLogEntry(
#                     actor=actor,
#                     action=action,
#                     target=target,
#                     result=result,
#                     error_message=error_message,
#                     changes=changes,
#                     request_id=ctx.get("request_id"),
#                     severity=severity,
#                 )
#                 try:
#                     audit_manager.log_buffered(entry)  # Non-blocking
#                 except Exception:
#                     pass  # Don't fail the request if audit logging fails

#             return return_value

#         return wrapper
#     return decorator


def compute_changes(old_data: dict, new_data: dict, fields: list[str]) -> list[AuditChange]:
    """
    Helper to compute changes between old and new data.

    Usage:
        changes = compute_changes(
            old_data=existing_agent,
            new_data=update_payload,
            fields=["name", "description", "system_prompt", "tools"]
        )
    """
    changes = []
    for field in fields:
        old_val = old_data.get(field)
        new_val = new_data.get(field)
        if old_val != new_val and new_val is not None:
            changes.append(AuditChange(
                field=field,
                old_value=old_val,
                new_value=new_val,
            ))
    return changes


class AuditLogger:
    """
    Simplified audit logger for manual logging in code (non-blocking).

    All methods are synchronous and return immediately - writes happen
    in background tasks that don't block the main thread.

    Usage:
        audit = AuditLogger(audit_manager)
        audit.log_create(AuditResource.AGENT, agent_id, name=agent_name)
    """

    def __init__(self, manager):
        self.manager = manager

    def _get_actor_info(self) -> dict:
        """Get actor info from context."""
        ctx = get_audit_context()
        return {
            "org_id": ctx.get("org_id", "unknown"),
            "user_id": ctx.get("user_id"),
            "api_key": ctx.get("api_key"),
            "ip_address": ctx.get("ip_address"),
        }

    def log_create(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a create action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_create(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_update(
        self,
        resource_type: AuditResource,
        resource_id: str,
        changes: list[AuditChange],
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an update action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_update(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            changes=changes,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_delete(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a delete action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_delete(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_execute(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an execute action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_execute(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_parse(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a parse action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_parse(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_train(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a train action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_train(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_upload(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an upload action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_upload(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_download(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a download action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_download(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_reset(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a reset action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_reset(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_share(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a share action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_share(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )
    

    def log_sign_in(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a sign-in action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_sign_in(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_sign_out(
        self,
        resource_type: AuditResource,
        resource_id: str,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a sign-out action using context (non-blocking)."""
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_sign_out(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
            resource_id=resource_id,
            resource_name=resource_name,
            metadata=metadata,
        )

    def log_auth(
        self,
        resource_type: AuditResource,
    ) -> None:
        """
        Log a auth action using context (non-blocking)
        """
        actor = self._get_actor_info()
        if not actor["org_id"] or actor["org_id"] == "unknown":
            return
        self.manager.log_auth(
            org_id=actor["org_id"],
            user_id=actor["user_id"],
            api_key=actor["api_key"],
            ip_address=actor["ip_address"],
            resource_type=resource_type,
        )