"""
Shared YOLO Training Module

Unified training logic for both CLI and Web API.
"""

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

logger = logging.getLogger(__name__)


@dataclass
class TrainingConfig:
    """Training configuration."""

    # Model settings
    model_path: str = "yolo26s.pt"  # Base model or path to trained model
    data_yaml: str = ""  # Path to data.yaml

    # Training hyperparameters
    epochs: int = 100
    batch_size: int = 16
    image_size: int = 640
    learning_rate: float = 0.01
    device: str = "0"

    # Output settings
    project: str = "runs/train"
    name: str = "invoice_fields"

    # Performance settings
    workers: int = 4
    cache: bool = False

    # Resume settings
    resume: bool = False
    resume_from: str | None = None  # Path to checkpoint

    # Fine-tuning specific
    freeze: int = 0  # Number of backbone layers to freeze (0 = none)
    cos_lr: bool = False  # Use cosine learning rate scheduler

    # Document-specific augmentation (optimized for invoices)
    augmentation: dict[str, Any] = field(default_factory=lambda: {
        "degrees": 5.0,
        "translate": 0.05,
        "scale": 0.2,
        "shear": 0.0,
        "perspective": 0.0,
        "flipud": 0.0,
        "fliplr": 0.0,
        "mosaic": 0.0,
        "mixup": 0.0,
        "hsv_h": 0.0,
        "hsv_s": 0.1,
        "hsv_v": 0.2,
    })


@dataclass
class TrainingResult:
    """Training result."""

    success: bool
    model_path: str | None = None
    metrics: dict[str, float] = field(default_factory=dict)
    error: str | None = None
    save_dir: str | None = None


class YOLOTrainer:
    """Unified YOLO trainer for CLI and Web API."""

    def __init__(
        self,
        config: TrainingConfig,
        log_callback: Callable[[str, str], None] | None = None,
    ):
        """
        Initialize trainer.

        Args:
            config: Training configuration
            log_callback: Optional callback for logging (level, message)
        """
        self.config = config
        self._log_callback = log_callback

    def _log(self, level: str, message: str) -> None:
        """Log a message."""
        if self._log_callback:
            self._log_callback(level, message)
        if level == "INFO":
            logger.info(message)
        elif level == "ERROR":
            logger.error(message)
        elif level == "WARNING":
            logger.warning(message)

    def validate_config(self) -> tuple[bool, str | None]:
        """
        Validate training configuration.

        Returns:
            Tuple of (is_valid, error_message)
        """
        # Check model path
        model_path = Path(self.config.model_path)
        if not model_path.suffix == ".pt":
            # Could be a model name like "yolo26s.pt" which is downloaded
            if not model_path.name.startswith("yolo"):
                return False, f"Invalid model: {self.config.model_path}"
        elif not model_path.exists():
            return False, f"Model file not found: {self.config.model_path}"

        # Check data.yaml
        if not self.config.data_yaml:
            return False, "data_yaml is required"
        data_yaml = Path(self.config.data_yaml)
        if not data_yaml.exists():
            return False, f"data.yaml not found: {self.config.data_yaml}"

        return True, None

    def train(self) -> TrainingResult:
        """
        Run YOLO training.

        Returns:
            TrainingResult with model path and metrics
        """
        try:
            from ultralytics import YOLO
        except ImportError:
            return TrainingResult(
                success=False,
                error="Ultralytics (YOLO) not installed. Install with: pip install ultralytics",
            )

        # Validate config
        is_valid, error = self.validate_config()
        if not is_valid:
            return TrainingResult(success=False, error=error)

        self._log("INFO", f"Starting YOLO training")
        self._log("INFO", f"  Model: {self.config.model_path}")
        self._log("INFO", f"  Data: {self.config.data_yaml}")
        self._log("INFO", f"  Epochs: {self.config.epochs}")
        self._log("INFO", f"  Batch size: {self.config.batch_size}")
        self._log("INFO", f"  Image size: {self.config.image_size}")
        if self.config.freeze > 0:
            self._log("INFO", f"  Freeze layers: {self.config.freeze}")
        if self.config.cos_lr:
            self._log("INFO", f"  Cosine LR: enabled")

        try:
            # Load model
            if self.config.resume and self.config.resume_from:
                resume_path = Path(self.config.resume_from)
                if resume_path.exists():
                    self._log("INFO", f"Resuming from: {resume_path}")
                    model = YOLO(str(resume_path))
                else:
                    model = YOLO(self.config.model_path)
            else:
                model = YOLO(self.config.model_path)

            # Build training arguments
            train_args = {
                "data": str(Path(self.config.data_yaml).absolute()),
                "epochs": self.config.epochs,
                "batch": self.config.batch_size,
                "imgsz": self.config.image_size,
                "lr0": self.config.learning_rate,
                "device": self.config.device,
                "project": self.config.project,
                "name": self.config.name,
                "exist_ok": True,
                "pretrained": True,
                "verbose": True,
                "workers": self.config.workers,
                "cache": self.config.cache,
                "resume": self.config.resume and self.config.resume_from is not None,
            }

            # Add fine-tuning settings
            if self.config.freeze > 0:
                train_args["freeze"] = self.config.freeze
            if self.config.cos_lr:
                train_args["cos_lr"] = True

            # Add augmentation settings
            train_args.update(self.config.augmentation)

            # Train
            results = model.train(**train_args)

            # Get best model path
            best_model = Path(results.save_dir) / "weights" / "best.pt"

            # Extract metrics
            metrics = {}
            if hasattr(results, "results_dict"):
                metrics = {
                    "mAP50": results.results_dict.get("metrics/mAP50(B)", 0),
                    "mAP50-95": results.results_dict.get("metrics/mAP50-95(B)", 0),
                    "precision": results.results_dict.get("metrics/precision(B)", 0),
                    "recall": results.results_dict.get("metrics/recall(B)", 0),
                }

            self._log("INFO", f"Training completed successfully")
            self._log("INFO", f"  Best model: {best_model}")
            self._log("INFO", f"  mAP@0.5: {metrics.get('mAP50', 'N/A')}")

            return TrainingResult(
                success=True,
                model_path=str(best_model) if best_model.exists() else None,
                metrics=metrics,
                save_dir=str(results.save_dir),
            )

        except Exception as e:
            self._log("ERROR", f"Training failed: {e}")
            return TrainingResult(success=False, error=str(e))

    def validate(self, split: str = "val") -> dict[str, float]:
        """
        Run validation on trained model.

        Args:
            split: Dataset split to validate on ("val" or "test")

        Returns:
            Validation metrics
        """
        try:
            from ultralytics import YOLO

            model = YOLO(self.config.model_path)
            metrics = model.val(data=self.config.data_yaml, split=split)

            return {
                "mAP50": metrics.box.map50,
                "mAP50-95": metrics.box.map,
                "precision": metrics.box.mp,
                "recall": metrics.box.mr,
            }
        except Exception as e:
            self._log("ERROR", f"Validation failed: {e}")
            return {}
