Exception Handling & Status Codes

from starlette.exceptions import HTTPException, WebSocketException
from starlette.responses import Response
from collections.abc import Mapping
from typing import Any

class HTTPException(Exception):
    """
    HTTP error exception with status code and optional details.
    
    Used to raise HTTP errors that will be converted to appropriate
    HTTP responses by exception middleware.
    """
    
    def __init__(
        self,
        status_code: int,
        detail: str | None = None,
        headers: Mapping[str, str] | None = None,
    ) -> None:
        """
        Initialize HTTP exception.
        
        Args:
            status_code: HTTP status code (400-599)
            detail: Error detail message or structured data
            headers: Additional HTTP headers for error response
        """
        self.status_code = status_code
        self.detail = detail
        self.headers = headers or {}
        
        # Set exception message
        if detail is None:
            super().__init__(f"HTTP {status_code}")
        else:
            super().__init__(f"HTTP {status_code}: {detail}")

from starlette.requests import ClientDisconnect

class ClientDisconnect(Exception):
    """
    Exception raised when client disconnects during request processing.
    
    This can occur during streaming request body reads or when 
    the client closes the connection unexpectedly.
    """
    pass

from starlette.websockets import WebSocketDisconnect
from starlette.exceptions import WebSocketException

class WebSocketDisconnect(Exception):
    """
    Exception raised when WebSocket connection closes.
    
    Contains close code and optional reason for disconnection.
    """
    
    def __init__(self, code: int = 1000, reason: str | None = None) -> None:
        """
        Initialize WebSocket disconnect exception.
        
        Args:
            code: WebSocket close code (1000 = normal closure)
            reason: Optional close reason string
        """
        self.code = code
        self.reason = reason
        
        super().__init__(f"WebSocket disconnect: {code}")

class WebSocketException(Exception):
    """
    WebSocket error exception with close code and reason.
    
    Used to signal WebSocket errors that should close the connection
    with a specific code and reason.
    """
    
    def __init__(self, code: int, reason: str | None = None) -> None:
        """
        Initialize WebSocket exception.
        
        Args:
            code: WebSocket close code
            reason: Optional close reason
        """
        self.code = code
        self.reason = reason
        
        super().__init__(f"WebSocket error: {code}")

from starlette import status

# 1xx Informational responses
HTTP_100_CONTINUE = 100
HTTP_101_SWITCHING_PROTOCOLS = 101  
HTTP_102_PROCESSING = 102
HTTP_103_EARLY_HINTS = 103

# 2xx Success responses
HTTP_200_OK = 200
HTTP_201_CREATED = 201
HTTP_202_ACCEPTED = 202
HTTP_203_NON_AUTHORITATIVE_INFORMATION = 203
HTTP_204_NO_CONTENT = 204
HTTP_205_RESET_CONTENT = 205
HTTP_206_PARTIAL_CONTENT = 206
HTTP_207_MULTI_STATUS = 207
HTTP_208_ALREADY_REPORTED = 208
HTTP_226_IM_USED = 226

# 3xx Redirection responses  
HTTP_300_MULTIPLE_CHOICES = 300
HTTP_301_MOVED_PERMANENTLY = 301
HTTP_302_FOUND = 302
HTTP_303_SEE_OTHER = 303
HTTP_304_NOT_MODIFIED = 304
HTTP_305_USE_PROXY = 305
HTTP_306_RESERVED = 306
HTTP_307_TEMPORARY_REDIRECT = 307
HTTP_308_PERMANENT_REDIRECT = 308

# 4xx Client error responses
HTTP_400_BAD_REQUEST = 400
HTTP_401_UNAUTHORIZED = 401
HTTP_402_PAYMENT_REQUIRED = 402
HTTP_403_FORBIDDEN = 403
HTTP_404_NOT_FOUND = 404
HTTP_405_METHOD_NOT_ALLOWED = 405
HTTP_406_NOT_ACCEPTABLE = 406
HTTP_407_PROXY_AUTHENTICATION_REQUIRED = 407
HTTP_408_REQUEST_TIMEOUT = 408
HTTP_409_CONFLICT = 409
HTTP_410_GONE = 410
HTTP_411_LENGTH_REQUIRED = 411
HTTP_412_PRECONDITION_FAILED = 412
HTTP_413_REQUEST_ENTITY_TOO_LARGE = 413
HTTP_414_REQUEST_URI_TOO_LONG = 414
HTTP_415_UNSUPPORTED_MEDIA_TYPE = 415
HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE = 416
HTTP_417_EXPECTATION_FAILED = 417
HTTP_418_IM_A_TEAPOT = 418
HTTP_421_MISDIRECTED_REQUEST = 421
HTTP_422_UNPROCESSABLE_ENTITY = 422
HTTP_423_LOCKED = 423
HTTP_424_FAILED_DEPENDENCY = 424
HTTP_425_TOO_EARLY = 425
HTTP_426_UPGRADE_REQUIRED = 426
HTTP_428_PRECONDITION_REQUIRED = 428
HTTP_429_TOO_MANY_REQUESTS = 429
HTTP_431_REQUEST_HEADER_FIELDS_TOO_LARGE = 431
HTTP_451_UNAVAILABLE_FOR_LEGAL_REASONS = 451

# 5xx Server error responses
HTTP_500_INTERNAL_SERVER_ERROR = 500
HTTP_501_NOT_IMPLEMENTED = 501
HTTP_502_BAD_GATEWAY = 502
HTTP_503_SERVICE_UNAVAILABLE = 503
HTTP_504_GATEWAY_TIMEOUT = 504
HTTP_505_HTTP_VERSION_NOT_SUPPORTED = 505
HTTP_506_VARIANT_ALSO_NEGOTIATES = 506
HTTP_507_INSUFFICIENT_STORAGE = 507
HTTP_508_LOOP_DETECTED = 508
HTTP_510_NOT_EXTENDED = 510
HTTP_511_NETWORK_AUTHENTICATION_REQUIRED = 511

# WebSocket close codes
WS_1000_NORMAL_CLOSURE = 1000
WS_1001_GOING_AWAY = 1001
WS_1002_PROTOCOL_ERROR = 1002
WS_1003_UNSUPPORTED_DATA = 1003
WS_1005_NO_STATUS_RCVD = 1005
WS_1006_ABNORMAL_CLOSURE = 1006
WS_1007_INVALID_FRAME_PAYLOAD_DATA = 1007
WS_1008_POLICY_VIOLATION = 1008
WS_1009_MESSAGE_TOO_BIG = 1009
WS_1010_MANDATORY_EXT = 1010
WS_1011_INTERNAL_ERROR = 1011
WS_1012_SERVICE_RESTART = 1012
WS_1013_TRY_AGAIN_LATER = 1013
WS_1014_BAD_GATEWAY = 1014
WS_1015_TLS_HANDSHAKE = 1015

from starlette.exceptions import HTTPException
from starlette import status
from starlette.responses import JSONResponse

async def get_user(request):
    user_id = request.path_params["user_id"]
    
    # Validate input
    try:
        user_id = int(user_id)
    except ValueError:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid user ID format"
        )
    
    # Check authorization
    if not request.user.is_authenticated:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Authentication required",
            headers={"WWW-Authenticate": "Bearer"}
        )
    
    # Find user
    user = await find_user_by_id(user_id)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"User {user_id} not found"
        )
    
    # Check permissions
    if user.id != request.user.id and not request.user.is_admin():
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied"
        )
    
    return JSONResponse({"user": user.to_dict()})

async def create_user(request):
    try:
        data = await request.json()
    except ValueError:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid JSON in request body"
        )
    
    # Validation errors with structured details
    errors = []
    
    if not data.get("email"):
        errors.append({"field": "email", "message": "Email is required"})
    elif not is_valid_email(data["email"]):
        errors.append({"field": "email", "message": "Invalid email format"})
    
    if not data.get("username"):
        errors.append({"field": "username", "message": "Username is required"})
    elif len(data["username"]) < 3:
        errors.append({"field": "username", "message": "Username too short"})
    
    if errors:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail={
                "message": "Validation failed",
                "errors": errors
            }
        )
    
    # Create user...
    user = await create_user_in_db(data)
    return JSONResponse(user.to_dict(), status_code=status.HTTP_201_CREATED)

class ValidationError(HTTPException):
    """Validation error with field-specific details."""
    
    def __init__(self, errors: list[dict]):
        super().__init__(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail={
                "message": "Validation failed",
                "errors": errors
            }
        )

class NotFoundError(HTTPException):
    """Resource not found error."""
    
    def __init__(self, resource: str, identifier: str | int):
        super().__init__(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"{resource} '{identifier}' not found"
        )

class PermissionDeniedError(HTTPException):
    """Permission denied error."""
    
    def __init__(self, action: str, resource: str = "resource"):
        super().__init__(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"Permission denied: cannot {action} {resource}"
        )

class RateLimitError(HTTPException):
    """Rate limit exceeded error."""
    
    def __init__(self, retry_after: int = 60):
        super().__init__(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail="Rate limit exceeded",
            headers={"Retry-After": str(retry_after)}
        )

# Usage
async def delete_user(request):
    user_id = request.path_params["user_id"]
    
    user = await find_user_by_id(user_id)
    if not user:
        raise NotFoundError("User", user_id)
    
    if not can_delete_user(request.user, user):
        raise PermissionDeniedError("delete", "user")
    
    await delete_user_from_db(user_id)
    return JSONResponse({"message": "User deleted"})

from starlette.websockets import WebSocket, WebSocketDisconnect
from starlette.exceptions import WebSocketException
from starlette import status as ws_status

async def websocket_endpoint(websocket: WebSocket):
    try:
        # Authentication check
        token = websocket.query_params.get("token")
        if not token:
            raise WebSocketException(
                code=ws_status.WS_1008_POLICY_VIOLATION,
                reason="Authentication token required"
            )
        
        user = await authenticate_token(token)
        if not user:
            raise WebSocketException(
                code=ws_status.WS_1008_POLICY_VIOLATION,
                reason="Invalid authentication token"
            )
        
        await websocket.accept()
        
        # Handle messages
        while True:
            try:
                data = await websocket.receive_json()
                
                # Validate message
                if not isinstance(data, dict):
                    await websocket.send_json({
                        "type": "error",
                        "message": "Message must be JSON object"
                    })
                    continue
                
                # Process message
                await process_websocket_message(websocket, user, data)
                
            except ValueError:
                # Invalid JSON
                await websocket.send_json({
                    "type": "error", 
                    "message": "Invalid JSON format"
                })
            
    except WebSocketDisconnect:
        print(f"Client disconnected")
    
    except WebSocketException as e:
        # Close with specific code and reason
        await websocket.close(code=e.code, reason=e.reason)
    
    except Exception as e:
        # Unexpected error - close with internal error code
        print(f"Unexpected error: {e}")
        await websocket.close(
            code=ws_status.WS_1011_INTERNAL_ERROR,
            reason="Internal server error"
        )

async def process_websocket_message(websocket: WebSocket, user, data):
    message_type = data.get("type")
    
    if message_type == "ping":
        await websocket.send_json({"type": "pong"})
    
    elif message_type == "chat_message":
        if not data.get("content"):
            await websocket.send_json({
                "type": "error",
                "message": "Message content required"
            })
            return
        
        # Process chat message...
        await broadcast_chat_message(user, data["content"])
    
    else:
        await websocket.send_json({
            "type": "error",
            "message": f"Unknown message type: {message_type}"
        })

from starlette.applications import Starlette
from starlette.middleware.exceptions import ExceptionMiddleware
from starlette.responses import JSONResponse, PlainTextResponse
from starlette.requests import Request

# Custom exception handlers
async def http_exception_handler(request: Request, exc: HTTPException):
    """Handle HTTPException instances."""
    return JSONResponse(
        content={
            "error": {
                "status_code": exc.status_code,
                "detail": exc.detail,
                "type": "HTTPException"
            }
        },
        status_code=exc.status_code,
        headers=exc.headers
    )

async def validation_error_handler(request: Request, exc: ValidationError):
    """Handle custom validation errors."""
    return JSONResponse(
        content={
            "error": {
                "type": "ValidationError",
                "message": "Request validation failed",
                "details": exc.detail
            }
        },
        status_code=exc.status_code
    )

async def generic_error_handler(request: Request, exc: Exception):
    """Handle unexpected exceptions."""
    print(f"Unexpected error: {exc}")
    
    return JSONResponse(
        content={
            "error": {
                "type": "InternalServerError",
                "message": "An unexpected error occurred"
            }
        },
        status_code=500
    )

async def not_found_handler(request: Request, exc):
    """Handle 404 errors."""
    return JSONResponse(
        content={
            "error": {
                "type": "NotFound",
                "message": "The requested resource was not found",
                "path": request.url.path
            }
        },
        status_code=404
    )

# Configure exception handlers
exception_handlers = {
    HTTPException: http_exception_handler,
    ValidationError: validation_error_handler,
    404: not_found_handler,  # By status code
    500: generic_error_handler,  # By status code
}

app = Starlette(
    routes=routes,
    exception_handlers=exception_handlers
)

# Add exception handlers dynamically
app.add_exception_handler(ValueError, validation_error_handler)
app.add_exception_handler(KeyError, generic_error_handler)

from starlette.config import Config
from starlette.responses import HTMLResponse
import traceback

config = Config()
DEBUG = config("DEBUG", cast=bool, default=False)

async def debug_exception_handler(request: Request, exc: Exception):
    """Detailed error handler for development."""
    error_html = f"""
    <html>
    <head><title>Error: {exc.__class__.__name__}</title></head>
    <body>
        <h1>Error: {exc.__class__.__name__}</h1>
        <p><strong>Message:</strong> {str(exc)}</p>
        <h2>Traceback:</h2>
        <pre>{traceback.format_exc()}</pre>
        <h2>Request Info:</h2>
        <pre>
Method: {request.method}
URL: {request.url}
Headers: {dict(request.headers)}
        </pre>
    </body>
    </html>
    """
    return HTMLResponse(error_html, status_code=500)

async def production_exception_handler(request: Request, exc: Exception):
    """Safe error handler for production."""
    # Log the error (use proper logging in production)
    print(f"ERROR: {exc}")
    
    return JSONResponse(
        content={
            "error": {
                "type": "InternalServerError",
                "message": "An internal error occurred"
            }
        },
        status_code=500
    )

# Choose handler based on environment
exception_handler = debug_exception_handler if DEBUG else production_exception_handler

app.add_exception_handler(Exception, exception_handler)

from starlette.middleware.base import BaseHTTPMiddleware
import logging

class ErrorLoggingMiddleware(BaseHTTPMiddleware):
    """Middleware to log exceptions with context."""
    
    async def dispatch(self, request, call_next):
        try:
            response = await call_next(request)
            return response
        
        except HTTPException as exc:
            # Log HTTP exceptions with request context
            logging.warning(
                f"HTTP {exc.status_code}: {exc.detail}",
                extra={
                    "method": request.method,
                    "path": request.url.path,
                    "client": request.client.host if request.client else None,
                    "user_agent": request.headers.get("user-agent"),
                }
            )
            raise  # Re-raise to be handled by exception handlers
        
        except Exception as exc:
            # Log unexpected exceptions
            logging.error(
                f"Unexpected error: {exc}",
                exc_info=True,
                extra={
                    "method": request.method,
                    "path": request.url.path,
                    "client": request.client.host if request.client else None,
                }
            )
            raise

# Add to middleware stack
app.add_middleware(ErrorLoggingMiddleware)

from starlette.testclient import TestClient
import pytest

def test_http_exceptions():
    client = TestClient(app)
    
    # Test 404 Not Found
    response = client.get("/nonexistent")
    assert response.status_code == 404
    assert "error" in response.json()
    
    # Test 400 Bad Request
    response = client.post("/users", json={"invalid": "data"})
    assert response.status_code == 400
    
    # Test 401 Unauthorized
    response = client.get("/protected")
    assert response.status_code == 401
    
    # Test 403 Forbidden
    response = client.get("/admin", headers={"Authorization": "Bearer user-token"})
    assert response.status_code == 403

def test_custom_exceptions():
    client = TestClient(app)
    
    # Test validation error
    response = client.post("/users", json={})
    assert response.status_code == 422
    
    error_data = response.json()["error"]
    assert error_data["type"] == "ValidationError"
    assert "errors" in error_data["details"]

def test_exception_handlers():
    client = TestClient(app)
    
    # Test that exceptions are properly handled
    response = client.get("/trigger-error")
    assert response.status_code == 500
    
    # Verify error format
    error = response.json()["error"]
    assert error["type"] == "InternalServerError"

def test_websocket_exceptions():
    client = TestClient(app)
    
    # Test authentication failure
    with pytest.raises(WebSocketDenialResponse):
        with client.websocket_connect("/ws"):
            pass
    
    # Test successful connection
    with client.websocket_connect("/ws?token=valid") as websocket:
        websocket.send_json({"type": "ping"})
        response = websocket.receive_json()
        assert response["type"] == "pong"
        
        # Test invalid message handling
        websocket.send_text("invalid json")
        error_response = websocket.receive_json()
        assert error_response["type"] == "error"