"""
Structured JSON logging for atomingest.

This module provides a JSON formatter for logging that outputs structured logs
suitable for ingestion by log aggregation systems like ELK, Splunk, or Loki.
"""

import logging
import json
import traceback
from datetime import datetime
from typing import Any, Dict, Optional


class JSONFormatter(logging.Formatter):
    """
    Custom logging formatter that outputs JSON-structured logs.

    Each log record includes:
    - timestamp: ISO 8601 timestamp
    - level: Log level (INFO, WARNING, ERROR, etc.)
    - logger: Logger name
    - message: Log message
    - context: Additional context fields (if provided)
    - exception: Exception info (if present)
    - source: Source location (module, function, line number)

    Usage:
        handler = logging.StreamHandler()
        handler.setFormatter(JSONFormatter())
        logger = logging.getLogger()
        logger.addHandler(handler)
    """

    def __init__(
        self,
        fmt_keys: Optional[Dict[str, str]] = None,
        include_traceback: bool = True,
        include_source: bool = True,
    ):
        """
        Initialize JSON formatter.

        Args:
            fmt_keys: Optional mapping of output keys to record attributes.
            include_traceback: Whether to include full traceback for exceptions.
            include_source: Whether to include source location info.
        """
        super().__init__()
        self.fmt_keys = fmt_keys or {
            "timestamp": "asctime",
            "level": "levelname",
            "logger": "name",
            "message": "message",
        }
        self.include_traceback = include_traceback
        self.include_source = include_source

    def format(self, record: logging.LogRecord) -> str:
        """
        Format a log record as JSON.

        Args:
            record: The log record to format.

        Returns:
            JSON-formatted log string.
        """
        # Build base log dict
        log_dict: Dict[str, Any] = {}

        # Add standard fields
        for output_key, record_attr in self.fmt_keys.items():
            if hasattr(record, record_attr):
                value = getattr(record, record_attr)
                log_dict[output_key] = value

        # Add timestamp if not already present
        if "timestamp" not in log_dict:
            log_dict["timestamp"] = datetime.fromtimestamp(record.created).isoformat()

        # Format message
        log_dict["message"] = record.getMessage()

        # Add source location
        if self.include_source:
            log_dict["source"] = {
                "module": record.module,
                "function": record.funcName,
                "line": record.lineno,
                "pathname": record.pathname,
            }

        # Add thread/process info
        log_dict["thread"] = {
            "id": record.thread,
            "name": record.threadName,
        }
        log_dict["process"] = {
            "id": record.process,
            "name": record.processName,
        }

        # Add exception info if present
        if record.exc_info:
            log_dict["exception"] = {
                "type": record.exc_info[0].__name__,
                "message": str(record.exc_info[1]),
            }

            if self.include_traceback:
                log_dict["exception"]["traceback"] = "".join(
                    traceback.format_exception(*record.exc_info)
                )

        # Add custom context fields from record.__dict__
        # These are extra fields added via logger.info("msg", extra={...})
        context = {}
        for key, value in record.__dict__.items():
            # Skip standard logging attributes
            if key not in [
                "name",
                "msg",
                "args",
                "created",
                "filename",
                "funcName",
                "levelname",
                "levelno",
                "lineno",
                "module",
                "msecs",
                "message",
                "pathname",
                "process",
                "processName",
                "relativeCreated",
                "thread",
                "threadName",
                "exc_info",
                "exc_text",
                "stack_info",
                "asctime",
            ]:
                # Try to serialize the value
                try:
                    json.dumps(value)  # Test if serializable
                    context[key] = value
                except (TypeError, ValueError):
                    context[key] = str(value)

        if context:
            log_dict["context"] = context

        # Serialize to JSON
        try:
            return json.dumps(log_dict, default=str)
        except Exception as e:
            # Fallback to simple dict if serialization fails
            return json.dumps(
                {
                    "timestamp": datetime.now().isoformat(),
                    "level": "ERROR",
                    "logger": "JSONFormatter",
                    "message": f"Failed to serialize log record: {e}",
                    "original_message": str(record.getMessage()),
                }
            )


class StructuredLogger:
    """
    Wrapper around Python logger that makes it easier to log with context.

    Usage:
        logger = StructuredLogger("my_module")
        logger.info("Processing record", record_id=123, table="users")
        logger.error("Failed to save", error=str(e), record_id=123)
    """

    def __init__(self, name: str):
        """
        Initialize structured logger.

        Args:
            name: Logger name (typically __name__).
        """
        self.logger = logging.getLogger(name)

    def _log(self, level: int, message: str, **context):
        """
        Internal log method with context.

        Args:
            level: Log level (logging.INFO, logging.ERROR, etc.).
            message: Log message.
            **context: Additional context fields.
        """
        self.logger.log(level, message, extra=context)

    def debug(self, message: str, **context):
        """Log debug message with context."""
        self._log(logging.DEBUG, message, **context)

    def info(self, message: str, **context):
        """Log info message with context."""
        self._log(logging.INFO, message, **context)

    def warning(self, message: str, **context):
        """Log warning message with context."""
        self._log(logging.WARNING, message, **context)

    def error(self, message: str, exc_info=None, **context):
        """
        Log error message with context.

        Args:
            message: Error message.
            exc_info: Exception info (from sys.exc_info()).
            **context: Additional context fields.
        """
        if exc_info:
            self.logger.error(message, exc_info=exc_info, extra=context)
        else:
            self._log(logging.ERROR, message, **context)

    def critical(self, message: str, exc_info=None, **context):
        """
        Log critical message with context.

        Args:
            message: Critical message.
            exc_info: Exception info (from sys.exc_info()).
            **context: Additional context fields.
        """
        if exc_info:
            self.logger.critical(message, exc_info=exc_info, extra=context)
        else:
            self._log(logging.CRITICAL, message, **context)


def configure_structured_logging(
    level: int = logging.INFO,
    include_traceback: bool = True,
    include_source: bool = True,
) -> None:
    """
    Configure the root logger to use structured JSON logging.

    Args:
        level: Logging level (default: INFO).
        include_traceback: Whether to include full traceback for exceptions.
        include_source: Whether to include source location info.
    """
    # Get root logger
    root_logger = logging.getLogger()
    root_logger.setLevel(level)

    # Remove existing handlers
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)

    # Add JSON handler to stdout
    handler = logging.StreamHandler()
    handler.setFormatter(
        JSONFormatter(
            include_traceback=include_traceback,
            include_source=include_source,
        )
    )
    root_logger.addHandler(handler)


def get_logger(name: str) -> StructuredLogger:
    """
    Get a structured logger.

    Args:
        name: Logger name (typically __name__).

    Returns:
        StructuredLogger instance.
    """
    return StructuredLogger(name)
