"""Logging middleware for request/response tracking."""

from __future__ import annotations

import time
from typing import Any, Dict
import uuid

from ..logging import get_logger
from .base import Middleware

logger = get_logger(__name__)


class LoggingMiddleware(Middleware):
    """Middleware for logging all requests and responses.
    
    Logs request details, execution time, and response status.
    Adds request IDs for correlation.
    """
    
    def __init__(self, log_request_body: bool = False, log_response_body: bool = False):
        """Initialize logging middleware.
        
        Args:
            log_request_body: Whether to log request arguments (may be large)
            log_response_body: Whether to log response data (may be large)
        """
        super().__init__("LoggingMiddleware")
        self.log_request_body = log_request_body
        self.log_response_body = log_response_body
    
    async def before_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Log incoming request and add request ID.
        
        Args:
            request: The request data
            
        Returns:
            Request with added metadata
        """
        # Add request ID if not present
        if "request_id" not in request:
            request["request_id"] = str(uuid.uuid4())
        
        # Add timestamp
        request["start_time"] = time.time()
        
        # Log request
        log_data = {
            "request_id": request["request_id"],
            "tool_name": request.get("tool_name"),
        }
        
        if self.log_request_body:
            log_data["arguments"] = request.get("arguments")
        
        logger.info("mcp_request_received", **log_data)
        
        return request
    
    async def after_response(self, response: Any, request: Dict[str, Any]) -> Any:
        """Log response and execution time.
        
        Args:
            response: The response data
            request: The original request
            
        Returns:
            Unmodified response
        """
        # Calculate execution time
        start_time = request.get("start_time", time.time())
        execution_time = time.time() - start_time
        
        # Determine if response indicates success
        is_success = True
        if hasattr(response, "ok"):
            is_success = response.ok
        elif isinstance(response, dict):
            is_success = response.get("ok", True)
        
        # Log response
        log_data = {
            "request_id": request.get("request_id"),
            "tool_name": request.get("tool_name"),
            "execution_time_ms": round(execution_time * 1000, 2),
            "success": is_success,
        }
        
        if self.log_response_body:
            log_data["response"] = response
        
        logger.info("mcp_request_completed", **log_data)
        
        return response

