"""Audit logging middleware for compliance and security."""

import time
import json
import logging
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional
from datetime import datetime, timezone
from pathlib import Path
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware

LOGGER = logging.getLogger(__name__)


class AuditLogWriter(ABC):
    """Abstract interface for audit log writing."""
    
    @abstractmethod
    async def write(self, entry: Dict[str, Any]) -> None:
        """Write audit log entry.
        
        Args:
            entry: Audit log entry dictionary
        """
        pass


class FileAuditLogWriter(AuditLogWriter):
    """Write audit logs to file (JSONL format)."""
    
    def __init__(self, log_file: Path):
        """Initialize file audit log writer.
        
        Args:
            log_file: Path to audit log file
        """
        self.log_file = log_file
        self.log_file.parent.mkdir(parents=True, exist_ok=True)
    
    async def write(self, entry: Dict[str, Any]) -> None:
        """Write audit log entry to file.
        
        Args:
            entry: Audit log entry dictionary
        """
        with open(self.log_file, "a") as f:
            f.write(json.dumps(entry) + "\n")


class LoggerAuditLogWriter(AuditLogWriter):
    """Write audit logs using Python logging."""
    
    def __init__(self, logger_name: str = "azure_discovery.audit"):
        """Initialize logger audit log writer.
        
        Args:
            logger_name: Name of logger to use
        """
        self.logger = logging.getLogger(logger_name)
    
    async def write(self, entry: Dict[str, Any]) -> None:
        """Write audit log entry using logger.
        
        Args:
            entry: Audit log entry dictionary
        """
        self.logger.info(
            "API Request",
            extra={"audit": entry},
        )


class AuditLogMiddleware(BaseHTTPMiddleware):
    """Log all API requests for security and compliance."""
    
    def __init__(
        self,
        app,
        writer: AuditLogWriter,
        include_request_body: bool = False,
        include_response_body: bool = False,
    ):
        """Initialize audit log middleware.
        
        Args:
            app: FastAPI application
            writer: Audit log writer implementation
            include_request_body: Whether to log request bodies
            include_response_body: Whether to log response bodies
        """
        super().__init__(app)
        self.writer = writer
        self.include_request_body = include_request_body
        self.include_response_body = include_response_body
    
    async def dispatch(self, request: Request, call_next):
        """Process request with audit logging.
        
        Args:
            request: Incoming HTTP request
            call_next: Next middleware in chain
            
        Returns:
            HTTP response
        """
        start_time = time.time()
        
        # Capture request details
        audit_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "method": request.method,
            "path": request.url.path,
            "query_params": dict(request.query_params),
            "client_ip": request.client.host if request.client else None,
            "user_agent": request.headers.get("user-agent"),
        }
        
        # Add authenticated user info if available
        if hasattr(request.state, "user"):
            audit_entry["user"] = {
                "id": request.state.user.id,
                "email": request.state.user.email,
                "source": request.state.user.source,
            }
        
        # Optionally capture request body
        if self.include_request_body and request.method in ["POST", "PUT", "PATCH"]:
            try:
                body = await request.body()
                # Avoid logging large bodies
                if len(body) < 10000:
                    audit_entry["request_body"] = body.decode("utf-8")
            except Exception as exc:
                LOGGER.debug(f"Failed to capture request body: {exc}")
        
        # Process request
        error: Optional[Exception] = None
        try:
            response = await call_next(request)
            audit_entry["status_code"] = response.status_code
        except Exception as exc:
            error = exc
            audit_entry["error"] = str(exc)
            audit_entry["error_type"] = type(exc).__name__
            raise
        finally:
            # Record duration
            audit_entry["duration_ms"] = (time.time() - start_time) * 1000
            
            # Write audit log
            try:
                await self.writer.write(audit_entry)
            except Exception as exc:
                LOGGER.error(f"Failed to write audit log: {exc}")
        
        return response
