package mcp.server.zap.core.exception;

import lombok.extern.slf4j.Slf4j;
import mcp.gateway.core.policy.ToolPolicyDeniedException;
import mcp.server.zap.core.logging.RequestCorrelationHolder;
import mcp.server.zap.core.logging.RequestLogContext;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.bind.support.WebExchangeBindException;
import org.springframework.web.reactive.resource.NoResourceFoundException;
import org.springframework.web.server.MissingRequestValueException;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.ServerWebInputException;

import java.time.Instant;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Global exception handler for the MCP ZAP Server.
 * Provides consistent error responses and prevents stack trace leakage.
 */
@Slf4j
@RestControllerAdvice
public class GlobalExceptionHandler {

    /**
     * Handle ZAP API exceptions.
     */
    @ExceptionHandler(ZapApiException.class)
    public ResponseEntity<Map<String, Object>> handleZapApiException(ZapApiException ex, ServerWebExchange exchange) {
        log.error("ZAP API error: {}", ex.getMessage(), ex);
        return buildErrorResponse(exchange, HttpStatus.SERVICE_UNAVAILABLE, "ZAP API Error", ex.getMessage());
    }

    /**
     * Handle validation exceptions.
     */
    @ExceptionHandler(WebExchangeBindException.class)
    public ResponseEntity<Map<String, Object>> handleValidationException(WebExchangeBindException ex, ServerWebExchange exchange) {
        String errors = ex.getBindingResult().getFieldErrors().stream()
                .map(error -> error.getField() + ": " + error.getDefaultMessage())
                .collect(Collectors.joining(", "));
        log.warn("Validation error: {}", errors);
        return buildErrorResponse(exchange, HttpStatus.BAD_REQUEST, "Validation Error", errors);
    }

    /**
     * Handle malformed/missing request inputs (headers/query/body values).
     */
    @ExceptionHandler({MissingRequestValueException.class, ServerWebInputException.class})
    public ResponseEntity<Map<String, Object>> handleRequestInputException(Exception ex, ServerWebExchange exchange) {
        log.warn("Request input error: {}", ex.getMessage());
        String message = ex.getMessage() != null ? ex.getMessage() : "Invalid request input";
        return buildErrorResponse(exchange, HttpStatus.BAD_REQUEST, "Invalid Request", message);
    }

    /**
     * Handle illegal argument exceptions (URL validation, etc.).
     */
    @ExceptionHandler(IllegalArgumentException.class)
    public ResponseEntity<Map<String, Object>> handleIllegalArgumentException(IllegalArgumentException ex, ServerWebExchange exchange) {
        log.warn("Invalid argument: {}", ex.getMessage());
        return buildErrorResponse(exchange, HttpStatus.BAD_REQUEST, "Invalid Request", ex.getMessage());
    }

    /**
     * Handle illegal state exceptions.
     */
    @ExceptionHandler(IllegalStateException.class)
    public ResponseEntity<Map<String, Object>> handleIllegalStateException(IllegalStateException ex, ServerWebExchange exchange) {
        log.error("Illegal state: {}", ex.getMessage(), ex);
        return buildErrorResponse(
                exchange,
                HttpStatus.INTERNAL_SERVER_ERROR,
                "Server Error",
                "The server could not complete the request."
        );
    }

    /**
     * Handle runtime policy denials as explicit forbidden responses.
     */
    @ExceptionHandler(ToolPolicyDeniedException.class)
    public ResponseEntity<Map<String, Object>> handleToolPolicyDeniedException(ToolPolicyDeniedException ex,
                                                                               ServerWebExchange exchange) {
        log.warn("Tool policy denied request: {}", ex.getMessage());
        return buildErrorResponse(exchange, HttpStatus.FORBIDDEN, "Policy Denied", ex.getMessage());
    }

    /**
     * Preserve 404 semantics for missing routes/static resources.
     */
    @ExceptionHandler(NoResourceFoundException.class)
    public ResponseEntity<Map<String, Object>> handleNoResourceFoundException(NoResourceFoundException ex,
                                                                              ServerWebExchange exchange) {
        log.warn("Resource not found: {}", ex.getMessage());
        return buildErrorResponse(exchange, HttpStatus.NOT_FOUND, "Not Found", ex.getMessage());
    }

    /**
     * Handle all other unexpected exceptions.
     * Note: Does not expose stack traces in production.
     */
    @ExceptionHandler(Exception.class)
    public ResponseEntity<Map<String, Object>> handleGenericException(Exception ex, ServerWebExchange exchange) {
        log.error("Unexpected error: {}", ex.getMessage(), ex);
        return buildErrorResponse(
                exchange,
                HttpStatus.INTERNAL_SERVER_ERROR,
                "Internal Server Error",
                "An unexpected error occurred. Please try again later."
        );
    }

    /**
     * Build a consistent error response.
     */
    private ResponseEntity<Map<String, Object>> buildErrorResponse(ServerWebExchange exchange,
                                                                  HttpStatus status,
                                                                  String error,
                                                                  String message) {
        Map<String, Object> body = new HashMap<>();
        body.put("timestamp", Instant.now().toString());
        body.put("status", status.value());
        body.put("error", error);
        body.put("message", message);
        String correlationId = exchange != null ? RequestLogContext.correlationId(exchange) : null;
        body.put("correlationId", correlationId != null ? correlationId : RequestCorrelationHolder.currentCorrelationId());
        return ResponseEntity.status(status).body(body);
    }
}
