"""Unit tests for snapshot export."""
from __future__ import annotations

import json
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock

import pytest

from libs.export.snapshot_export import _temporal_split, export_dataset_snapshot


@pytest.mark.unit
class TestTemporalSplit:
    def test_split_proportions(self) -> None:
        """Temporal split produces correct proportions for 100 rows."""
        rows = [{"event_date": "2026-01-01"} for _ in range(100)]
        splits = _temporal_split(rows, "temporal_70_15_15")
        assert len(splits["train"]) == 70
        assert len(splits["valid"]) == 15
        assert len(splits["test"]) == 15

    def test_split_preserves_temporal_order(self) -> None:
        """Train set contains earliest dates, test contains latest."""
        rows = [{"event_date": f"2026-{m:02d}-01"} for m in range(1, 13)]
        splits = _temporal_split(rows, "temporal_70_15_15")
        if splits["train"] and splits["test"]:
            assert splits["train"][-1]["event_date"] <= splits["test"][0]["event_date"]

    def test_empty_rows_returns_empty_splits(self) -> None:
        splits = _temporal_split([], "temporal_70_15_15")
        assert splits == {"train": [], "valid": [], "test": []}


@pytest.mark.unit
class TestExportDatasetSnapshot:
    @pytest.mark.asyncio
    async def test_manifest_is_written(self) -> None:
        """export_dataset_snapshot writes a manifest.json with expected fields."""
        # Mock DB session returning empty results (no feature+label pairs)
        mock_session = AsyncMock()
        mock_result = MagicMock()
        mock_result.all.return_value = []  # no rows
        mock_session.execute = AsyncMock(return_value=mock_result)

        with tempfile.TemporaryDirectory() as tmpdir:
            manifest = await export_dataset_snapshot(
                session=mock_session,
                snapshot_id="test-snapshot-001",
                split_policy="temporal_70_15_15",
                output_dir=tmpdir,
            )

        assert manifest["snapshot_id"] == "test-snapshot-001"
        assert "created_at_utc" in manifest
        assert "row_counts" in manifest
        assert manifest["split_policy"] == "temporal_70_15_15"
        assert manifest["total_rows"] == 0

    @pytest.mark.asyncio
    async def test_parquet_files_created(self) -> None:
        """Parquet files are created for each split partition."""
        mock_session = AsyncMock()
        mock_result = MagicMock()
        mock_result.all.return_value = []
        mock_session.execute = AsyncMock(return_value=mock_result)

        with tempfile.TemporaryDirectory() as tmpdir:
            await export_dataset_snapshot(
                session=mock_session,
                snapshot_id="test-parquet-002",
                split_policy="temporal_70_15_15",
                output_dir=tmpdir,
            )

            snap_dir = Path(tmpdir) / "test-parquet-002"
            assert (snap_dir / "train.parquet").exists()
            assert (snap_dir / "valid.parquet").exists()
            assert (snap_dir / "test.parquet").exists()
            assert (snap_dir / "manifest.json").exists()

            manifest_data = json.loads((snap_dir / "manifest.json").read_text())
            assert manifest_data["snapshot_id"] == "test-parquet-002"