"""
pipeline.py
-----------
High-level orchestrator that wires preprocessing → fine-tuning → checkpoint
cleanup into a single ``train()`` call.

This class is fully self-contained: it has no dependency on any internal
framework class.  ``BasePipeline`` has been replaced with a lightweight
abstract base defined in this module.
"""

from __future__ import annotations

import logging
import shutil
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any

from fiesta.modelling import (
    MultilingualE5LargeInstructModelling,
    LoraSettings,
    TrainingSettings,
    TrainResult,
)
from fiesta.preprocessing import (
    MultilingualE5LargeInstructPreProcessing,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Minimal abstract base (replaces the removed internal BasePipeline)
# ---------------------------------------------------------------------------

class BasePipeline(ABC):
    """
    Minimal abstract base for training pipelines.

    Subclasses must implement :meth:`train`.
    """

    @abstractmethod
    def train(self, data: list[dict[str, Any]], **kwargs: Any) -> Any:
        """Run the full training pipeline and return results."""


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

class MultilingualE5LargeInstructPipeline(BasePipeline):
    """
    End-to-end pipeline for fine-tuning *intfloat/multilingual-e5-large-instruct*.

    Parameters
    ----------
    kb_id:
        Knowledge-base identifier (used to organise the save directory).
    lang:
        Language code (e.g. ``"en"``, ``"pt"``).
    base_dir:
        Root directory where model checkpoints are saved.
        Defaults to ``~/models/fiesta/``.
    lora_settings:
        Optional :class:`~fiesta.modelling.LoraSettings`
        to override LoRA defaults.
    training_settings:
        Optional :class:`~fiesta.modelling.TrainingSettings`
        to override training hyper-parameters.
    preprocessing_settings:
        Optional kwargs forwarded to
        :class:`~fiesta.preprocessing.MultilingualE5LargeInstructPreProcessing`.

    Example
    -------
    >>> from fiesta import MultilingualE5LargeInstructPipeline
    >>> pipeline = MultilingualE5LargeInstructPipeline(kb_id="my-kb", lang="en")
    >>> results = pipeline.train(raw_data)
    >>> print(results.improvement())
    """

    LEARNER_NAME: str = "fiesta"

    def __init__(
        self,
        kb_id: str,
        lang: str,
        base_dir: str | Path | None = None,
        lora_settings: LoraSettings | None = None,
        training_settings: TrainingSettings | None = None,
        preprocessing_settings: dict[str, Any] | None = None,
    ) -> None:
        self.kb_id = kb_id
        self.lang = lang.lower()

        # Resolve save directory ─────────────────────────────────────────────
        root = Path(base_dir) if base_dir else Path.home() / "models"
        self.save_dir: Path = (
            root
            / self.LEARNER_NAME
            / self.lang
            / f"{kb_id.lower()}.{self.lang}"
        )
        self.save_dir.mkdir(parents=True, exist_ok=True)
        logger.info("Model will be saved to: %s", self.save_dir)

        # Sub-components ─────────────────────────────────────────────────────
        self.modeller = MultilingualE5LargeInstructModelling(
            lora_settings=lora_settings,
            training_settings=training_settings,
        )
        pp_kwargs: dict[str, Any] = preprocessing_settings or {}
        self.preprocessor = MultilingualE5LargeInstructPreProcessing(**pp_kwargs)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def train(
        self,
        data: list[dict[str, Any]],
        **kwargs: Any,
    ) -> TrainResult:
        """
        Run the full fine-tuning pipeline.

        Parameters
        ----------
        data:
            List of knowledge-base samples.  See
            :mod:`fiesta.preprocessing` for the
            expected schema.

        Returns
        -------
        TrainResult
            Dataclass with ``baseline``, ``final_metrics``, and an
            ``improvement()`` helper method.
        """
        logger.info("Starting preprocessing ...")
        train_ds, test_ds, evaluator_data = self.preprocessor.preprocess(data)

        logger.info("Starting training ...")
        result = self.modeller.train(
            train_dataset=train_ds,
            test_dataset=test_ds,
            evaluator_data=evaluator_data,
            save_dir=str(self.save_dir),
        )

        self._delete_checkpoints()
        return result

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _delete_checkpoints(self) -> None:
        """
        Remove all checkpoint sub-directories from ``save_dir`` using
        cross-platform :func:`shutil.rmtree` (replaces the old ``rm -rf``
        shell command).
        """
        checkpoint_dirs = list(self.save_dir.glob("*checkpoint*"))
        if not checkpoint_dirs:
            logger.debug("No checkpoint directories found in %s.", self.save_dir)
            return

        for ckpt_dir in checkpoint_dirs:
            if ckpt_dir.is_dir():
                try:
                    shutil.rmtree(ckpt_dir)
                    logger.info("Deleted checkpoint: %s", ckpt_dir)
                except OSError as exc:
                    logger.error(
                        "Failed to delete checkpoint %s: %s", ckpt_dir, exc
                    )



