import pytest
from pathlib import Path
import numpy as np
import nibabel as nib

from cad_benchmark_space.validator import validate_submission


@pytest.fixture
def mock_gt():
    return {
        "canonical_orientation": "RAS+",
        "cases": {
            "case001": {
                "labels": {"2c": 1, "2d": 0},
                "shape": (32, 64, 64),
                "spacing": (1.5, 0.75, 0.75),
                "orientation": ("R", "A", "S"),
                "findings": {
                    "case001_f0": {"class": "2c", "mask_path": "masks/case001_f0.nii.gz"}
                },
                "gt_masks": {"2c": np.zeros((32, 64, 64), dtype=bool), "2d": np.zeros((32, 64, 64), dtype=bool)},
            },
            "case002": {
                "labels": {"2c": 0, "2d": 1},
                "shape": (32, 64, 64),
                "spacing": (1.5, 0.75, 0.75),
                "orientation": ("R", "A", "S"),
                "findings": {
                    "case002_f0": {"class": "2d", "mask_path": "masks/case002_f0.nii.gz"}
                },
                "gt_masks": {"2c": np.zeros((32, 64, 64), dtype=bool), "2d": np.zeros((32, 64, 64), dtype=bool)},
            },
        },
    }


def test_reject_invalid_model_type(mock_gt):
    manifest = {"model_name": "TestModel", "model_type": "invalid_type"}
    res = validate_submission(manifest, mock_gt)
    assert not res.is_valid
    assert any("model_type" in err for err in res.errors)


def test_classification_reject_floats(mock_gt):
    manifest = {
        "model_name": "TestModel",
        "model_type": "classification",
        "predictions": [
            {"scan_id": "case001", "2c": 0.85, "2d": 0}
        ],
    }
    res = validate_submission(manifest, mock_gt)
    assert not res.is_valid
    assert any("must be int 0 or 1" in err for err in res.errors)


def test_classification_reject_unknown_scan_id(mock_gt):
    manifest = {
        "model_name": "TestModel",
        "model_type": "classification",
        "predictions": [
            {"scan_id": "case999", "2c": 1, "2d": 0}
        ],
    }
    res = validate_submission(manifest, mock_gt)
    assert not res.is_valid
    assert any("case999" in err for err in res.errors)


def test_classification_reject_duplicates(mock_gt):
    manifest = {
        "model_name": "TestModel",
        "model_type": "classification",
        "predictions": [
            {"scan_id": "case001", "2c": 1, "2d": 0},
            {"scan_id": "case001", "2c": 0, "2d": 1},
        ],
    }
    res = validate_submission(manifest, mock_gt)
    assert not res.is_valid
    assert any("duplicate scan_id" in err for err in res.errors)


def test_classification_valid(mock_gt):
    manifest = {
        "model_name": "TestModel",
        "model_type": "classification",
        "predictions": [
            {"scan_id": "case001", "2c": 1, "2d": 0},
            {"scan_id": "case002", "2c": 0, "2d": 1},
        ],
    }
    res = validate_submission(manifest, mock_gt)
    assert res.is_valid
    assert len(res.classification_samples) == 2


def test_segmentation_validation(mock_gt, tmp_path):
    shape = (32, 64, 64)
    spacing = (1.5, 0.75, 0.75)
    affine = np.diag([spacing[0], spacing[1], spacing[2], 1.0])
    mask_data = np.zeros(shape, dtype=np.uint8)
    mask_data[10:15, 20:30, 20:30] = 1

    mask_file = tmp_path / "case001_f0.nii.gz"
    nii = nib.Nifti1Image(mask_data, affine)
    nib.save(nii, str(mask_file))

    manifest = {
        "model_name": "SegModel",
        "model_type": "segmentation",
        "predictions": [
            {
                "scan_id": "case001",
                "finding_id": "case001_f0",
                "class": "2c",
                "mask_path": "case001_f0.nii.gz",
                "is_soft_mask": False,
            }
        ],
    }

    res = validate_submission(manifest, mock_gt, files_dir=tmp_path)
    assert res.is_valid
    assert len(res.localization_samples) == 1


def test_segmentation_reject_shape_mismatch(mock_gt, tmp_path):
    wrong_shape = (20, 64, 64)
    spacing = (1.5, 0.75, 0.75)
    affine = np.diag([spacing[0], spacing[1], spacing[2], 1.0])
    mask_data = np.zeros(wrong_shape, dtype=np.uint8)

    mask_file = tmp_path / "case001_wrong.nii.gz"
    nii = nib.Nifti1Image(mask_data, affine)
    nib.save(nii, str(mask_file))

    manifest = {
        "model_name": "SegModel",
        "model_type": "segmentation",
        "predictions": [
            {
                "scan_id": "case001",
                "finding_id": "case001_f0",
                "class": "2c",
                "mask_path": "case001_wrong.nii.gz",
            }
        ],
    }

    res = validate_submission(manifest, mock_gt, files_dir=tmp_path)
    assert not res.is_valid
    assert any("mask shape" in err for err in res.errors)