from __future__ import annotations

from pathlib import Path

from kneiff.datasets.export.types import (
    TrainingExportResult,
    TrainingSubsetExportResult,
)
from kneiff.training.lora.sampling_report import build_training_sampling_result


def _export_result() -> TrainingExportResult:
    """Return export counts with one non-training generated SFW subset."""
    return TrainingExportResult(
        export_root=Path("/tmp/export"),
        source_image_count=4,
        output_image_count=28,
        caption_file_count=28,
        min_outputs_per_source=2,
        max_outputs_per_source=14,
        average_outputs_per_source=7.0,
        subset_results=(
            _subset("fullbody", output_images=10),
            _subset("details", output_images=10),
            _subset("sfw", output_images=8),
        ),
        dry_run=False,
    )


def _subset(name: str, *, output_images: int) -> TrainingSubsetExportResult:
    """Return one subset export count row."""
    return TrainingSubsetExportResult(
        subset=name,
        source_image_count=output_images,
        output_image_count=output_images,
        caption_file_count=output_images,
        min_outputs_per_source=1,
        max_outputs_per_source=1,
        average_outputs_per_source=1.0,
    )


def test_training_sampling_report_uses_auto_weighting_and_excludes_sfw(
    tmp_path: Path,
) -> None:
    config_path = tmp_path / "CONFIG-test.yaml"
    config_path.write_text(
        """
training:
  enabled: true
  simpletuner:
    trainer:
      data_backend_sampling: auto-weighting
    subsets:
      fullbody:
        probability: 2.0
      details:
        probability: 1.0
""",
        encoding="utf-8",
    )

    result = build_training_sampling_result(
        project_config_path=config_path,
        export_result=_export_result(),
    )

    assert result is not None
    by_subset = result.by_subset()
    assert set(by_subset) == {"fullbody", "details"}
    assert by_subset["fullbody"].probability == 2.0
    assert by_subset["fullbody"].relative_percent == 66.66666666666666
    assert by_subset["details"].relative_percent == 33.33333333333333


def test_training_sampling_report_includes_sfw_when_training_config_lists_it(
    tmp_path: Path,
) -> None:
    config_path = tmp_path / "CONFIG-test.yaml"
    config_path.write_text(
        """
training:
  enabled: true
  simpletuner:
    trainer:
      data_backend_sampling: auto-weighting
    subsets:
      fullbody:
        probability: 2.0
      details:
        probability: 1.0
      sfw:
        probability: 1.0
""",
        encoding="utf-8",
    )

    result = build_training_sampling_result(
        project_config_path=config_path,
        export_result=_export_result(),
    )

    assert result is not None
    by_subset = result.by_subset()
    assert set(by_subset) == {"fullbody", "details", "sfw"}
    assert round(by_subset["sfw"].relative_percent, 1) == 21.1


def test_training_sampling_report_uses_uniform_probability_weights(
    tmp_path: Path,
) -> None:
    config_path = tmp_path / "CONFIG-test.yaml"
    config_path.write_text(
        """
training:
  enabled: true
  simpletuner:
    trainer:
      data_backend_sampling: uniform
    subsets:
      fullbody:
        probability: 2.0
      details:
        probability: 1.0
""",
        encoding="utf-8",
    )

    result = build_training_sampling_result(
        project_config_path=config_path,
        export_result=_export_result(),
    )

    assert result is not None
    by_subset = result.by_subset()
    assert by_subset["fullbody"].relative_percent == 66.66666666666666
    assert by_subset["details"].relative_percent == 33.33333333333333