src/services/training/training_service.py

"""
Training service for flower classification models.
"""

import os

from core.constants import IMAGES_DIR
from utils.file_utils import count_training_images


class TrainingService:
    """Service for managing model training."""

    def __init__(self):
        pass

    def start_training(
        self, epochs: int = 5, batch_size: int = 8, learning_rate: float = 1e-5
    ) -> str:
        """Start the training process."""
        try:
            # Check if training data exists
            if not os.path.exists(IMAGES_DIR):
                return "❌ Training directory not found. Please create training_data/images/ and add your data."

            # Count images
            total_images, _ = count_training_images()

            if total_images < 10:
                return f"❌ Need at least 10 training images. Found {total_images}. Add more images to training_data/images/"

            # Import and run training (lazy import to avoid startup issues)
            try:
                from training.simple_train import simple_train

                model_path = simple_train()

                if model_path:
                    return f"✅ Training completed! Model saved to: {model_path}"
                else:
                    return "❌ Training failed. Check the console for details."
            except ImportError:
                # Fallback to old training method
                try:
                    from simple_train import simple_train as legacy_train

                    model_path = legacy_train()

                    if model_path:
                        return f"✅ Training completed! Model saved to: {model_path}"
                    else:
                        return "❌ Training failed. Check the console for details."
                except ImportError:
                    return "❌ Training module not found. Please ensure training scripts are available."

        except Exception as e:
            return f"❌ Training error: {e!s}"


# Global service instance
training_service = TrainingService()