import sys
import tempfile
from multiprocessing import Process
from pathlib import Path
from typing import List

import pytest
import requests
import toml
import zarr

from ultrack.cli.main import main
from ultrack.config import load_config
from ultrack.utils.data import make_config_content


def _run_server(instance_config_path: str):
    _run_command(["server", "--port", "54123", "-cfg", instance_config_path])


def _run_command(command_and_args: List[str]) -> None:
    try:
        main(command_and_args)
    except SystemExit as exit:
        assert exit.code == 0, f"{command_and_args} failed with exit code {exit.code}"


@pytest.mark.usefixtures("zarr_dataset_paths")
@pytest.mark.parametrize(
    "instance_config_path",
    [
        {},  # defaults
        {
            "segmentation.n_workers": 2,
            "linking.n_workers": 2,
            "data.n_workers": 2,
        },
    ],
    indirect=True,
)
class TestCommandLine:
    @pytest.fixture(scope="class")
    def instance_config_path(self, request) -> str:
        """Created this fixture so configuration are shared between methods."""
        with tempfile.TemporaryDirectory() as tmpdir:
            kwargs = {"data.working_dir": tmpdir}
            if hasattr(request, "param"):
                kwargs.update(request.param)

            content = make_config_content(kwargs)
            config_path = f"{tmpdir}/config.toml"
            with open(config_path, mode="w") as f:
                toml.dump(content, f)
            yield config_path

    def test_segment(
        self, instance_config_path: str, zarr_dataset_paths: List[str]
    ) -> None:
        _run_command(
            [
                "segment",
                "-cfg",
                instance_config_path,
                "-fl",
                "foreground",
                "-cl",
                "contours",
                "-il",
                "labels",
                "-il",
                "contours",
                "-p",
                "intensity_mean",
            ]
            + zarr_dataset_paths
        )

    def test_add_flow(self, instance_config_path: str, tmp_path: Path) -> None:
        config = load_config(instance_config_path)
        tmp_path = tmp_path / "flow.zarr"
        tmp_store = zarr.storage.LocalStore(tmp_path)
        zarr.zeros((2,) + tuple(config.data_config.metadata["shape"]), store=tmp_store)
        _run_command(["add_flow", "-cfg", str(instance_config_path), str(tmp_path)])

    def test_link_iou(self, instance_config_path: str) -> None:
        _run_command(["link", "-cfg", str(instance_config_path)])

    def test_link_with_images(
        self, instance_config_path: str, zarr_dataset_paths: List[str]
    ) -> None:
        # using foreground and contours layer to simulate image channel
        _run_command(
            ["link", "-cfg", str(instance_config_path), "-ow"] + zarr_dataset_paths[:2]
        )

    def test_fit_and_add_probs(
        self, instance_config_path: str, tmp_path: Path, zarr_dataset_paths: List[str]
    ) -> None:
        # required by match gt with model output
        pytest.importorskip("catboost")
        pytest.importorskip("sklearn")

        model_path = tmp_path / "model.pkl"
        new_cfg_path = tmp_path / "new_config.toml"

        _run_command(
            [
                "match_gt",
                "-cfg",
                instance_config_path,
                "-gl",
                "labels",
                "-om",
                str(model_path),
                "-oc",
                str(new_cfg_path),
                "--is-tracking",
                "--is-segmentation",
                "--persistence",
            ]
            + zarr_dataset_paths
        )

        # testing loading new config
        load_config(new_cfg_path)

        for var in ["nodes", "links"]:
            _run_command(
                [
                    "add_probs",
                    str(model_path),
                    "-cfg",
                    instance_config_path,
                    "--persistence",
                    "--var",
                    var,
                ]
            )

    def test_solve(self, instance_config_path: str) -> None:
        with pytest.warns(UserWarning):
            # batch index with overwrite should trigger warning
            _run_command(["solve", "-cfg", instance_config_path, "-ow", "-b", "0"])

    def test_summary(self, instance_config_path: str, tmp_path: Path) -> None:
        _run_command(["data_summary", "-cfg", instance_config_path, "-o", tmp_path])

    def test_ctc_export(self, instance_config_path: str, tmp_path: Path) -> None:
        _run_command(
            [
                "export",
                "ctc",
                "-cfg",
                instance_config_path,
                "-s",
                "0.5,1,1",
                "-ma",
                "5",
                "-di",
                "1",
                "-o",
                str(tmp_path / "01_RES"),
            ]
        )

    def test_ilp_export(self, instance_config_path: str, tmp_path: Path) -> None:
        _run_command(
            [
                "export",
                "lp",
                "-cfg",
                instance_config_path,
                "-o",
                str(tmp_path / "model.lp"),
            ]
        )

    def test_trackmate_export(self, instance_config_path: str, tmp_path: Path) -> None:
        _run_command(
            [
                "export",
                "trackmate",
                "-cfg",
                instance_config_path,
                "-o",
                str(tmp_path / "tracks.xml"),
            ]
        )

    def test_zarr_napari_export(
        self,
        instance_config_path: str,
        tmp_path: Path,
        zarr_dataset_paths: List[str],
    ) -> None:
        _run_command(
            [
                "export",
                "zarr-napari",
                "-cfg",
                instance_config_path,
                "-o",
                str(tmp_path / "results"),
                "-i",
                zarr_dataset_paths[2],
            ]
        )

    @pytest.mark.parametrize("mode", ["gt", "solutions", "links", "all"])
    def test_clear_database(self, instance_config_path: str, mode: str) -> None:
        _run_command(
            [
                "clear_database",
                mode,
                "-cfg",
                instance_config_path,
            ]
        )

    @pytest.mark.skipif(
        sys.platform.startswith("darwin"),
        reason="Not supported on OSX",
    )
    def test_server(self, instance_config_path: str) -> None:
        # Start server in a background thread
        process = Process(target=_run_server, args=(instance_config_path,))
        process.start()

        # Wait for server to start
        import time

        time.sleep(10)

        response = requests.get("http://127.0.0.1:54123")
        print(response.content)

        assert process.is_alive()
        assert response.status_code == 200

        process.terminate()
        process.join()


def test_create_config(tmp_path: Path) -> None:
    _run_command(["create_config", str(tmp_path / "config.toml")])


def test_estimate_params(zarr_dataset_paths: List[str], tmp_path: Path) -> None:
    _run_command(["estimate_params", zarr_dataset_paths[2], "-o", str(tmp_path)])


def test_labels_to_contours(zarr_dataset_paths: List[str], tmp_path: Path) -> None:
    _run_command(
        ["labels_to_contours", zarr_dataset_paths[2], "-o", str(tmp_path / "output")]
    )


def test_check_gurobi() -> None:
    _run_command(["check_gurobi"])
