"""High-level training orchestration for the logistic classifier."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from llm_detector.baselines import (
    BaselineCache,
    BaselineSet,
    build_default_registry_baselines,
    divergence_baseline_overrides,
)
from llm_detector.features import FeatureVectorizer, register_default_features
from llm_detector.models import LogisticRegressionModel

from .dataset import build_dataset_from_registry
from .sources.registry import DEFAULT_REGISTRY, SourceRegistry


@dataclass(slots=True)
class TrainingArtifacts:
    """Outputs generated by the training pipeline."""

    model_path: Path
    baselines_path: Path | None
    train_accuracy: float
    metrics: dict[str, float]
    feature_names: list[str]


def _ensure_baselines(
    *,
    registry: SourceRegistry,
    baseline_path: Path | None,
    samples_per_source: int | None,
    version: str,
    enabled_only: bool,
) -> tuple[BaselineSet, Path | None]:
    if baseline_path is None:
        baselines = build_default_registry_baselines(
            registry=registry,
            samples_per_source=samples_per_source,
            version=version,
            enabled_only=enabled_only,
        )
        return baselines, None

    cache = Path(baseline_path)
    baselines = BaselineCache.get_or_compute(
        cache,
        lambda: build_default_registry_baselines(
            registry=registry,
            samples_per_source=samples_per_source,
            version=version,
            enabled_only=enabled_only,
        ),
    )
    return baselines, cache


def _vectorizer_from_baselines(
    baselines: BaselineSet,
    *,
    scale_invariant_only: bool,
) -> FeatureVectorizer:
    registry = register_default_features()
    overrides = divergence_baseline_overrides(baselines, cohort="human")
    return FeatureVectorizer(
        registry,
        scale_invariant_only=scale_invariant_only,
        baseline_overrides=overrides,
    )


def train_logistic_from_registry(
    *,
    model_path: Path,
    baseline_path: Path | None = None,
    registry: SourceRegistry | None = None,
    samples_per_source: int | None = None,
    baseline_samples_per_source: int | None = None,
    version: str = "v1",
    enabled_only: bool = True,
    balance: bool = True,
    shuffle: bool = True,
    seed: int | None = 42,
    scale_invariant_only: bool = True,
    test_ratio: float | None = 0.1,
    show_progress: bool = False,
) -> TrainingArtifacts:
    """Train a logistic classifier using the configured registry."""

    reg = registry or DEFAULT_REGISTRY

    baselines, cache_path = _ensure_baselines(
        registry=reg,
        baseline_path=baseline_path,
        samples_per_source=baseline_samples_per_source or samples_per_source,
        version=version,
        enabled_only=enabled_only,
    )

    vectorizer = _vectorizer_from_baselines(
        baselines,
        scale_invariant_only=scale_invariant_only,
    )

    dataset = build_dataset_from_registry(
        reg,
        vectorizer,
        samples_per_source=samples_per_source,
        enabled_only=enabled_only,
        balance=balance,
        shuffle=shuffle,
        seed=seed,
        keep_text=False,
        show_progress=show_progress,
    )

    if test_ratio is not None:
        train_set, test_set = dataset.split(test_ratio, seed=seed)
    else:
        train_set, test_set = dataset, None

    model = LogisticRegressionModel()
    result = model.fit(train_set)

    if test_set is not None and len(test_set.matrix) > 0:
        metrics = result.metrics.copy()
        metrics.update({f"test_{k}": v for k, v in model.evaluate(test_set).items()})
    else:
        metrics = result.metrics

    model.save(model_path)

    return TrainingArtifacts(
        model_path=Path(model_path),
        baselines_path=cache_path,
        train_accuracy=result.train_accuracy,
        metrics=metrics,
        feature_names=result.feature_names,
    )


__all__ = ["train_logistic_from_registry", "TrainingArtifacts"]
