response.py

from datetime import datetime
import json

class Response:
    """
    HTTP Response class that handles creating and formatting HTTP responses
    according to RFC 7231 specifications
    """
    def __init__(self, version='1.1', status=200, message='OK'):
        self.status = status
        self.message = message
        self.version = version
        self.headers = {}
        self.body = ""
        self.binary_data = None
        
        # Set default headers
        self._set_default_headers()

    def _set_default_headers(self):
        """Set default HTTP headers for all responses"""
        self.headers['Date'] = datetime.utcnow().strftime('%a, %d %b %Y %H:%M:%S GMT')
        self.headers['Server'] = 'Multi-threaded HTTP Server'
        self.headers['Connection'] = 'keep-alive'
        self.headers['Keep-Alive'] = 'timeout=30, max=100'

    def __str__(self) -> str:
        """
        Convert response to HTTP string format
        Returns the complete HTTP response as a string
        """
        # Create status line
        response_line = f'HTTP/{self.version} {self.status} {self.message}'
        
        # Format headers
        header_lines = []
        for key, value in self.headers.items():
            header_lines.append(f'{key}: {value}')
        
        # Join response line and headers
        response = response_line + '\r\n' + '\r\n'.join(header_lines) + '\r\n\r\n'
        
        # Add body if it exists
        if self.body:
            response += self.body
            
        return response

    def to_bytes(self) -> bytes:
        """
        Convert response to bytes format for binary data transmission
        """
        # Create status line and headers as string
        response_line = f'HTTP/{self.version} {self.status} {self.message}'
        header_lines = []
        for key, value in self.headers.items():
            header_lines.append(f'{key}: {value}')
        
        headers_str = response_line + '\r\n' + '\r\n'.join(header_lines) + '\r\n\r\n'
        headers_bytes = headers_str.encode('utf-8')
        
        # Add binary data if it exists
        if self.binary_data is not None:
            return headers_bytes + self.binary_data
        elif self.body:
            return headers_bytes + self.body.encode('utf-8')
        else:
            return headers_bytes

    def add_header(self, key: str, value: str) -> None:
        """Add a header to the response"""
        self.headers[key] = value

    def set_body(self, data: str):
        """Set text body data and Content-Length header"""
        self.body = data
        self.headers['Content-Length'] = str(len(data.encode('utf-8')))

    def set_binary_data(self, data: bytes):
        """Set binary body data and Content-Length header"""
        self.binary_data = data
        self.headers['Content-Length'] = str(len(data))

    def set_json_body(self, data: dict):
        """Set JSON body data with appropriate headers"""
        json_str = json.dumps(data)
        self.set_body(json_str)
        self.add_header('Content-Type', 'application/json')

    def set_html_body(self, html: str):
        """Set HTML body data with appropriate headers"""
        self.set_body(html)
        self.add_header('Content-Type', 'text/html; charset=utf-8')

    def set_file_download(self, filename: str, data: bytes):
        """Set response for file download with appropriate headers"""
        self.set_binary_data(data)
        self.add_header('Content-Type', 'application/octet-stream')
        self.add_header('Content-Disposition', f'attachment; filename="{filename}"')