"""Structured logging configuration with correlation IDs for OCLAWMA.

This module provides JSON structured logging with correlation IDs that flow
through the entire job lifecycle for easier debugging.

Example:
    >>> from oclawma.logging_config import get_logger, correlation_id
    >>> with correlation_id("job-123"):
    ...     logger = get_logger(__name__)
    ...     logger.info("Processing started")
"""

from __future__ import annotations

import contextvars
import json
import logging
import logging.handlers
import os
import sys
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# Context variable for correlation ID
_correlation_id_var: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "correlation_id", default=None
)


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

    Outputs log records as JSON with standardized fields for easy parsing
    and ingestion into log aggregation systems.

    Attributes:
        default_fields: Default fields to include in every log entry.
    """

    def __init__(
        self,
        *,
        fmt: str | None = None,
        datefmt: str | None = None,
        style: str = "%",
        validate: bool = True,
        indent: int | None = None,
    ) -> None:
        """Initialize the JSON formatter.

        Args:
            fmt: Unused (for compatibility).
            datefmt: Date format string (ISO 8601 by default).
            style: Format style (for compatibility).
            validate: Whether to validate the format.
            indent: JSON indentation (None for compact, 2 for pretty-print).
        """
        super().__init__(fmt=fmt, datefmt=datefmt, style=style, validate=validate)
        self.indent = indent
        self._hostname = os.uname().nodename if hasattr(os, "uname") else "unknown"
        self._pid = os.getpid()

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

        Args:
            record: The log record to format.

        Returns:
            JSON string representation of the log record.
        """
        log_entry: dict[str, Any] = {
            "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
            "module": record.module,
            "function": record.funcName,
            "line": record.lineno,
            "pid": self._pid,
            "hostname": self._hostname,
        }

        # Add correlation ID if available
        correlation_id = getattr(record, "correlation_id", None) or _correlation_id_var.get()
        if correlation_id:
            log_entry["correlation_id"] = correlation_id

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

        # Add stack info if present
        if record.stack_info:
            log_entry["stack_trace"] = record.stack_info

        # Add extra fields from record
        for key, value in record.__dict__.items():
            if key not in (
                "name",
                "msg",
                "args",
                "levelname",
                "levelno",
                "pathname",
                "filename",
                "module",
                "exc_info",
                "exc_text",
                "stack_info",
                "lineno",
                "funcName",
                "created",
                "msecs",
                "relativeCreated",
                "thread",
                "threadName",
                "processName",
                "process",
                "message",
                "asctime",
                "correlation_id",
            ):
                try:
                    # Try to serialize value
                    json.dumps({key: value})
                    log_entry[key] = value
                except (TypeError, ValueError):
                    log_entry[key] = str(value)

        return json.dumps(log_entry, indent=self.indent, default=str)


class CorrelationIdFilter(logging.Filter):
    """Filter that adds correlation ID to log records.

    This filter injects the current correlation ID into each log record,
    making it available to formatters and handlers.
    """

    def filter(self, record: logging.LogRecord) -> bool:
        """Add correlation ID to the log record.

        Args:
            record: The log record to process.

        Returns:
            True to include the record in log output.
        """
        correlation_id = _correlation_id_var.get()
        if correlation_id:
            record.correlation_id = correlation_id
        return True


class TraceContext:
    """Trace context for job execution flow.

    Provides a way to track the execution flow through async boundaries
    by maintaining a trace ID and span information.
    """

    def __init__(
        self,
        trace_id: str | None = None,
        span_id: str | None = None,
        parent_span_id: str | None = None,
    ) -> None:
        """Initialize trace context.

        Args:
            trace_id: Unique trace identifier (generated if not provided).
            span_id: Current span identifier (generated if not provided).
            parent_span_id: Parent span identifier for nested spans.
        """
        self.trace_id = trace_id or self._generate_id()
        self.span_id = span_id or self._generate_id()
        self.parent_span_id = parent_span_id

    @staticmethod
    def _generate_id() -> str:
        """Generate a unique identifier.

        Returns:
            Unique hex string.
        """
        return uuid.uuid4().hex[:16]

    def to_dict(self) -> dict[str, str]:
        """Convert trace context to dictionary.

        Returns:
            Dictionary with trace context fields.
        """
        result: dict[str, str] = {
            "trace_id": self.trace_id,
            "span_id": self.span_id,
        }
        if self.parent_span_id:
            result["parent_span_id"] = self.parent_span_id
        return result

    def child_span(self) -> TraceContext:
        """Create a child span context.

        Returns:
            New TraceContext with current span as parent.
        """
        return TraceContext(
            trace_id=self.trace_id,
            parent_span_id=self.span_id,
        )


# Context variable for trace context
_trace_context_var: contextvars.ContextVar[TraceContext | None] = contextvars.ContextVar(
    "trace_context", default=None
)


class TraceContextFilter(logging.Filter):
    """Filter that adds trace context to log records."""

    def filter(self, record: logging.LogRecord) -> bool:
        """Add trace context to the log record.

        Args:
            record: The log record to process.

        Returns:
            True to include the record in log output.
        """
        trace_ctx = _trace_context_var.get()
        if trace_ctx:
            record.trace_id = trace_ctx.trace_id
            record.span_id = trace_ctx.span_id
            if trace_ctx.parent_span_id:
                record.parent_span_id = trace_ctx.parent_span_id
        return True


@contextmanager
def correlation_id(cid: str | None = None) -> Iterator[str]:
    """Context manager for setting correlation ID.

    Sets a correlation ID for the duration of the context, making it
    available to all loggers within that context.

    Args:
        cid: Correlation ID to use (generated if not provided).

    Yields:
        The correlation ID being used.

    Example:
        >>> with correlation_id("job-123") as cid:
        ...     logger.info("Starting job")
        ...     # All logs within this block include correlation_id="job-123"
    """
    generated_id = cid or uuid.uuid4().hex
    token = _correlation_id_var.set(generated_id)
    try:
        yield generated_id
    finally:
        _correlation_id_var.reset(token)


@contextmanager
def trace_context(
    trace_id: str | None = None,
    span_id: str | None = None,
    parent_span_id: str | None = None,
) -> Iterator[TraceContext]:
    """Context manager for setting trace context.

    Sets a trace context for distributed tracing across async boundaries.

    Args:
        trace_id: Trace ID (generated if not provided).
        span_id: Span ID (generated if not provided).
        parent_span_id: Parent span ID for nested traces.

    Yields:
        The TraceContext being used.

    Example:
        >>> with trace_context() as ctx:
        ...     logger.info("Operation started")
        ...     with trace_context(parent_span_id=ctx.span_id) as child_ctx:
        ...         logger.info("Sub-operation")
    """
    ctx = TraceContext(trace_id, span_id, parent_span_id)
    token = _trace_context_var.set(ctx)
    try:
        yield ctx
    finally:
        _trace_context_var.reset(token)


def get_correlation_id() -> str | None:
    """Get the current correlation ID.

    Returns:
        Current correlation ID or None if not set.
    """
    return _correlation_id_var.get()


def get_trace_context() -> TraceContext | None:
    """Get the current trace context.

    Returns:
        Current TraceContext or None if not set.
    """
    return _trace_context_var.get()


def setup_logging(
    level: str | int = logging.INFO,
    json_format: bool = True,
    log_file: str | Path | None = None,
    max_bytes: int = 10 * 1024 * 1024,  # 10MB
    backup_count: int = 5,
    retention_days: int = 30,
    enable_console: bool = True,
) -> None:
    """Configure structured logging for OCLAWMA.

    Sets up the root logger with JSON formatting, correlation ID injection,
    and rotation policies.

    Args:
        level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL or int).
        json_format: Whether to use JSON formatting (vs plain text).
        log_file: Path to log file (if None, file logging is disabled).
        max_bytes: Maximum size of log file before rotation.
        backup_count: Number of backup files to keep.
        retention_days: Days to retain old log files (for cleanup).
        enable_console: Whether to log to console (stdout).
    """
    # Convert string level to int
    if isinstance(level, str):
        level = getattr(logging, level.upper(), logging.INFO)

    # Configure root logger
    root_logger = logging.getLogger()
    root_logger.setLevel(level)

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

    # Create formatter
    if json_format:
        formatter: logging.Formatter = JSONFormatter()
    else:
        formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")

    # Console handler
    if enable_console:
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(level)
        console_handler.setFormatter(formatter)
        console_handler.addFilter(CorrelationIdFilter())
        console_handler.addFilter(TraceContextFilter())
        root_logger.addHandler(console_handler)

    # File handler with rotation
    if log_file:
        log_path = Path(log_file)
        log_path.parent.mkdir(parents=True, exist_ok=True)

        file_handler = logging.handlers.RotatingFileHandler(
            log_file,
            maxBytes=max_bytes,
            backupCount=backup_count,
            encoding="utf-8",
        )
        file_handler.setLevel(level)
        file_handler.setFormatter(formatter)
        file_handler.addFilter(CorrelationIdFilter())
        file_handler.addFilter(TraceContextFilter())
        root_logger.addHandler(file_handler)

        # Schedule old log cleanup (simple implementation)
        _cleanup_old_logs(log_path.parent, retention_days)


def _cleanup_old_logs(log_dir: Path, retention_days: int) -> None:
    """Clean up log files older than retention period.

    Args:
        log_dir: Directory containing log files.
        retention_days: Number of days to retain logs.
    """
    if retention_days <= 0:
        return

    from datetime import datetime, timedelta

    cutoff = datetime.now() - timedelta(days=retention_days)

    try:
        for log_file in log_dir.glob("*.log*"):
            try:
                if log_file.stat().st_mtime < cutoff.timestamp():
                    log_file.unlink()
            except OSError:
                pass  # Ignore permission errors
    except OSError:
        pass  # Ignore directory access errors


def get_logger(name: str) -> logging.Logger:
    """Get a configured logger instance.

    Returns a logger that will automatically include correlation IDs
    and trace context in its output.

    Args:
        name: Logger name (typically __name__).

    Returns:
        Configured logger instance.
    """
    logger = logging.getLogger(name)

    # Ensure logger has correlation ID filter
    has_correlation_filter = any(isinstance(f, CorrelationIdFilter) for f in logger.filters)
    if not has_correlation_filter:
        logger.addFilter(CorrelationIdFilter())

    # Ensure logger has trace context filter
    has_trace_filter = any(isinstance(f, TraceContextFilter) for f in logger.filters)
    if not has_trace_filter:
        logger.addFilter(TraceContextFilter())

    return logger


# Convenience function for job lifecycle logging
def log_job_start(
    logger: logging.Logger,
    job_id: str,
    job_type: str,
    **kwargs: Any,
) -> None:
    """Log job start event with correlation ID.

    Args:
        logger: Logger instance.
        job_id: Unique job identifier.
        job_type: Type of job being executed.
        **kwargs: Additional context to log.
    """
    with correlation_id(job_id):
        logger.info(
            f"Job started: {job_type}",
            extra={"job_id": job_id, "job_type": job_type, **kwargs},
        )


def log_job_end(
    logger: logging.Logger,
    job_id: str,
    job_type: str,
    success: bool = True,
    duration_ms: float | None = None,
    **kwargs: Any,
) -> None:
    """Log job completion event.

    Args:
        logger: Logger instance.
        job_id: Unique job identifier.
        job_type: Type of job being executed.
        success: Whether the job completed successfully.
        duration_ms: Duration in milliseconds.
        **kwargs: Additional context to log.
    """
    with correlation_id(job_id):
        status = "completed" if success else "failed"
        extra = {"job_id": job_id, "job_type": job_type, "success": success, **kwargs}
        if duration_ms is not None:
            extra["duration_ms"] = duration_ms

        if success:
            logger.info(f"Job {status}: {job_type}", extra=extra)
        else:
            logger.error(f"Job {status}: {job_type}", extra=extra)
