"""Structured JSON logging utilities for observability.

Provides JSON-formatted logging with consistent schema for machine-readable logs
that can be easily ingested into log analysis systems.
"""

import json
import logging
from datetime import UTC, datetime
from typing import Any


class StructuredFormatter(logging.Formatter):
    """JSON formatter for structured logging.

    Outputs logs in JSON format with consistent fields:
    - timestamp: ISO 8601 timestamp in UTC
    - log_level: Log level (INFO, ERROR, etc.)
    - logger_name: Logger name
    - message: Log message
    - module: Python module name
    - function: Function name
    - line: Line number
    - trace_id: Optimization trace ID (if available)
    - event_type: Event type for categorization (if available)
    - metadata: Additional structured data (if available)
    """

    def format(self, record: logging.LogRecord) -> str:
        """Format log record as JSON."""
        log_data: dict[str, Any] = {
            "timestamp": datetime.now(UTC).isoformat(),
            "log_level": record.levelname,
            "logger_name": record.name,
            "message": record.getMessage(),
            "module": record.module,
            "function": record.funcName,
            "line": record.lineno,
        }

        # Add trace_id if present (critical for correlating logs)
        if hasattr(record, "trace_id"):
            log_data["trace_id"] = record.trace_id

        # Add event_type if present (for categorization)
        if hasattr(record, "event_type"):
            log_data["event_type"] = record.event_type

        # Add metadata if present (for additional context)
        if hasattr(record, "metadata"):
            log_data["metadata"] = record.metadata

        # Add user_id if present
        if hasattr(record, "user_id"):
            log_data["user_id"] = record.user_id

        # Add exception info if present
        if record.exc_info:
            log_data["exception"] = self.formatException(record.exc_info)

        return json.dumps(log_data)


def get_observability_logger(name: str, component: str = "aiservice") -> logging.Logger:
    """Get a logger configured for observability.

    Args:
        name: Logger name (typically __name__ of the calling module)
        component: Component name for organization (default: "aiservice")

    Returns:
        Configured logger instance

    Usage:
        logger = get_observability_logger(__name__, "optimizer")
        logger.info(
            "Optimization started",
            extra={
                "trace_id": "abc-123",
                "event_type": "optimization.started",
                "metadata": {"model": "gpt-4", "n_candidates": 5}
            }
        )

    """
    logger_name = f"{component}.{name}"
    logger = logging.getLogger(logger_name)

    # Logger configuration will be set globally in settings.py
    # This just returns the logger with the correct name hierarchy
    return logger


def log_with_context(
    logger: logging.Logger,
    level: str,
    message: str,
    trace_id: str | None = None,
    event_type: str | None = None,
    user_id: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> None:
    """Log a message with structured context.

    Helper function for logging with consistent structured context fields.

    Args:
        logger: Logger instance
        level: Log level ('DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL')
        message: Log message
        trace_id: Optimization trace ID for correlation
        event_type: Event type for categorization
        user_id: User ID for user tracking
        metadata: Additional structured data

    Usage:
        log_with_context(
            logger,
            'INFO',
            'LLM call completed',
            trace_id='abc-123',
            event_type='llm.call_success',
            metadata={'model': 'gpt-4', 'latency_ms': 1250}
        )

    """
    extra: dict[str, Any] = {}

    if trace_id:
        extra["trace_id"] = trace_id
    if event_type:
        extra["event_type"] = event_type
    if user_id:
        extra["user_id"] = user_id
    if metadata:
        extra["metadata"] = metadata

    log_level = getattr(logging, level.upper())
    logger.log(log_level, message, extra=extra)


class LogContext:
    """Context manager for logging with persistent context fields.

    Allows setting context fields (like trace_id) that will be included
    in all log messages within the context.

    Usage:
        with LogContext(logger, trace_id="abc-123", user_id="user-456") as log:
            log.info("Step 1 started")
            # ... do work ...
            log.info("Step 1 completed")
            # All logs will include trace_id and user_id
    """

    def __init__(self, logger: logging.Logger, trace_id: str | None = None, user_id: str | None = None, **kwargs: Any):
        """Initialize log context.

        Args:
            logger: Logger instance
            trace_id: Optimization trace ID
            user_id: User ID
            **kwargs: Additional context fields to include in all logs

        """
        self.logger = logger
        self.context: dict[str, Any] = {}

        if trace_id:
            self.context["trace_id"] = trace_id
        if user_id:
            self.context["user_id"] = user_id
        self.context.update(kwargs)

    def __enter__(self) -> "LogContext":
        """Enter context."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context."""
        # Context is automatically released

    def log(
        self, level: str, message: str, event_type: str | None = None, metadata: dict[str, Any] | None = None
    ) -> None:
        """Log message with context."""
        extra = dict(self.context)
        if event_type:
            extra["event_type"] = event_type
        if metadata:
            extra["metadata"] = metadata

        log_level = getattr(logging, level.upper())
        self.logger.log(log_level, message, extra=extra)

    def debug(self, message: str, **kwargs):
        """Log debug message."""
        self.log("DEBUG", message, **kwargs)

    def info(self, message: str, **kwargs):
        """Log info message."""
        self.log("INFO", message, **kwargs)

    def warning(self, message: str, **kwargs):
        """Log warning message."""
        self.log("WARNING", message, **kwargs)

    def error(self, message: str, **kwargs):
        """Log error message."""
        self.log("ERROR", message, **kwargs)

    def critical(self, message: str, **kwargs):
        """Log critical message."""
        self.log("CRITICAL", message, **kwargs)


# Event type constants for consistent categorization
class EventType:
    """Standard event types for optimization flow."""

    # Validation events
    VALIDATION_STARTED = "validation.started"
    VALIDATION_SUCCESS = "validation.success"
    VALIDATION_FAILED = "validation.failed"

    # LLM events
    LLM_CALL_STARTED = "llm.call_started"
    LLM_CALL_SUCCESS = "llm.call_success"
    LLM_CALL_FAILED = "llm.call_failed"
    LLM_PARSE_ERROR = "llm.parse_error"

    # Testing events
    TEST_GENERATION_STARTED = "test.generation_started"
    TEST_GENERATION_SUCCESS = "test.generation_success"
    TEST_GENERATION_FAILED = "test.generation_failed"
    TEST_EXECUTION_STARTED = "test.execution_started"
    TEST_EXECUTION_SUCCESS = "test.execution_success"
    TEST_EXECUTION_FAILED = "test.execution_failed"

    # Correctness events
    CORRECTNESS_VERIFIED = "correctness.verified"
    CORRECTNESS_FAILED = "correctness.failed"

    # Performance events
    BENCHMARK_STARTED = "benchmark.started"
    BENCHMARK_SUCCESS = "benchmark.success"
    BENCHMARK_FAILED = "benchmark.failed"

    # Ranking events
    RANKING_STARTED = "ranking.started"
    RANKING_SUCCESS = "ranking.success"
    RANKING_FAILED = "ranking.failed"

    # Optimization flow events
    OPTIMIZATION_STARTED = "optimization.started"
    OPTIMIZATION_SUCCESS = "optimization.success"
    OPTIMIZATION_FAILED = "optimization.failed"
