from pythonjsonlogger.orjson import OrjsonFormatter
from datetime import datetime, timezone
from typing import Dict, Any, TYPE_CHECKING
import warnings

if TYPE_CHECKING:
    from logging import LogRecord


class CustomJsonFormatter(OrjsonFormatter):
    """Custom JSON formatter that adds timestamps and normalizes log levels.
    
    This formatter extends the OrjsonFormatter to add:
    - Automatic timestamp generation in ISO 8601 format with microseconds
    - Log level normalization to uppercase
    - Correlation ID string conversion
    """

    def add_fields(self, log_record: Dict[str, Any], record: 'LogRecord', message_dict: Dict[str, Any]) -> None:
        """Add custom fields to the log record.
        
        Args:
            log_record: The dictionary that will be serialized to JSON
            record: The original LogRecord instance
            message_dict: Additional message fields
        """
        super().add_fields(log_record, record, message_dict)
        
        # Add timestamp if not present
        if not log_record.get('timestamp'):
            now = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%fZ')
            log_record['timestamp'] = now
        
        # Normalize log level
        if log_record.get('level'):
            log_record['level'] = log_record['level'].upper()
        else:
            log_record['level'] = record.levelname
        
        # Convert correlation_id to string if present and not None
        if log_record.get('correlation_id') is not None:
            log_record['correlation_id'] = str(log_record['correlation_id'])


class AccessJsonFormatter(CustomJsonFormatter):
    """Specialized JSON formatter for HTTP access logs.
    
    This formatter expects log record args in the following order:
    (client_ip, http_method, path, protocol, status_code)
    
    Example:
        access_logger.info("Request", "192.168.1.1", "GET", "/api/users", "HTTP/1.1", 200)
    """

    def add_fields(self, log_record: Dict[str, Any], record: 'LogRecord', message_dict: Dict[str, Any]) -> None:
        """Add access log fields and standard fields.
        
        Args:
            log_record: The dictionary that will be serialized to JSON
            record: The original LogRecord instance  
            message_dict: Additional message fields
            
        Raises:
            ValueError: If record.args doesn't have exactly 5 elements
            IndexError: If record.args is malformed
        """
        # First add standard fields from parent
        super().add_fields(log_record, record, message_dict)
        
        # Parse access log arguments
        try:
            if len(record.args) != 5:
                raise ValueError(
                    f"AccessJsonFormatter expects exactly 5 args (ip, method, path, protocol, status_code), "
                    f"got {len(record.args)}: {record.args}"
                )
            
            ip, method, path, protocol, status_code = record.args
            
            # Add access-specific fields
            log_record['client_addr'] = str(ip)
            log_record['method'] = str(method)
            log_record['path'] = str(path)
            log_record['protocol'] = str(protocol)
            log_record['status_code'] = int(status_code) if str(status_code).isdigit() else status_code
            
        except (ValueError, TypeError, IndexError) as e:
            # Log the error but don't break the logging system
            warnings.warn(f"Failed to parse access log arguments: {e}. Args: {record.args}")
            
            # Add error information to the log record
            log_record['access_log_parse_error'] = str(e)
            log_record['raw_args'] = str(record.args)
