"""Logging utilities."""
import logging
import sys
from pathlib import Path
from typing import Optional
from datetime import datetime


def setup_logger(
    name: str,
    log_file: Optional[str] = None,
    level: int = logging.INFO,
    format_string: Optional[str] = None
) -> logging.Logger:
    """
    Set up a logger with console and optional file output.
    
    Args:
        name: Logger name
        log_file: Optional path to log file
        level: Logging level
        format_string: Optional custom format string
    
    Returns:
        Configured logger
    """
    if format_string is None:
        format_string = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    
    formatter = logging.Formatter(format_string)
    
    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.handlers.clear()
    
    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(level)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    # File handler
    if log_file:
        Path(log_file).parent.mkdir(parents=True, exist_ok=True)
        file_handler = logging.FileHandler(log_file)
        file_handler.setLevel(level)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    
    return logger


def get_logger(name: str) -> logging.Logger:
    """Get or create a logger."""
    return logging.getLogger(name)


class TrainingLogger:
    """Logger specifically for training runs."""
    
    def __init__(self, run_name: Optional[str] = None, log_dir: str = "logs"):
        """
        Initialize training logger.
        
        Args:
            run_name: Name for this training run
            log_dir: Directory for log files
        """
        if run_name is None:
            run_name = f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        self.run_name = run_name
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        log_file = self.log_dir / f"{run_name}.log"
        self.logger = setup_logger(
            name=f"training.{run_name}",
            log_file=str(log_file)
        )
    
    def info(self, message: str) -> None:
        """Log info message."""
        self.logger.info(message)
    
    def warning(self, message: str) -> None:
        """Log warning message."""
        self.logger.warning(message)
    
    def error(self, message: str) -> None:
        """Log error message."""
        self.logger.error(message)
    
    def debug(self, message: str) -> None:
        """Log debug message."""
        self.logger.debug(message)
    
    def log_config(self, config: dict) -> None:
        """Log configuration."""
        self.info("=" * 80)
        self.info("Training Configuration:")
        self.info("=" * 80)
        for key, value in config.items():
            if isinstance(value, dict):
                self.info(f"{key}:")
                for k, v in value.items():
                    self.info(f"  {k}: {v}")
            else:
                self.info(f"{key}: {value}")
        self.info("=" * 80)
    
    def log_episode(self, episode: int, metrics: dict) -> None:
        """Log episode metrics."""
        metric_str = ", ".join([f"{k}={v:.4f}" for k, v in metrics.items()])
        self.info(f"Episode {episode}: {metric_str}")