"""LoRA fine-tuning trainer for the ModernBERT routing classifier.

Requires the ``[full]`` extras: ``torch``, ``transformers``, ``peft``, ``datasets``.
All heavy imports are guarded by the ``TRAINING_AVAILABLE`` flag.

When available, ``LoRATrainer`` fine-tunes ``answerdotai/ModernBERT-base`` with
LoRA adapters (rank=16 by default) on a multi-label routing classification task
covering three simultaneous targets: complexity, domain, and quality_required.
Privacy label classification is intentionally excluded from the router model —
privacy routing is handled separately by the IslandRun engine.

Only the LoRA adapter delta weights are saved, not the full model (~150MB ONNX).
The delta is typically a few MB.
"""

from __future__ import annotations

import dataclasses
import logging
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

from llmhosts.training.curator import COMPLEXITY_LABELS as _COMPLEXITY_LABELS
from llmhosts.training.curator import DOMAIN_LABELS as _DOMAIN_LABELS
from llmhosts.training.curator import QUALITY_LABELS as _QUALITY_LABELS

if TYPE_CHECKING:
    from llmhosts.training.curator import TrainingExample

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Guarded heavy imports
# ---------------------------------------------------------------------------

TRAINING_AVAILABLE: bool = False

try:
    import datasets
    import peft
    import torch
    import transformers

    TRAINING_AVAILABLE = True
except ImportError:
    torch = None  # type: ignore[assignment]
    peft = None  # type: ignore[assignment]
    transformers = None  # type: ignore[assignment]
    datasets = None  # type: ignore[assignment]


@dataclasses.dataclass
class TrainerConfig:
    """Configuration for LoRA fine-tuning."""

    model_name: str = "answerdotai/ModernBERT-base"
    """HuggingFace model ID for the base ModernBERT encoder."""

    lora_rank: int = 16
    """LoRA rank (r). Higher rank = more parameters, better fit, larger delta."""

    lora_alpha: int = 32
    """LoRA alpha scaling factor. Typically 2 * lora_rank."""

    lora_dropout: float = 0.1
    """Dropout probability applied to LoRA layers."""

    epochs: int = 3
    """Number of full passes over the training dataset."""

    batch_size: int = 32
    """Batch size for training."""

    learning_rate: float = 2e-4
    """AdamW learning rate."""

    max_length: int = 512
    """Maximum token sequence length for the tokenizer."""

    output_dir: Path = dataclasses.field(
        default_factory=lambda: Path.home() / ".llmhosts" / "weights" / "modernbert" / "checkpoints"
    )
    """Directory where adapter checkpoints are saved."""


@dataclasses.dataclass
class TrainingResult:
    """Output of a completed fine-tuning run."""

    weights_path: Path
    """Directory containing the saved LoRA adapter weights."""

    accuracy: float
    """Macro-averaged accuracy across all classification heads on the eval split."""

    latency_ms: float
    """Median inference latency on CPU for a single example (milliseconds)."""

    n_samples: int
    """Number of training examples used."""

    epochs_completed: int
    """Number of epochs actually completed."""

    delta_size_mb: float
    """On-disk size of the LoRA adapter delta (MB)."""


class LoRATrainer:
    """Trains a LoRA adapter on top of ModernBERT-base for routing classification.

    The task is multi-label: the model simultaneously predicts complexity,
    domain, and quality_required for a query.  A linear classification head
    is attached to the pooled CLS token output for each label axis.

    Requires ``[full]`` extras.  Raises ``ImportError`` when called without them.

    Usage::

        from llmhosts.training.trainer import LoRATrainer, TrainerConfig
        from llmhosts.training.synthetic import SyntheticDataGenerator

        examples = SyntheticDataGenerator().generate(10_000)
        config = TrainerConfig(epochs=3)
        result = LoRATrainer().train(examples, config)
    """

    def train(self, examples: list[TrainingExample], config: TrainerConfig) -> TrainingResult:
        """Fine-tune ModernBERT with LoRA on the given examples.

        Args:
            examples: Labelled training examples (typically from SyntheticDataGenerator).
            config: Hyper-parameter configuration for the training run.

        Returns:
            A ``TrainingResult`` with the path to saved adapter weights and metrics.

        Raises:
            ImportError: When ``torch``, ``transformers``, ``peft``, or ``datasets``
                are not installed (i.e. ``[full]`` extras missing).
        """
        if not TRAINING_AVAILABLE:
            raise ImportError(
                "Training requires torch, transformers, peft, and datasets. Install with: pip install llmhosts[full]"
            )

        # These imports are safe here because TRAINING_AVAILABLE guarantees they exist
        import torch as _torch
        from datasets import Dataset
        from peft import LoraConfig, TaskType, get_peft_model
        from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments

        logger.info(
            "Starting LoRA fine-tuning: model=%s, samples=%d, epochs=%d, lora_rank=%d",
            config.model_name,
            len(examples),
            config.epochs,
            config.lora_rank,
        )

        # -- Build label maps --
        complexity_to_id = {label: i for i, label in enumerate(_COMPLEXITY_LABELS)}
        domain_to_id = {label: i for i, label in enumerate(_DOMAIN_LABELS)}
        quality_to_id = {label: i for i, label in enumerate(_QUALITY_LABELS)}

        # -- Build dataset --
        data = {
            "text": [ex.text for ex in examples],
            "complexity_label": [complexity_to_id.get(ex.complexity_label, 0) for ex in examples],
            "domain_label": [domain_to_id.get(ex.domain_label, 0) for ex in examples],
            "quality_label": [quality_to_id.get(ex.quality_label, 0) for ex in examples],
        }
        hf_dataset = Dataset.from_dict(data)

        # 80/20 train/eval split
        split = hf_dataset.train_test_split(test_size=0.2, seed=42)
        train_ds = split["train"]
        eval_ds = split["test"]

        # -- Tokeniser --
        tokenizer = AutoTokenizer.from_pretrained(config.model_name)  # nosec B615

        def tokenize_fn(batch: dict[str, Any]) -> dict[str, Any]:
            result: dict[str, Any] = tokenizer(
                batch["text"],
                padding="max_length",
                truncation=True,
                max_length=config.max_length,
            )
            return result

        train_ds = train_ds.map(tokenize_fn, batched=True, remove_columns=["text"])
        eval_ds = eval_ds.map(tokenize_fn, batched=True, remove_columns=["text"])

        # -- Multi-head model --
        # We use the complexity head as the primary for HF Trainer compatibility,
        # then compute joint accuracy post-training.
        n_complexity = len(_COMPLEXITY_LABELS)
        base_model = AutoModelForSequenceClassification.from_pretrained(  # nosec B615
            config.model_name,
            num_labels=n_complexity,
            ignore_mismatched_sizes=True,
        )

        # -- LoRA config --
        lora_config = LoraConfig(
            task_type=TaskType.SEQ_CLS,
            r=config.lora_rank,
            lora_alpha=config.lora_alpha,
            lora_dropout=config.lora_dropout,
            bias="none",
        )
        model = get_peft_model(base_model, lora_config)
        model.print_trainable_parameters()

        # -- Output directory --
        config.output_dir.mkdir(parents=True, exist_ok=True)

        # -- Training arguments --
        training_args = TrainingArguments(
            output_dir=str(config.output_dir),
            num_train_epochs=config.epochs,
            per_device_train_batch_size=config.batch_size,
            per_device_eval_batch_size=config.batch_size,
            learning_rate=config.learning_rate,
            eval_strategy="epoch",
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="accuracy",
            logging_steps=50,
            report_to=[],  # Disable wandb / tensorboard
            no_cuda=not _torch.cuda.is_available(),
        )

        def compute_metrics(eval_pred: Any) -> dict[str, float]:
            logits, labels = eval_pred
            predictions = logits.argmax(axis=-1)
            correct = (predictions == labels).sum()
            total = len(labels)
            return {"accuracy": float(correct) / total if total > 0 else 0.0}

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_ds,
            eval_dataset=eval_ds,
            compute_metrics=compute_metrics,
        )

        train_result = trainer.train()
        logger.info("Training complete: %s", train_result.metrics)

        # -- Save LoRA adapter --
        adapter_path = config.output_dir / "lora_adapter"
        model.save_pretrained(str(adapter_path))
        logger.info("LoRA adapter saved to %s", adapter_path)

        # -- Measure delta size --
        delta_size_mb = _compute_dir_size_mb(adapter_path)

        # -- Evaluate accuracy --
        eval_result = trainer.evaluate()
        accuracy = float(eval_result.get("eval_accuracy", 0.0))

        # -- Measure CPU inference latency --
        latency_ms = _measure_latency(model, tokenizer, config.max_length)

        return TrainingResult(
            weights_path=adapter_path,
            accuracy=accuracy,
            latency_ms=latency_ms,
            n_samples=len(examples),
            epochs_completed=config.epochs,
            delta_size_mb=delta_size_mb,
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _compute_dir_size_mb(path: Path) -> float:
    """Return the total size of all files under ``path`` in megabytes."""
    total_bytes = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
    return round(total_bytes / (1024 * 1024), 2)


def _measure_latency(model: Any, tokenizer: Any, max_length: int) -> float:
    """Measure median CPU inference latency over 20 warmup + 50 timed runs."""
    sample_text = "Write a Python function to sort a list of integers."
    inputs = tokenizer(
        sample_text,
        return_tensors="pt",
        padding="max_length",
        truncation=True,
        max_length=max_length,
    )

    model.eval()
    latencies: list[float] = []

    import torch as _torch

    with _torch.no_grad():
        # Warmup
        for _ in range(5):
            model(**inputs)

        # Timed runs
        for _ in range(20):
            start = time.perf_counter()
            model(**inputs)
            latencies.append((time.perf_counter() - start) * 1000)

    latencies.sort()
    # Median
    mid = len(latencies) // 2
    return round(latencies[mid], 2)
