"""GPU-accelerated training integration for the self-improvement pipeline.

Wraps fusionagi.gpu.training to provide a self-improvement-aware training
interface that integrates with AutoTrainer and reflective memory.
"""

from __future__ import annotations

from typing import Any, Protocol

from fusionagi._logger import logger


class ReflectiveMemoryLike(Protocol):
    """Protocol for reflective memory access."""

    def get_lessons(self, limit: int = 50) -> list[dict[str, Any]]: ...
    def get_all_heuristics(self) -> dict[str, Any]: ...
    def set_heuristic(self, key: str, value: Any) -> None: ...


def run_gpu_enhanced_training(
    reflective_memory: ReflectiveMemoryLike,
    epochs: int = 10,
    learning_rate: float = 0.01,
) -> dict[str, Any]:
    """Run GPU-accelerated training on reflective memory lessons.

    Optimizes heuristic scoring weights using gradient descent on GPU,
    then applies the learned improvements back to reflective memory.

    Args:
        reflective_memory: Source of training data and target for updates.
        epochs: Number of training epochs.
        learning_rate: Learning rate for optimization.

    Returns:
        Training result dict with loss history and update count.
    """
    try:
        from fusionagi.gpu.training import (
            TrainingConfig,
            run_gpu_training,
        )

        config = TrainingConfig(
            learning_rate=learning_rate,
            epochs=epochs,
        )
        result = run_gpu_training(reflective_memory, config=config)

        if result.weights_updated > 0:
            reflective_memory.set_heuristic(
                "gpu_training_last_loss", result.final_loss
            )
            reflective_memory.set_heuristic(
                "gpu_training_epochs", result.epochs_run
            )

        logger.info(
            "GPU-enhanced training complete",
            extra={
                "initial_loss": result.initial_loss,
                "final_loss": result.final_loss,
                "weights_updated": result.weights_updated,
            },
        )
        return {
            "initial_loss": result.initial_loss,
            "final_loss": result.final_loss,
            "epochs_run": result.epochs_run,
            "weights_updated": result.weights_updated,
            "gpu_accelerated": True,
            "metadata": result.metadata,
        }
    except ImportError:
        logger.debug("GPU training not available; skipping")
        return {
            "gpu_accelerated": False,
            "reason": "GPU dependencies not installed",
        }


def can_gpu_train() -> bool:
    """Check if GPU training is available."""
    try:
        from fusionagi.gpu.backend import get_backend

        get_backend()
        return True
    except ImportError:
        return False