"""
Admin Training Scheduler

Background scheduler for training tasks using APScheduler.
"""

import logging
import threading
from datetime import datetime
from pathlib import Path
from typing import Any

from backend.data.repositories import (
    TrainingTaskRepository,
    DatasetRepository,
    ModelVersionRepository,
    DocumentRepository,
    AnnotationRepository,
    FineTunePoolRepository,
)
from backend.web.core.task_interface import TaskRunner, TaskStatus
from backend.web.services.storage_helpers import get_storage_helper

logger = logging.getLogger(__name__)


class TrainingScheduler(TaskRunner):
    """Scheduler for training tasks."""

    def __init__(
        self,
        check_interval_seconds: int = 60,
    ):
        """
        Initialize training scheduler.

        Args:
            check_interval_seconds: Interval to check for pending tasks
        """
        self._check_interval = check_interval_seconds
        self._running = False
        self._thread: threading.Thread | None = None
        self._stop_event = threading.Event()
        self._lock = threading.Lock()
        # Repositories
        self._training_tasks = TrainingTaskRepository()
        self._datasets = DatasetRepository()
        self._model_versions = ModelVersionRepository()
        self._documents = DocumentRepository()
        self._annotations = AnnotationRepository()
        self._pool = FineTunePoolRepository()

    @property
    def name(self) -> str:
        """Unique identifier for this runner."""
        return "training_scheduler"

    @property
    def is_running(self) -> bool:
        """Check if the scheduler is currently active."""
        return self._running

    def get_status(self) -> TaskStatus:
        """Get current status of the scheduler."""
        try:
            pending_tasks = self._training_tasks.get_pending()
            pending_count = len(pending_tasks)
        except Exception:
            pending_count = 0

        return TaskStatus(
            name=self.name,
            is_running=self._running,
            pending_count=pending_count,
            processing_count=1 if self._running else 0,
        )

    def start(self) -> None:
        """Start the scheduler."""
        with self._lock:
            if self._running:
                logger.warning("Training scheduler already running")
                return

            self._running = True
            self._stop_event.clear()
            self._thread = threading.Thread(target=self._run_loop, daemon=True)
            self._thread.start()
            logger.info("Training scheduler started")

    def stop(self, timeout: float | None = None) -> None:
        """Stop the scheduler.

        Args:
            timeout: Maximum time to wait for graceful shutdown.
                    If None, uses default of 5 seconds.
        """
        # Minimize lock scope to avoid potential deadlock
        with self._lock:
            if not self._running:
                return

            self._running = False
            self._stop_event.set()
            thread_to_join = self._thread

        effective_timeout = timeout if timeout is not None else 5.0
        if thread_to_join:
            thread_to_join.join(timeout=effective_timeout)

        with self._lock:
            self._thread = None

        logger.info("Training scheduler stopped")

    def _run_loop(self) -> None:
        """Main scheduler loop."""
        while self._running:
            try:
                self._check_pending_tasks()
            except Exception as e:
                logger.error(f"Error in scheduler loop: {e}")

            # Wait for next check interval
            self._stop_event.wait(timeout=self._check_interval)

    def _check_pending_tasks(self) -> None:
        """Check and execute pending training tasks."""
        try:
            tasks = self._training_tasks.get_pending()

            for task in tasks:
                task_id = str(task.task_id)

                # Check if scheduled time has passed
                if task.scheduled_at and task.scheduled_at > datetime.utcnow():
                    continue

                logger.info(f"Starting training task: {task_id}")

                try:
                    dataset_id = getattr(task, "dataset_id", None)
                    self._execute_task(task_id, task.config or {}, dataset_id=dataset_id)
                except Exception as e:
                    logger.error(f"Training task {task_id} failed: {e}")
                    self._training_tasks.update_status(
                        task_id=task_id,
                        status="failed",
                        error_message=str(e),
                    )

        except Exception as e:
            logger.error(f"Error checking pending tasks: {e}")

    def _execute_task(
        self, task_id: str, config: dict[str, Any], dataset_id: str | None = None
    ) -> None:
        """Execute a training task."""
        # Update status to running
        self._training_tasks.update_status(task_id, "running")
        self._training_tasks.add_log(task_id, "INFO", "Training task started")

        # Update dataset training status to running
        if dataset_id:
            self._datasets.update_training_status(
                dataset_id,
                training_status="running",
                active_training_task_id=task_id,
            )

        try:
            # Get training configuration
            model_name = config.get("model_name", "yolo26s.pt")
            base_model_path = config.get("base_model_path")  # For incremental training
            epochs = config.get("epochs", 100)
            batch_size = config.get("batch_size", 16)
            image_size = config.get("image_size", 640)
            learning_rate = config.get("learning_rate", 0.01)
            device = config.get("device", "0")
            project_name = config.get("project_name", "invoice_fields")

            # Get augmentation config if present
            augmentation_config = config.get("augmentation")
            augmentation_multiplier = config.get("augmentation_multiplier", 2)

            # Determine which model to use as base
            is_finetune = bool(base_model_path)
            if is_finetune:
                # Fine-tuning: use existing trained model as base
                if not Path(base_model_path).exists():
                    raise ValueError(f"Base model not found: {base_model_path}")
                effective_model = base_model_path
                # Override parameters for fine-tuning (best practices)
                epochs = config.get("epochs", 10)
                learning_rate = config.get("learning_rate", 0.001)
                self._training_tasks.add_log(
                    task_id, "INFO",
                    f"Fine-tuning from: {base_model_path} "
                    f"(epochs={epochs}, freeze=10, cos_lr=true)",
                )
            else:
                # Train from pretrained model
                effective_model = model_name

            # Use dataset if available, otherwise export from scratch
            if dataset_id:
                dataset = self._datasets.get(dataset_id)
                if not dataset or not dataset.dataset_path:
                    raise ValueError(f"Dataset {dataset_id} not found or has no path")
                data_yaml = str(Path(dataset.dataset_path) / "data.yaml")
                dataset_path = Path(dataset.dataset_path)
                self._training_tasks.add_log(
                    task_id, "INFO",
                    f"Using pre-built dataset: {dataset.name} ({dataset.total_images} images)",
                )
            else:
                export_result = self._export_training_data(task_id)
                if not export_result:
                    raise ValueError("Failed to export training data")
                data_yaml = export_result["data_yaml"]
                dataset_path = Path(data_yaml).parent
                self._training_tasks.add_log(
                    task_id, "INFO",
                    f"Exported {export_result['total_images']} images for training",
                )

            # Apply augmentation if config is provided
            if augmentation_config and self._has_enabled_augmentations(augmentation_config):
                aug_result = self._apply_augmentation(
                    task_id, dataset_path, augmentation_config, augmentation_multiplier
                )
                if aug_result:
                    self._training_tasks.add_log(
                        task_id, "INFO",
                        f"Augmentation complete: {aug_result['augmented_images']} new images "
                        f"(total: {aug_result['total_images']})",
                    )

            # Build mixed dataset for fine-tuning (pool samples + old data)
            if is_finetune and dataset_id:
                data_yaml, dataset_path = self._build_mixed_finetune_dataset(
                    task_id, dataset_path, data_yaml,
                )

            # Run YOLO training
            result = self._run_yolo_training(
                task_id=task_id,
                model_name=effective_model,
                data_yaml=data_yaml,
                epochs=epochs,
                batch_size=batch_size,
                image_size=image_size,
                learning_rate=learning_rate,
                device=device,
                project_name=project_name,
                freeze=10 if is_finetune else 0,
                cos_lr=is_finetune,
            )

            # Update task with results
            self._training_tasks.update_status(
                task_id=task_id,
                status="completed",
                result_metrics=result.get("metrics"),
                model_path=result.get("model_path"),
            )
            self._training_tasks.add_log(task_id, "INFO", "Training completed successfully")

            # Update dataset training status to completed and main status to trained
            if dataset_id:
                self._datasets.update_training_status(
                    dataset_id,
                    training_status="completed",
                    active_training_task_id=None,
                    update_main_status=True,  # Set main status to 'trained'
                )

            # Auto-create model version for the completed training
            model_version = self._create_model_version_from_training(
                task_id=task_id,
                config=config,
                dataset_id=dataset_id,
                result=result,
            )

            # Auto-run gating validation for fine-tuned models
            if is_finetune and model_version:
                self._run_gating_after_finetune(
                    task_id=task_id,
                    model_version=model_version,
                    config=config,
                    data_yaml=data_yaml,
                    result=result,
                )

        except Exception as e:
            logger.error(f"Training task {task_id} failed: {e}")
            self._training_tasks.add_log(task_id, "ERROR", f"Training failed: {e}")
            # Update dataset training status to failed
            if dataset_id:
                self._datasets.update_training_status(
                    dataset_id,
                    training_status="failed",
                    active_training_task_id=None,
                )
            raise

    def _create_model_version_from_training(
        self,
        task_id: str,
        config: dict[str, Any],
        dataset_id: str | None,
        result: dict[str, Any],
    ) -> Any | None:
        """Create a model version entry from completed training.

        Returns the created model version, or None on failure.
        """
        try:
            model_path = result.get("model_path")
            if not model_path:
                logger.warning(f"No model path in training result for task {task_id}")
                return None

            # Get task info for name
            task = self._training_tasks.get(task_id)
            task_name = task.name if task else f"Task {task_id[:8]}"

            # Generate version number based on existing versions
            existing_versions = self._model_versions.get_paginated(limit=1, offset=0)
            version_count = existing_versions[1] if existing_versions else 0
            version = f"v{version_count + 1}.0"

            # Extract metrics from result
            metrics = result.get("metrics", {})
            metrics_mAP = metrics.get("mAP50") or metrics.get("mAP")
            metrics_precision = metrics.get("precision")
            metrics_recall = metrics.get("recall")

            # Get file size if possible
            file_size = None
            model_file = Path(model_path)
            if model_file.exists():
                file_size = model_file.stat().st_size

            # Get document count from dataset if available
            document_count = 0
            if dataset_id:
                dataset = self._datasets.get(dataset_id)
                if dataset:
                    document_count = dataset.total_documents

            # Determine model lineage
            is_finetune = bool(config.get("base_model_path"))
            model_type = "finetune" if is_finetune else "base"
            base_model_version_id = config.get("base_model_version_id") if is_finetune else None
            gating_status = "pending" if is_finetune else "skipped"

            # Create model version
            model_version = self._model_versions.create(
                version=version,
                name=task_name,
                model_path=str(model_path),
                description=f"Auto-created from training task {task_id[:8]}",
                task_id=task_id,
                dataset_id=dataset_id,
                metrics_mAP=metrics_mAP,
                metrics_precision=metrics_precision,
                metrics_recall=metrics_recall,
                document_count=document_count,
                training_config=config,
                file_size=file_size,
                trained_at=datetime.utcnow(),
                model_type=model_type,
                base_model_version_id=base_model_version_id,
                base_training_dataset_id=dataset_id if not is_finetune else None,
                gating_status=gating_status,
            )

            logger.info(
                f"Created model version {version} (ID: {model_version.version_id}) "
                f"from training task {task_id}"
            )
            mAP_display = f"{metrics_mAP:.3f}" if metrics_mAP else "N/A"
            self._training_tasks.add_log(
                task_id, "INFO",
                f"Model version {version} created (mAP: {mAP_display})",
            )

            return model_version

        except Exception as e:
            logger.error(f"Failed to create model version for task {task_id}: {e}")
            self._training_tasks.add_log(
                task_id, "WARNING",
                f"Failed to auto-create model version: {e}",
            )
            return None

    def _build_mixed_finetune_dataset(
        self,
        task_id: str,
        base_dataset_path: Path,
        original_data_yaml: str,
    ) -> tuple[str, Path]:
        """Build a mixed dataset for fine-tuning.

        Combines verified pool samples with randomly sampled old training data.

        Returns:
            Tuple of (data_yaml path, dataset_path) for the mixed dataset.
            Falls back to original if mixing fails or pool is empty.
        """
        try:
            from backend.web.services.data_mixer import build_mixed_dataset

            pool_doc_ids = self._pool.get_all_document_ids(verified_only=True)
            if not pool_doc_ids:
                self._training_tasks.add_log(
                    task_id, "INFO",
                    "No verified pool entries found, using original dataset",
                )
                return original_data_yaml, base_dataset_path

            mixed_output = base_dataset_path.parent / f"mixed_{task_id[:8]}"
            mix_result = build_mixed_dataset(
                pool_document_ids=pool_doc_ids,
                base_dataset_path=base_dataset_path,
                output_dir=mixed_output,
            )

            self._training_tasks.add_log(
                task_id, "INFO",
                f"Data mixing: {mix_result['new_images']} new + "
                f"{mix_result['old_images']} old = {mix_result['total_images']} total "
                f"(ratio: {mix_result['mixing_ratio']}x)",
            )

            return mix_result["data_yaml"], mixed_output

        except Exception as e:
            logger.error(f"Data mixing failed for task {task_id}: {e}")
            self._training_tasks.add_log(
                task_id, "WARNING",
                f"Data mixing failed: {e}. Using original dataset.",
            )
            return original_data_yaml, base_dataset_path

    def _run_gating_after_finetune(
        self,
        task_id: str,
        model_version: Any,
        config: dict[str, Any],
        data_yaml: str,
        result: dict[str, Any],
    ) -> None:
        """Run gating validation after a fine-tune training completes."""
        try:
            from backend.web.services.gating_validator import run_gating_validation

            model_path = result.get("model_path")
            base_model_version_id = config.get("base_model_version_id")
            version_id = str(model_version.version_id)

            self._training_tasks.add_log(
                task_id, "INFO", "Running deployment gating validation...",
            )

            gating_result = run_gating_validation(
                model_version_id=version_id,
                new_model_path=model_path,
                base_model_version_id=base_model_version_id,
                data_yaml=data_yaml,
                task_id=task_id,
            )

            self._training_tasks.add_log(
                task_id, "INFO",
                f"Gating result: {gating_result.overall_status} "
                f"(gate1={gating_result.gate1_status}, gate2={gating_result.gate2_status})",
            )

        except Exception as e:
            logger.error(f"Gating validation failed for task {task_id}: {e}")
            self._training_tasks.add_log(
                task_id, "WARNING",
                f"Gating validation failed: {e}. Model remains in 'pending' state.",
            )

    def _export_training_data(self, task_id: str) -> dict[str, Any] | None:
        """Export training data for a task."""
        from pathlib import Path
        from shared.fields import FIELD_CLASSES
        from backend.web.services.storage_helpers import get_storage_helper

        # Get storage helper for reading images
        storage = get_storage_helper()

        # Get all labeled documents
        documents = self._documents.get_labeled_for_export()

        if not documents:
            self._training_tasks.add_log(task_id, "ERROR", "No labeled documents available")
            return None

        # Create export directory using StorageHelper
        training_base = storage.get_training_data_path()
        if training_base is None:
            self._training_tasks.add_log(task_id, "ERROR", "Storage not configured for local access")
            return None
        export_dir = training_base / task_id
        export_dir.mkdir(parents=True, exist_ok=True)

        # YOLO format directories
        (export_dir / "images" / "train").mkdir(parents=True, exist_ok=True)
        (export_dir / "images" / "val").mkdir(parents=True, exist_ok=True)
        (export_dir / "labels" / "train").mkdir(parents=True, exist_ok=True)
        (export_dir / "labels" / "val").mkdir(parents=True, exist_ok=True)

        # 80/20 train/val split
        total_docs = len(documents)
        train_count = int(total_docs * 0.8)
        train_docs = documents[:train_count]
        val_docs = documents[train_count:]

        total_images = 0
        total_annotations = 0

        # Export documents
        for split, docs in [("train", train_docs), ("val", val_docs)]:
            for doc in docs:
                annotations = self._annotations.get_for_document(str(doc.document_id))

                if not annotations:
                    continue

                for page_num in range(1, doc.page_count + 1):
                    page_annotations = [a for a in annotations if a.page_number == page_num]

                    # Get image from storage
                    doc_id = str(doc.document_id)
                    if not storage.admin_image_exists(doc_id, page_num):
                        continue

                    # Download image and save to export directory
                    image_name = f"{doc.document_id}_page{page_num}.png"
                    dst_image = export_dir / "images" / split / image_name
                    image_content = storage.get_admin_image(doc_id, page_num)
                    dst_image.write_bytes(image_content)
                    total_images += 1

                    # Write YOLO label
                    label_name = f"{doc.document_id}_page{page_num}.txt"
                    label_path = export_dir / "labels" / split / label_name

                    with open(label_path, "w") as f:
                        for ann in page_annotations:
                            line = f"{ann.class_id} {ann.x_center:.6f} {ann.y_center:.6f} {ann.width:.6f} {ann.height:.6f}\n"
                            f.write(line)
                            total_annotations += 1

        # Create data.yaml
        yaml_path = export_dir / "data.yaml"
        yaml_content = f"""path: {export_dir.absolute()}
train: images/train
val: images/val

nc: {len(FIELD_CLASSES)}
names: {list(FIELD_CLASSES.values())}
"""
        yaml_path.write_text(yaml_content)

        return {
            "data_yaml": str(yaml_path),
            "total_images": total_images,
            "total_annotations": total_annotations,
        }

    def _run_yolo_training(
        self,
        task_id: str,
        model_name: str,
        data_yaml: str,
        epochs: int,
        batch_size: int,
        image_size: int,
        learning_rate: float,
        device: str,
        project_name: str,
        freeze: int = 0,
        cos_lr: bool = False,
    ) -> dict[str, Any]:
        """Run YOLO training using shared trainer."""
        from shared.training import YOLOTrainer, TrainingConfig as SharedTrainingConfig

        # Create log callback that writes to DB
        def log_callback(level: str, message: str) -> None:
            self._training_tasks.add_log(task_id, level, message)

        # Create shared training config
        # Note: Model outputs go to local runs/train directory (not STORAGE_BASE_PATH)
        # because models need to be accessible by inference service on any platform
        # Note: workers=0 to avoid multiprocessing issues when running in scheduler thread
        config = SharedTrainingConfig(
            model_path=model_name,
            data_yaml=data_yaml,
            epochs=epochs,
            batch_size=batch_size,
            image_size=image_size,
            learning_rate=learning_rate,
            device=device,
            project="runs/train",
            name=f"{project_name}/task_{task_id[:8]}",
            workers=0,
            freeze=freeze,
            cos_lr=cos_lr,
        )

        # Run training using shared trainer
        trainer = YOLOTrainer(config=config, log_callback=log_callback)
        result = trainer.train()

        if not result.success:
            raise ValueError(result.error or "Training failed")

        return {
            "model_path": result.model_path,
            "metrics": result.metrics,
        }

    def _has_enabled_augmentations(self, aug_config: dict[str, Any]) -> bool:
        """Check if any augmentations are enabled in the config."""
        augmentation_fields = [
            "perspective_warp", "wrinkle", "edge_damage", "stain",
            "lighting_variation", "shadow", "gaussian_blur", "motion_blur",
            "gaussian_noise", "salt_pepper", "paper_texture", "scanner_artifacts",
        ]
        for field in augmentation_fields:
            if field in aug_config:
                field_config = aug_config[field]
                if isinstance(field_config, dict) and field_config.get("enabled", False):
                    return True
        return False

    def _apply_augmentation(
        self,
        task_id: str,
        dataset_path: Path,
        aug_config: dict[str, Any],
        multiplier: int,
    ) -> dict[str, int] | None:
        """Apply augmentation to dataset before training."""
        try:
            from shared.augmentation import DatasetAugmenter

            self._training_tasks.add_log(
                task_id, "INFO",
                f"Applying augmentation with multiplier={multiplier}",
            )

            augmenter = DatasetAugmenter(aug_config)
            result = augmenter.augment_dataset(dataset_path, multiplier=multiplier)

            return result

        except Exception as e:
            logger.error(f"Augmentation failed for task {task_id}: {e}")
            self._training_tasks.add_log(
                task_id, "WARNING",
                f"Augmentation failed: {e}. Continuing with original dataset.",
            )
            return None


# Global scheduler instance
_scheduler: TrainingScheduler | None = None
_scheduler_lock = threading.Lock()


def get_training_scheduler() -> TrainingScheduler:
    """Get the training scheduler instance.

    Uses double-checked locking pattern for thread safety.
    """
    global _scheduler

    if _scheduler is None:
        with _scheduler_lock:
            if _scheduler is None:
                _scheduler = TrainingScheduler()

    return _scheduler


def start_scheduler() -> None:
    """Start the global training scheduler."""
    scheduler = get_training_scheduler()
    scheduler.start()


def stop_scheduler() -> None:
    """Stop the global training scheduler."""
    global _scheduler
    if _scheduler:
        _scheduler.stop()
        _scheduler = None
