import os
import pytest
from httpx import AsyncClient
from main import app
import h5py
import numpy as np

from loguru import logger
from survos2.model import DataModel
from survos2.api.workspace import create as create_workspace
from survos2.api.workspace import add_dataset, add_data, delete, get
from survos2.improc.utils import DatasetManager

tmp_ws_name = "testworkspace_tmp1"


@pytest.fixture(scope="session")
def datamodel():
    # make test vol
    map_fullpath = os.path.join("./tmp/testvol_4x4x4b.h5")
    testvol = np.array(
        [
            [
                [0.1761602, 0.6701295, 0.13151232, 0.95726678],
                [0.4795476, 0.48114134, 0.0410548, 0.29893265],
                [0.49127266, 0.70298447, 0.42751211, 0.08101552],
                [0.73805652, 0.83111601, 0.36852477, 0.38732476],
            ],
            [
                [0.2847222, 0.96054574, 0.25430756, 0.35403861],
                [0.54439093, 0.65897414, 0.1959487, 0.90714872],
                [0.84462152, 0.90754182, 0.02455657, 0.26180662],
                [0.1711208, 0.40122666, 0.54562598, 0.01419861],
            ],
            [
                [0.59280376, 0.42706895, 0.86637913, 0.87831645],
                [0.57991401, 0.31989204, 0.85869799, 0.6333411],
                [0.21539274, 0.63780214, 0.64204493, 0.74425482],
                [0.1903691, 0.81962537, 0.31774673, 0.34812628],
            ],
            [
                [0.40880077, 0.595773, 0.28856063, 0.19316746],
                [0.03195766, 0.62475541, 0.50762591, 0.34700798],
                [0.98913461, 0.07883111, 0.96534233, 0.57697606],
                [0.71496714, 0.70764578, 0.92294417, 0.91300531],
            ],
        ]
    )
    with h5py.File(map_fullpath, "w") as hf:
        hf.create_dataset("data", data=testvol)

    result = get(workspace=tmp_ws_name)
    if not isinstance(result, bool):
        logger.debug("tmp exists, deleting and recreating")
        delete(workspace=tmp_ws_name)
        logger.debug("workspace deleted")

    create_workspace(workspace=tmp_ws_name)
    logger.debug("workspace recreated")
    add_data(workspace=tmp_ws_name, data_fname=map_fullpath)
    DataModel.g.current_workspace = tmp_ws_name

    return DataModel


class Tests(object):
    @pytest.mark.asyncio
    async def test_feature_gaussian_blur(self, datamodel):
        async with AsyncClient(app=app, base_url="http://test") as ac:
            DataModel = datamodel
            src = DataModel.g.dataset_uri("__data__", None)
            dst = DataModel.g.dataset_uri("001_gaussian_blur", group="features")

            params = {"workspace": "testworkspace_tmp1", "src": src, "dst": dst, "sigma": 1}
            response = await ac.get("/features/gaussian_blur", params=params)
            assert (
                response.text
                == '{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[1],"source":"__data__","id":"001_gaussian_blur"}'
            )

            with DatasetManager(src, out=dst, dtype="float32", fillvalue=0) as DM:
                src_dataset = DM.sources[0]
                dst_dataset = DM.out
                src_arr = src_dataset[:]
                gblur_arr = dst_dataset[:]

            assert src_arr.shape == gblur_arr.shape

    @pytest.mark.asyncio
    async def test_feature_composite(self, datamodel):
        async with AsyncClient(app=app, base_url="http://test") as ac:
            DataModel = datamodel
            src = DataModel.g.dataset_uri("__data__", None)
            dst = DataModel.g.dataset_uri("001_gaussian_blur", group="features")

            params = {"workspace": "testworkspace_tmp1", "src": src, "dst": dst, "sigma": 1}
            response = await ac.get("/features/gaussian_blur", params=params)
            assert (
                response.text
                == '{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[1],"source":"__data__","id":"001_gaussian_blur"}'
            )

            dst = DataModel.g.dataset_uri("002_gaussian_blur", group="features")
            params = {"workspace": "testworkspace_tmp1", "src": src, "dst": dst, "sigma": 2}
            response = await ac.get("/features/gaussian_blur", params=params)
            assert (
                response.text
                == '{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[2],"source":"__data__","id":"002_gaussian_blur"}'
            )

            params = {"workspace": "testworkspace_tmp1"}
            response = await ac.get("/features/existing", params=params)
            assert (
                response.text
                == '{"001_gaussian_blur":{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[1],"source":"__data__","id":"001_gaussian_blur"},"002_gaussian_blur":{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[2],"source":"__data__","id":"002_gaussian_blur"}}'
            )

            src = DataModel.g.dataset_uri("__data__", None)
            dst = DataModel.g.dataset_uri("003_feature_composite", group="features")

            params = {
                "workspace": "testworkspace_tmp1",
                "src": src,
                "dst": dst,
                "op": "+",
                "feature_A": "001_gaussian_blur",
                "feature_B": "002_gaussian_blur",
            }
            response = await ac.get("/features/feature_composite", params=params)
            assert (
                response.text
                == '{"kind":"feature_composite","name":"feature_composite","workspace":"testworkspace_tmp1","feature_A":"001_gaussian_blur","feature_B":"002_gaussian_blur","op":"+","source":"__data__","id":"003_feature_composite"}'
            )

            with DatasetManager(src, out=dst, dtype="float32", fillvalue=0) as DM:
                src_dataset = DM.sources[0]
                dst_dataset = DM.out
                src_arr = src_dataset[:]
                gblur_arr = dst_dataset[:]

            assert src_arr.shape == gblur_arr.shape

    @pytest.mark.asyncio
    async def test_remove(self, datamodel):
        async with AsyncClient(app=app, base_url="http://test") as ac:
            params = {"workspace": "testworkspace_tmp1", "feature_id": "003_feature_composite"}
            response = await ac.get("/features/remove", params=params)

            params = {"workspace": "testworkspace_tmp1"}
            response = await ac.get("/features/existing", params=params)
            assert (
                response.text
                == '{"001_gaussian_blur":{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[1],"source":"__data__","id":"001_gaussian_blur"},"002_gaussian_blur":{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[2],"source":"__data__","id":"002_gaussian_blur"}}'
            )

    @pytest.mark.asyncio
    async def test_rename(self, datamodel):
        async with AsyncClient(app=app, base_url="http://test") as ac:
            DataModel = datamodel
            src = DataModel.g.dataset_uri("__data__", None)
            dst = DataModel.g.dataset_uri("001_gaussian_blur", group="features")

            params = {
                "workspace": "testworkspace_tmp1",
                "src": src,
                "feature_id": "001_gaussian_blur",
                "new_name": "001_gblur",
            }
            response = await ac.get("/features/rename", params=params)

            params = {"workspace": "testworkspace_tmp1"}
            response = await ac.get("/features/existing", params=params)
            assert (
                response.text
                == '{"001_gaussian_blur":{"kind":"gaussian_blur","name":"001_gblur","sigma":[1],"source":"__data__","id":"001_gaussian_blur"},"002_gaussian_blur":{"kind":"gaussian_blur","name":"gaussian_blur","sigma":[2],"source":"__data__","id":"002_gaussian_blur"}}'
            )


if __name__ == "__main__":
    pytest.main()
