import json
import logging
from pathlib import Path
from typing import Literal

import numpy as np
import pytest
from ndtiff import Dataset
from tifffile import TiffFile

from iohub.convert import TIFFConverter, _adjust_chunks_for_divisibility
from iohub.mm_fov import MicroManagerFOV
from iohub.ngff import Position, open_ome_zarr
from iohub.reader import MMStack, NDTiffDataset
from tests.conftest import (
    mm2gamma_ome_tiffs,
    ndtiff_v2_datasets,
    ndtiff_v3_labeled_positions,
)


def pytest_generate_tests(metafunc):
    if "mm2gamma_ome_tiff" in metafunc.fixturenames:
        metafunc.parametrize("mm2gamma_ome_tiff", mm2gamma_ome_tiffs)
    if "ndtiff_datasets" in metafunc.fixturenames:
        metafunc.parametrize(
            "ndtiff_datasets",
            ndtiff_v2_datasets + [ndtiff_v3_labeled_positions],
        )
    if "grid_layout" in metafunc.fixturenames:
        metafunc.parametrize("grid_layout", [True, False])
    if "chunks" in metafunc.fixturenames:
        metafunc.parametrize(
            "chunks", ["XY", "XYZ", (1, 1, 3, 256, 256), None]
        )


def _check_scale_transform(fov: MicroManagerFOV, position: Position) -> None:
    """Check scale transformation of the highest resolution level."""
    tf = (
        position.metadata.multiscales[0]
        .datasets[0]
        .coordinate_transformations[0]
    )
    assert tf.type == "scale"
    assert tf.scale[0] == fov.t_scale
    assert tf.scale[1] == 1.0
    assert tf.scale[2:] == list(fov.zyx_scale)


def _check_chunks(
    position: Position, chunks: Literal["XY", "XYZ"] | tuple[int] | None
) -> None:
    """Check chunk size of the highest resolution level."""
    img = position["0"]
    match chunks:
        case "XY":
            assert img.chunks == (1,) * 3 + img.shape[-2:]
        case "XYZ" | None:
            assert img.chunks == (1,) * 2 + img.shape[-3:]
        case tuple():
            expected = _adjust_chunks_for_divisibility(img.shape, chunks)
            assert img.chunks == expected
        case _:
            assert False


def _check_result(
    output: Path,
    expected_sum: float,
    converter: TIFFConverter,
    grid_layout: bool,
    chunks: Literal["XY", "XYZ"] | tuple[int] | None,
) -> None:
    with open_ome_zarr(output, mode="r") as result:
        intensity = 0
        if grid_layout and converter.p > 1:
            assert len(result) < converter.p
        for (_, example_fov), (pos_name, pos) in zip(
            converter.reader, result.positions(), strict=True
        ):
            _check_scale_transform(example_fov, pos)
            _check_chunks(pos, chunks)
            intensity += pos["0"][:].sum()
            with open(
                output / pos_name / "0" / "image_plane_metadata.json"
            ) as f:
                metadata = json.load(f)
                key = "0/0/0"
                assert key in metadata
                assert "ElapsedTime-ms" in metadata[key]
        assert intensity == expected_sum


def test_converter_inputs(mm2gamma_ome_tiff, tmpdir):
    # Test that output directory needs to end with ".zarr"
    output = tmpdir / "converted"
    with pytest.raises(ValueError):
        _ = TIFFConverter(mm2gamma_ome_tiff, output)


def test_converter_ometiff(mm2gamma_ome_tiff, grid_layout, chunks, tmpdir):
    logging.getLogger("tifffile").setLevel(logging.ERROR)
    output = tmpdir / "converted.zarr"
    converter = TIFFConverter(
        mm2gamma_ome_tiff, output, grid_layout=grid_layout, chunks=chunks
    )
    assert isinstance(converter.reader, MMStack)
    with TiffFile(next(mm2gamma_ome_tiff.glob("*.tif*"))) as tf:
        raw_array = tf.asarray()
        assert (
            converter.summary_metadata == tf.micromanager_metadata["Summary"]
        )
    assert np.prod([d for d in converter.dim if d > 0]) == np.prod(
        raw_array.shape
    )
    assert list(converter.metadata.keys()) == [
        "iohub_version",
        "Summary",
    ]
    converter()
    _check_result(
        output=output,
        expected_sum=raw_array.sum(),
        converter=converter,
        grid_layout=grid_layout,
        chunks=chunks,
    )


@pytest.fixture(scope="module")
def example_ome_tiff() -> Path:
    for d in mm2gamma_ome_tiffs:
        if d.name == "mm2.0-20201209_4p_2t_5z_1c_512k_1":
            return d
    raise FileNotFoundError("Corrupted test data directory")


@pytest.fixture(scope="function")
def mock_hcs_ome_tiff_reader(
    example_ome_tiff, monkeypatch: pytest.MonkeyPatch
):
    # dataset with 4 positions without HCS site names
    mock_stage_positions = [
        {"Label": "A1-Site_0"},
        {"Label": "A1-Site_1"},
        {"Label": "B4-Site_0"},
        {"Label": "H12-Site_0"},
    ]
    expected_ngff_name = {"A/1/0", "A/1/1", "B/4/0", "H/12/0"}
    monkeypatch.setattr(
        "iohub.convert.MMStack.stage_positions",
        mock_stage_positions,
    )
    return example_ome_tiff, expected_ngff_name


@pytest.fixture(scope="function")
def mock_non_hcs_ome_tiff_reader(
    example_ome_tiff, monkeypatch: pytest.MonkeyPatch
):
    # dataset with 4 positions without HCS site names
    mock_stage_positions = [
        {"Label": "0"},
        {"Label": "1"},
        {"Label": "2"},
        {"Label": "3"},
    ]
    monkeypatch.setattr(
        "iohub.convert.MMStack.stage_positions",
        mock_stage_positions,
    )
    return example_ome_tiff


def test_converter_ometiff_mock_hcs(mock_hcs_ome_tiff_reader, tmpdir):
    data, expected_ngff_name = mock_hcs_ome_tiff_reader
    output = tmpdir / "converted.zarr"
    converter = TIFFConverter(data, output, hcs_plate=True)
    converter()
    with open_ome_zarr(output, mode="r") as plate:
        assert expected_ngff_name == {name for name, _ in plate.positions()}


def test_converter_ometiff_mock_non_hcs(mock_non_hcs_ome_tiff_reader, tmpdir):
    data = mock_non_hcs_ome_tiff_reader
    output = tmpdir / "converted.zarr"
    with pytest.raises(ValueError, match="HCS position label"):
        TIFFConverter(data, output, hcs_plate=True)


def test_converter_ometiff_hcs_numerical(example_ome_tiff, tmpdir):
    output = tmpdir / "converted.zarr"
    converter = TIFFConverter(example_ome_tiff, output, hcs_plate=True)
    converter()
    with open_ome_zarr(output, mode="r") as plate:
        for name, _ in plate.positions():
            for segment in name.split("/"):
                assert segment.isdigit()


def test_converter_ndtiff(ndtiff_datasets: Path, grid_layout, chunks, tmpdir):
    logging.getLogger("tifffile").setLevel(logging.ERROR)
    output = tmpdir / "converted.zarr"
    converter = TIFFConverter(
        ndtiff_datasets, output, grid_layout=grid_layout, chunks=chunks
    )
    assert isinstance(converter.reader, NDTiffDataset)
    raw_array = np.asarray(Dataset(str(ndtiff_datasets)).as_array())
    assert np.prod([d for d in converter.dim if d > 0]) == np.prod(
        raw_array.shape
    )
    assert list(converter.metadata.keys()) == [
        "iohub_version",
        "Summary",
    ]
    converter()
    _check_result(
        output=output,
        expected_sum=raw_array.sum(),
        converter=converter,
        grid_layout=grid_layout,
        chunks=chunks,
    )


def test_converter_ndtiff_v3_position_labels(tmpdir):
    output = tmpdir / "converted.zarr"
    converter = TIFFConverter(ndtiff_v3_labeled_positions, output)
    converter()
    with open_ome_zarr(output, mode="r") as result:
        assert result.channel_names == ["Channel0"]
        assert [name.split("/")[1] for name, _ in result.positions()] == [
            "Pos0",
            "Pos1",
            "Pos2",
        ]


class TestGenChunks:
    """Unit tests for the TIFFConverter._gen_chunks method."""

    @pytest.fixture
    def make_mock_converter(self):
        """Factory fixture to create mock converter with dimensions."""

        def _make(t=1, c=1, z=10, y=256, x=256, dtype="uint16"):
            class MockConverter:
                pass

            converter = MockConverter()
            converter.t = t
            converter.c = c
            converter.z = z
            converter.y = y
            converter.x = x

            class MockReader:
                pass

            converter.reader = MockReader()
            converter.reader.dtype = dtype

            converter._gen_chunks = TIFFConverter._gen_chunks.__get__(
                converter
            )
            return converter

        return _make

    @pytest.mark.parametrize(
        "z_dim,input_chunk,expected_z",
        [
            (15, 10, 5),  # 10 doesn't divide 15, adjust to 5
            (7, 5, 1),  # prime - falls to 1
            (16, 8, 8),  # divides evenly, no change
            (100, 50, 50),  # divides evenly
            (9, 4, 3),  # odd non-prime, adjust to 3
        ],
    )
    def test_gen_chunks_divisibility(
        self, make_mock_converter, z_dim, input_chunk, expected_z
    ):
        """Test chunks are adjusted to divide evenly into dimensions."""
        converter = make_mock_converter(z=z_dim)
        chunks = converter._gen_chunks((1, 1, input_chunk, 256, 256))
        assert chunks[2] == expected_z
        assert z_dim % chunks[2] == 0

    def test_gen_chunks_max_chunk_size_then_divisibility(
        self, make_mock_converter
    ):
        """Test divisibility check happens AFTER MAX_CHUNK_SIZE adjustment."""
        # Large chunks that trigger MAX_CHUNK_SIZE, then need divisibility fix
        # Z=15, large XY to trigger size limit, chunk should end up dividing 15
        converter = make_mock_converter(z=15, y=2048, x=2048)
        chunks = converter._gen_chunks((1, 1, 15, 2048, 2048))
        assert 15 % chunks[2] == 0

    @pytest.mark.parametrize("z_dim", [7, 11, 13, 17, 19, 23])  # prime numbers
    def test_gen_chunks_prime_dimensions(self, make_mock_converter, z_dim):
        """Test handling of prime dimension sizes - should fall to 1."""
        converter = make_mock_converter(z=z_dim)
        chunks = converter._gen_chunks((1, 1, z_dim - 1, 256, 256))
        assert z_dim % chunks[2] == 0

    @pytest.mark.parametrize("input_chunks", ["XY", "XYZ", None])
    def test_gen_chunks_string_inputs(self, make_mock_converter, input_chunks):
        """Test string chunk specifications."""
        converter = make_mock_converter(z=10, y=256, x=256)
        chunks = converter._gen_chunks(input_chunks)
        assert len(chunks) == 5
        assert all(isinstance(c, (int, np.integer)) for c in chunks)
