from __future__ import annotations

import zipfile
from pathlib import Path

import numpy as np
import pytest

from k1link.datasets.rellis_admission import RellisAdmissionError, _member_index
from k1link.datasets.rellis_profile import RellisPatchworkProfile
from k1link.datasets.rellis_qualification import calibrate_rellis_sensor_height


def _frame(height_m: float) -> tuple[bytes, bytes]:
    point_count = 512
    points = np.zeros((point_count, 4), dtype="<f4")
    points[:, 0] = np.linspace(3.0, 12.0, point_count, dtype=np.float32)
    points[:, 2] = -height_m
    points[:, 3] = 0.5
    labels = np.full(point_count, 1, dtype="<u4")
    return points.tobytes(), labels.tobytes()


def test_rellis_height_calibration_uses_only_stratified_train_frames(
    tmp_path: Path,
) -> None:
    scan_path = tmp_path / "scans.zip"
    label_path = tmp_path / "labels.zip"
    pairs: list[tuple[str, str]] = []
    with (
        zipfile.ZipFile(scan_path, "w") as scans,
        zipfile.ZipFile(label_path, "w") as labels,
    ):
        for sequence, height in enumerate((1.20, 1.25, 1.30, 1.35, 1.40)):
            sequence_id = f"{sequence:05d}"
            point_member = f"{sequence_id}/os1_cloud_node_kitti_bin/000000.bin"
            label_member = (
                f"{sequence_id}/os1_cloud_node_semantickitti_label_id/000000.label"
            )
            point_bytes, label_bytes = _frame(height)
            scans.writestr(f"Rellis-3D/{point_member}", point_bytes)
            labels.writestr(f"Rellis-3D/{label_member}", label_bytes)
            pairs.append((point_member, label_member))

    result = calibrate_rellis_sensor_height(
        scan_path,
        label_path,
        tuple(pairs),
        archive_identity="a" * 64,
        frame_count=5,
    )

    assert result["split"] == "train"
    assert result["validation_labels_used"] is False
    assert result["frame_count"] == 5
    assert result["sensor_height_m"] == pytest.approx(1.30)
    assert len(result["identity_sha256"]) == 64


def test_rellis_calibration_accepts_the_official_train_sequence_subset(
    tmp_path: Path,
) -> None:
    scan_path = tmp_path / "scans.zip"
    label_path = tmp_path / "labels.zip"
    pairs: list[tuple[str, str]] = []
    train_sequences = ("00000", "00002", "00003", "00004")
    with (
        zipfile.ZipFile(scan_path, "w") as scans,
        zipfile.ZipFile(label_path, "w") as labels,
    ):
        for sequence_id in train_sequences:
            for frame_number in range(2):
                point_member = (
                    f"{sequence_id}/os1_cloud_node_kitti_bin/{frame_number:06d}.bin"
                )
                label_member = (
                    f"{sequence_id}/os1_cloud_node_semantickitti_label_id/"
                    f"{frame_number:06d}.label"
                )
                point_bytes, label_bytes = _frame(1.3)
                scans.writestr(f"Rellis-3D/{point_member}", point_bytes)
                labels.writestr(f"Rellis-3D/{label_member}", label_bytes)
                pairs.append((point_member, label_member))

    result = calibrate_rellis_sensor_height(
        scan_path,
        label_path,
        tuple(pairs),
        archive_identity="b" * 64,
        frame_count=8,
    )

    calibrated_sequences = {
        item["frame_id"].split("-")[1] for item in result["frame_estimates"]
    }
    assert calibrated_sequences == set(train_sequences)
    assert "00001" not in calibrated_sequences


def test_rellis_archive_index_rejects_path_traversal(tmp_path: Path) -> None:
    archive = tmp_path / "unsafe.zip"
    with zipfile.ZipFile(archive, "w") as target:
        target.writestr("../outside.bin", b"data")
    with (
        zipfile.ZipFile(archive) as source,
        pytest.raises(RellisAdmissionError, match="unsafe"),
    ):
        _member_index(source)


def test_rellis_patchwork_profile_rejects_unphysical_height() -> None:
    with pytest.raises(ValueError):
        RellisPatchworkProfile(
            patchwork_sensor_height_proxy_m=0.1,
            calibration_identity_sha256="a" * 64,
            calibration_frame_count=64,
            calibration_median_absolute_deviation_m=0.01,
        )