"""Base trainer class for SalTTS training."""

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path
from typing import Optional, Dict

from saltts.utils.checkpoint import CheckpointManager
from saltts.utils.logging import setup_logger


class BaseTrainer:
    """Base trainer class with common training utilities."""

    def __init__(
        self,
        model: nn.Module,
        optimizer: torch.optim.Optimizer,
        device: str = "cuda",
        checkpoint_dir: str = "checkpoints",
        log_dir: str = "logs",
    ):
        """Initialize base trainer.

        Args:
            model: Model to train
            optimizer: Optimizer
            device: Device (cuda or cpu)
            checkpoint_dir: Checkpoint directory
            log_dir: Log directory
        """
        self.model = model.to(device)
        self.optimizer = optimizer
        self.device = device

        # Checkpoint manager
        self.checkpoint_manager = CheckpointManager(checkpoint_dir)

        # Logger
        self.logger = setup_logger("trainer", log_dir=log_dir)

        # TensorBoard writer
        self.writer = SummaryWriter(log_dir)

        # Training state
        self.epoch = 0
        self.global_step = 0

    def train_epoch(self, dataloader) -> Dict[str, float]:
        """Train for one epoch.

        Args:
            dataloader: Training dataloader

        Returns:
            Dictionary of metrics
        """
        raise NotImplementedError

    def validate(self, dataloader) -> Dict[str, float]:
        """Validate model.

        Args:
            dataloader: Validation dataloader

        Returns:
            Dictionary of metrics
        """
        raise NotImplementedError

    def save_checkpoint(self, metrics: Optional[Dict[str, float]] = None) -> None:
        """Save checkpoint.

        Args:
            metrics: Training metrics
        """
        self.checkpoint_manager.save(
            self.model,
            self.optimizer,
            self.epoch,
            metrics,
        )

    def load_checkpoint(self, checkpoint_path: str) -> None:
        """Load checkpoint.

        Args:
            checkpoint_path: Path to checkpoint
        """
        checkpoint = self.checkpoint_manager.load(
            checkpoint_path,
            self.model,
            self.optimizer,
        )
        self.epoch = checkpoint.get("epoch", 0)
        self.global_step = checkpoint.get("global_step", 0)

    def log_metrics(self, metrics: Dict[str, float], prefix: str = "") -> None:
        """Log metrics to TensorBoard.

        Args:
            metrics: Metrics dictionary
            prefix: Prefix for metric names
        """
        for key, value in metrics.items():
            self.writer.add_scalar(f"{prefix}/{key}", value, self.global_step)
