"""Logging utilities for Myanmar Ghost project."""

import logging
import sys
from datetime import datetime
from pathlib import Path
from typing import Optional

from loguru import logger as _logger


def setup_logger(
    name: str = "myanmar_ghost",
    log_dir: Optional[str] = None,
    level: str = "INFO",
    format: str = None,
) -> logging.Logger:
    """Set up logger with file and console output.
    
    Args:
        name: Logger name
        log_dir: Directory for log files
        level: Logging level
        format: Custom log format
    
    Returns:
        Configured logger
    """
    if format is None:
        format = (
            "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
            "<level>{level: <8}</level> | "
            "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | "
            "<level>{message}</level>"
        )
    
    # Remove default handler
    _logger.remove()
    
    # Console output
    _logger.add(
        sys.stdout,
        format=format,
        level=level,
        colorize=True,
    )
    
    # File output
    if log_dir:
        log_path = Path(log_dir)
        log_path.mkdir(parents=True, exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = log_path / f"{name}_{timestamp}.log"
        
        _logger.add(
            log_file,
            format=format,
            level=level,
            rotation="100 MB",
            retention="30 days",
            compression="zip",
        )
    
    return _logger


def get_logger(name: str = None) -> logging.Logger:
    """Get logger instance.
    
    Args:
        name: Logger name (optional)
    
    Returns:
        Logger instance
    """
    return _logger


class TrainingLogger:
    """Logger for training metrics and progress."""
    
    def __init__(
        self,
        log_dir: str = "outputs/logs",
        experiment_name: str = "experiment",
    ):
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        self.experiment_name = experiment_name
        self.metrics_history = []
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_file = self.log_dir / f"{experiment_name}_{timestamp}.log"
    
    def log_metrics(self, metrics: dict, step: int) -> None:
        """Log metrics at a specific step."""
        entry = {
            "step": step,
            "timestamp": datetime.now().isoformat(),
            **metrics,
        }
        self.metrics_history.append(entry)
        
        log_line = f"Step {step}: " + ", ".join(
            f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
            for k, v in metrics.items()
        )
        _logger.info(log_line)
    
    def log_epoch(self, epoch: int, metrics: dict) -> None:
        """Log metrics at epoch end."""
        entry = {
            "epoch": epoch,
            "timestamp": datetime.now().isoformat(),
            **metrics,
        }
        self.metrics_history.append(entry)
        
        log_line = f"Epoch {epoch}: " + ", ".join(
            f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
            for k, v in metrics.items()
        )
        _logger.info(log_line)
    
    def save_history(self) -> str:
        """Save metrics history to file."""
        import json
        
        with open(self.log_file, "w", encoding="utf-8") as f:
            json.dump(self.metrics_history, f, indent=2)
        
        return str(self.log_file)


if __name__ == "__main__":
    logger = setup_logger("test_logger", "outputs/logs")
    logger.info("Logger initialized successfully")