"""
Structured logging configuration for ReadAnyBook.

Supports:
- JSON formatted logs for production
- Pretty console logs for development
- Request/document ID tracking
- Integration with observability platforms
"""

import json
import logging
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional
import contextvars

# Context variables for request tracking
request_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("request_id", default="")
document_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("document_id", default="")
chain_name_var: contextvars.ContextVar[str] = contextvars.ContextVar("chain_name", default="")


class JSONFormatter(logging.Formatter):
    """JSON log formatter for structured logging."""
    
    def format(self, record: logging.LogRecord) -> str:
        log_data = {
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
            "module": record.module,
            "function": record.funcName,
            "line": record.lineno,
        }
        
        # Add context variables
        if request_id := request_id_var.get():
            log_data["request_id"] = request_id
        if document_id := document_id_var.get():
            log_data["document_id"] = document_id
        if chain_name := chain_name_var.get():
            log_data["chain_name"] = chain_name
        
        # Add extra fields
        if hasattr(record, "extra_fields"):
            log_data.update(record.extra_fields)
        
        # Add exception info
        if record.exc_info:
            log_data["exception"] = self.formatException(record.exc_info)
        
        return json.dumps(log_data)


class PrettyFormatter(logging.Formatter):
    """Pretty console formatter for development."""
    
    COLORS = {
        "DEBUG": "\033[36m",     # Cyan
        "INFO": "\033[32m",      # Green
        "WARNING": "\033[33m",   # Yellow
        "ERROR": "\033[31m",     # Red
        "CRITICAL": "\033[35m",  # Magenta
    }
    RESET = "\033[0m"
    
    def format(self, record: logging.LogRecord) -> str:
        color = self.COLORS.get(record.levelname, "")
        
        # Build context string
        context_parts = []
        if request_id := request_id_var.get():
            context_parts.append(f"req={request_id[:8]}")
        if document_id := document_id_var.get():
            context_parts.append(f"doc={document_id[:8]}")
        if chain_name := chain_name_var.get():
            context_parts.append(f"chain={chain_name}")
        
        context_str = f" [{', '.join(context_parts)}]" if context_parts else ""
        
        timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
        
        return (
            f"{color}{timestamp} | {record.levelname:8s}{self.RESET} | "
            f"{record.name}:{record.lineno}{context_str} - {record.getMessage()}"
        )


class ContextLogger(logging.LoggerAdapter):
    """Logger adapter that adds context to all log messages."""
    
    def process(self, msg: str, kwargs: Dict[str, Any]) -> tuple:
        extra = kwargs.get("extra", {})
        
        # Add any extra fields to the record
        if "extra_fields" not in extra:
            extra["extra_fields"] = {}
        
        if additional := kwargs.pop("additional", None):
            extra["extra_fields"].update(additional)
        
        kwargs["extra"] = extra
        return msg, kwargs
    
    def with_context(
        self,
        request_id: Optional[str] = None,
        document_id: Optional[str] = None,
        chain_name: Optional[str] = None
    ) -> "ContextLogger":
        """Set context variables for subsequent log calls."""
        if request_id:
            request_id_var.set(request_id)
        if document_id:
            document_id_var.set(document_id)
        if chain_name:
            chain_name_var.set(chain_name)
        return self


_loggers: Dict[str, ContextLogger] = {}
_configured = False


def setup_logging(
    level: str = "INFO",
    log_format: str = "json",
    log_file: Optional[str] = None,
    structured: bool = True
) -> None:
    """
    Set up logging configuration.
    
    Args:
        level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
        log_format: Format type ("json" or "pretty")
        log_file: Optional file path for log output
        structured: Whether to use structured logging
    """
    global _configured
    
    root_logger = logging.getLogger("readanybook")
    root_logger.setLevel(getattr(logging, level.upper()))
    
    # Remove existing handlers
    root_logger.handlers.clear()
    
    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    if log_format == "json" and structured:
        console_handler.setFormatter(JSONFormatter())
    else:
        console_handler.setFormatter(PrettyFormatter())
    root_logger.addHandler(console_handler)
    
    # File handler
    if log_file:
        log_path = Path(log_file)
        log_path.parent.mkdir(parents=True, exist_ok=True)
        
        file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(JSONFormatter())  # Always JSON for files
        root_logger.addHandler(file_handler)
    
    _configured = True


def get_logger(name: str) -> ContextLogger:
    """
    Get a logger with context support.
    
    Args:
        name: Logger name (typically __name__)
        
    Returns:
        ContextLogger instance
    """
    global _configured
    
    if not _configured:
        setup_logging(level="INFO", log_format="pretty", structured=False)
    
    if name not in _loggers:
        logger = logging.getLogger(f"readanybook.{name}" if not name.startswith("readanybook") else name)
        _loggers[name] = ContextLogger(logger, {})
    
    return _loggers[name]


class LogContext:
    """Context manager for scoped logging context."""
    
    def __init__(
        self,
        request_id: Optional[str] = None,
        document_id: Optional[str] = None,
        chain_name: Optional[str] = None
    ):
        self.request_id = request_id
        self.document_id = document_id
        self.chain_name = chain_name
        self._tokens = []
    
    def __enter__(self) -> "LogContext":
        if self.request_id:
            self._tokens.append(("request_id", request_id_var.set(self.request_id)))
        if self.document_id:
            self._tokens.append(("document_id", document_id_var.set(self.document_id)))
        if self.chain_name:
            self._tokens.append(("chain_name", chain_name_var.set(self.chain_name)))
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        for var_name, token in self._tokens:
            if var_name == "request_id":
                request_id_var.reset(token)
            elif var_name == "document_id":
                document_id_var.reset(token)
            elif var_name == "chain_name":
                chain_name_var.reset(token)
