import re
import json
import logging
from unittest.mock import AsyncMock, patch
from pytest import mark, fixture, raises
from aiohttp import WSMessage, WSMsgType

import proximl.datasets as specimen
from proximl.exceptions import (
    ApiError,
    DatasetError,
    SpecificationError,
    ProxiMLException,
)

pytestmark = [mark.sdk, mark.unit, mark.datasets]


@fixture
def datasets(mock_proximl):
    yield specimen.Datasets(mock_proximl)


@fixture
def dataset(mock_proximl):
    yield specimen.Dataset(
        mock_proximl,
        dataset_uuid="1",
        project_uuid="proj-id-1",
        name="first one",
        type="evefs",
        status="downloading",
        size=100000,
        createdAt="2020-12-31T23:59:59.000Z",
    )


class DatasetsTests:
    @mark.asyncio
    async def test_get_dataset(
        self,
        datasets,
        mock_proximl,
    ):
        api_response = dict()
        mock_proximl._query = AsyncMock(return_value=api_response)
        await datasets.get("1234")
        mock_proximl._query.assert_called_once_with(
            "/dataset/1234",
            "GET",
            dict(),
        )

    @mark.asyncio
    async def test_list_datasets(
        self,
        datasets,
        mock_proximl,
    ):
        api_response = dict()
        mock_proximl._query = AsyncMock(return_value=api_response)
        await datasets.list()
        mock_proximl._query.assert_called_once_with(
            "/dataset",
            "GET",
            dict(),
        )

    @mark.asyncio
    async def test_list_public_datasets(self, datasets, mock_proximl):
        api_response = dict()
        mock_proximl._query = AsyncMock(return_value=api_response)
        await datasets.list_public()
        mock_proximl._query.assert_called_once_with(
            "/dataset/public",
            "GET",
            dict(),
        )

    @mark.asyncio
    async def test_remove_dataset(
        self,
        datasets,
        mock_proximl,
    ):
        api_response = dict()
        mock_proximl._query = AsyncMock(return_value=api_response)
        await datasets.remove("4567")
        mock_proximl._query.assert_called_once_with(
            "/dataset/4567",
            "DELETE",
            dict(force=True),
        )

    @mark.asyncio
    async def test_create_dataset_simple(self, datasets, mock_proximl):
        requested_config = dict(
            name="new dataset",
            source_type="aws",
            source_uri="s3://proximl-examples/data/cifar10",
        )
        expected_payload = dict(
            project_uuid="proj-id-1",
            name="new dataset",
            source_type="aws",
            source_uri="s3://proximl-examples/data/cifar10",
            type="evefs",
        )
        api_response = {
            "customer_uuid": "cus-id-1",
            "project_uuid": "proj-id-1",
            "dataset_uuid": "data-id-1",
            "name": "new dataset",
            "status": "new",
            "type": "evefs",
            "source_type": "aws",
            "source_uri": "s3://proximl-examples/data/cifar10",
            "createdAt": "2020-12-20T16:46:23.909Z",
        }

        mock_proximl._query = AsyncMock(return_value=api_response)
        response = await datasets.create(**requested_config)
        mock_proximl._query.assert_called_once_with(
            "/dataset", "POST", None, expected_payload
        )
        assert response.id == "data-id-1"


class DatasetTests:
    def test_dataset_properties(self, dataset):
        assert isinstance(dataset.id, str)
        assert isinstance(dataset.status, str)
        assert isinstance(dataset.name, str)
        assert isinstance(dataset.size, int)

    def test_dataset_str(self, dataset):
        string = str(dataset)
        regex = r"^{.*\"dataset_uuid\": \"" + dataset.id + r"\".*}$"
        assert isinstance(string, str)
        assert re.match(regex, string)

    def test_dataset_repr(self, dataset):
        string = repr(dataset)
        regex = (
            r"^Dataset\( proximl , \*\*{.*'dataset_uuid': '"
            + dataset.id
            + r"'.*}\)$"
        )
        assert isinstance(string, str)
        assert re.match(regex, string)

    def test_dataset_bool(self, dataset, mock_proximl):
        empty_dataset = specimen.Dataset(mock_proximl)
        assert bool(dataset)
        assert not bool(empty_dataset)

    @mark.asyncio
    async def test_dataset_get_log_url(self, dataset, mock_proximl):
        api_response = "https://trainml-jobs-dev.s3.us-east-2.amazonaws.com/1/logs/first_one.zip"
        mock_proximl._query = AsyncMock(return_value=api_response)
        response = await dataset.get_log_url()
        mock_proximl._query.assert_called_once_with(
            "/dataset/1/logs", "GET", dict(project_uuid="proj-id-1")
        )
        assert response == api_response

    @mark.asyncio
    async def test_dataset_get_details(self, dataset, mock_proximl):
        api_response = {
            "type": "directory",
            "name": "/",
            "count": "8",
            "size": "177M",
            "contents": [],
        }
        mock_proximl._query = AsyncMock(return_value=api_response)
        response = await dataset.get_details()
        mock_proximl._query.assert_called_once_with(
            "/dataset/1/details", "GET", dict(project_uuid="proj-id-1")
        )
        assert response == api_response

    @mark.asyncio
    async def test_dataset_connect_downloading_status(self, mock_proximl):
        dataset = specimen.Dataset(
            mock_proximl,
            dataset_uuid="1",
            project_uuid="proj-id-1",
            name="test dataset",
            status="downloading",
            auth_token="test-token",
            hostname="example.com",
            source_uri="/path/to/source",
        )

        with patch(
            "proximl.datasets.Dataset.refresh", new_callable=AsyncMock
        ) as mock_refresh:
            with patch(
                "proximl.datasets.upload", new_callable=AsyncMock
            ) as mock_upload:
                await dataset.connect()
                mock_refresh.assert_called_once()
                mock_upload.assert_called_once_with(
                    "example.com", "test-token", "/path/to/source"
                )

    @mark.asyncio
    async def test_dataset_connect_exporting_status(
        self, mock_proximl, tmp_path
    ):
        output_dir = str(tmp_path / "output")
        dataset = specimen.Dataset(
            mock_proximl,
            dataset_uuid="1",
            project_uuid="proj-id-1",
            name="test dataset",
            status="exporting",
            auth_token="test-token",
            hostname="example.com",
            output_uri=output_dir,
        )

        with patch(
            "proximl.datasets.Dataset.refresh", new_callable=AsyncMock
        ) as mock_refresh:
            with patch(
                "proximl.datasets.download", new_callable=AsyncMock
            ) as mock_download:
                await dataset.connect()
                mock_refresh.assert_called_once()
                mock_download.assert_called_once_with(
                    "example.com", "test-token", output_dir
                )

    @mark.asyncio
    async def test_dataset_connect_invalid_status(self, mock_proximl):
        dataset = specimen.Dataset(
            mock_proximl,
            dataset_uuid="1",
            project_uuid="proj-id-1",
            name="test dataset",
            status="ready",
        )

        with raises(
            SpecificationError,
            match="You can only connect to downloading or exporting datasets",
        ):
            await dataset.connect()

    @mark.asyncio
    async def test_dataset_remove(self, dataset, mock_proximl):
        api_response = dict()
        mock_proximl._query = AsyncMock(return_value=api_response)
        await dataset.remove()
        mock_proximl._query.assert_called_once_with(
            "/dataset/1",
            "DELETE",
            dict(project_uuid="proj-id-1", force=False),
        )

    def test_dataset_default_ws_msg_handler(self, dataset, capsys):
        data = {
            "msg": "download: s3://proximl-examples/data/cifar10/data_batch_2.bin to ./data_batch_2.bin\n",
            "time": 1613079345318,
            "type": "subscription",
            "stream": "worker-id-1",
            "job_worker_uuid": "worker-id-1",
        }

        handler = dataset._get_msg_handler(None)
        handler(data)
        captured = capsys.readouterr()
        assert (
            captured.out
            == "02/11/2021, 15:35:45: download: s3://proximl-examples/data/cifar10/data_batch_2.bin to ./data_batch_2.bin\n"
        )

    def test_dataset_custom_ws_msg_handler(self, dataset, capsys):
        def custom_handler(msg):
            print(msg.get("stream"))

        data = {
            "msg": "download: s3://proximl-examples/data/cifar10/data_batch_2.bin to ./data_batch_2.bin\n",
            "time": 1613079345318,
            "type": "subscription",
            "stream": "worker-id-1",
            "job_worker_uuid": "worker-id-1",
        }

        handler = dataset._get_msg_handler(custom_handler)
        handler(data)
        captured = capsys.readouterr()
        assert captured.out == "worker-id-1\n"

    @mark.asyncio
    async def test_dataset_attach(self, dataset, mock_proximl):
        api_response = None
        mock_proximl._ws_subscribe = AsyncMock(return_value=api_response)
        refresh_response = {
            "customer_uuid": "cus-id-1",
            "dataset_uuid": "data-id-1",
            "name": "new dataset",
            "status": "downloading",
            "source_type": "aws",
            "source_uri": "s3://proximl-examples/data/cifar10",
            "createdAt": "2020-12-20T16:46:23.909Z",
        }
        dataset.refresh = AsyncMock(return_value=refresh_response)
        await dataset.attach()
        mock_proximl._ws_subscribe.assert_called_once()

    @mark.asyncio
    async def test_dataset_attach_immediate_return(self, mock_proximl):
        dataset = specimen.Dataset(
            mock_proximl,
            dataset_uuid="1",
            name="first one",
            status="ready",
            createdAt="2020-12-31T23:59:59.000Z",
        )
        api_response = None
        mock_proximl._ws_subscribe = AsyncMock(return_value=api_response)
        refresh_response = {
            "customer_uuid": "cus-id-1",
            "dataset_uuid": "1",
            "name": "new dataset",
            "status": "ready",
            "createdAt": "2020-12-20T16:46:23.909Z",
        }
        dataset.refresh = AsyncMock(return_value=refresh_response)
        await dataset.attach()
        mock_proximl._ws_subscribe.assert_not_called()

    @mark.asyncio
    async def test_dataset_refresh(self, dataset, mock_proximl):
        api_response = {
            "customer_uuid": "cus-id-1",
            "dataset_uuid": "data-id-1",
            "name": "new dataset",
            "status": "ready",
            "source_type": "aws",
            "source_uri": "s3://proximl-examples/data/cifar10",
            "createdAt": "2020-12-20T16:46:23.909Z",
        }
        mock_proximl._query = AsyncMock(return_value=api_response)
        response = await dataset.refresh()
        mock_proximl._query.assert_called_once_with(
            f"/dataset/1", "GET", dict(project_uuid="proj-id-1")
        )
        assert dataset.id == "data-id-1"
        assert response.id == "data-id-1"

    @mark.asyncio
    async def test_dataset_wait_for_successful(self, dataset, mock_proximl):
        api_response = {
            "customer_uuid": "cus-id-1",
            "dataset_uuid": "data-id-1",
            "name": "new dataset",
            "status": "ready",
            "source_type": "aws",
            "source_uri": "s3://proximl-examples/data/cifar10",
            "createdAt": "2020-12-20T16:46:23.909Z",
        }
        mock_proximl._query = AsyncMock(return_value=api_response)
        response = await dataset.wait_for("ready")
        mock_proximl._query.assert_called_once_with(
            f"/dataset/1", "GET", dict(project_uuid="proj-id-1")
        )
        assert dataset.id == "data-id-1"
        assert response.id == "data-id-1"

    @mark.asyncio
    async def test_dataset_wait_for_current_status(self, mock_proximl):
        dataset = specimen.Dataset(
            mock_proximl,
            dataset_uuid="1",
            name="first one",
            status="ready",
            createdAt="2020-12-31T23:59:59.000Z",
        )
        api_response = None
        mock_proximl._query = AsyncMock(return_value=api_response)
        await dataset.wait_for("ready")
        mock_proximl._query.assert_not_called()

    @mark.asyncio
    async def test_dataset_wait_for_incorrect_status(
        self, dataset, mock_proximl
    ):
        api_response = None
        mock_proximl._query = AsyncMock(return_value=api_response)
        with raises(SpecificationError):
            await dataset.wait_for("stopped")
        mock_proximl._query.assert_not_called()

    @mark.asyncio
    async def test_dataset_wait_for_with_delay(self, dataset, mock_proximl):
        api_response_initial = dict(
            dataset_uuid="1",
            name="first one",
            status="new",
            createdAt="2020-12-31T23:59:59.000Z",
        )
        api_response_final = dict(
            dataset_uuid="1",
            name="first one",
            status="ready",
            createdAt="2020-12-31T23:59:59.000Z",
        )
        mock_proximl._query = AsyncMock()
        mock_proximl._query.side_effect = [
            api_response_initial,
            api_response_initial,
            api_response_final,
        ]
        response = await dataset.wait_for("ready")
        assert dataset.status == "ready"
        assert response.status == "ready"

    @mark.asyncio
    async def test_dataset_wait_for_timeout(self, dataset, mock_proximl):
        api_response = dict(
            dataset_uuid="1",
            name="first one",
            status="new",
            createdAt="2020-12-31T23:59:59.000Z",
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        with raises(ProxiMLException):
            await dataset.wait_for("ready", 10)
        mock_proximl._query.assert_called()

    @mark.asyncio
    async def test_dataset_wait_for_dataset_failed(
        self, dataset, mock_proximl
    ):
        api_response = dict(
            dataset_uuid="1",
            name="first one",
            status="failed",
            createdAt="2020-12-31T23:59:59.000Z",
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        with raises(DatasetError):
            await dataset.wait_for("ready")
        mock_proximl._query.assert_called()

    @mark.asyncio
    async def test_dataset_rename(self, dataset, mock_proximl):
        api_response = dict(
            dataset_uuid="1",
            name="renamed dataset",
            project_uuid="proj-id-1",
            status="ready",
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        result = await dataset.rename("renamed dataset")
        mock_proximl._query.assert_called_once_with(
            "/dataset/1",
            "PATCH",
            None,
            dict(name="renamed dataset"),
        )
        assert result == dataset
        assert dataset.name == "renamed dataset"

    @mark.asyncio
    async def test_dataset_export(self, dataset, mock_proximl):
        api_response = dict(
            dataset_uuid="1",
            name="first one",
            project_uuid="proj-id-1",
            status="exporting",
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        result = await dataset.export("aws", "s3://bucket/path", dict(key="value"))
        mock_proximl._query.assert_called_once_with(
            "/dataset/1/export",
            "POST",
            dict(project_uuid="proj-id-1"),
            dict(
                output_type="aws",
                output_uri="s3://bucket/path",
                output_options=dict(key="value"),
            ),
        )
        assert result == dataset
        assert dataset.status == "exporting"

    @mark.asyncio
    async def test_dataset_export_default_options(self, dataset, mock_proximl):
        api_response = dict(
            dataset_uuid="1",
            name="first one",
            project_uuid="proj-id-1",
            status="exporting",
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        result = await dataset.export("aws", "s3://bucket/path")
        mock_proximl._query.assert_called_once_with(
            "/dataset/1/export",
            "POST",
            dict(project_uuid="proj-id-1"),
            dict(
                output_type="aws",
                output_uri="s3://bucket/path",
                output_options=dict(),
            ),
        )
        assert result == dataset

    @mark.asyncio
    async def test_dataset_wait_for_timeout_validation(
        self, dataset, mock_proximl
    ):
        with raises(SpecificationError) as exc_info:
            await dataset.wait_for("ready", timeout=25 * 60 * 60)  # > 24 hours
        assert "timeout" in str(exc_info.value.attribute).lower()
        assert "less than" in str(exc_info.value.message).lower()

    @mark.asyncio
    async def test_dataset_connect_new_status_waits_for_downloading(
        self, dataset, mock_proximl
    ):
        """Test that connect() waits for downloading status when status is 'new'."""
        dataset._dataset["status"] = "new"
        dataset._status = "new"
        api_response_new = dict(
            dataset_uuid="1",
            name="first one",
            status="new",
        )
        api_response_downloading = dict(
            dataset_uuid="1",
            name="first one",
            status="downloading",
            auth_token="token",
            hostname="host",
            source_uri="s3://bucket/path",
        )
        # wait_for calls refresh multiple times, then connect calls refresh again
        # We need enough responses for wait_for polling and the final refresh
        mock_proximl._query = AsyncMock(
            side_effect=[
                api_response_new,  # wait_for refresh 1
                api_response_downloading,  # wait_for refresh 2 (status matches, wait_for returns)
                api_response_downloading,  # connect refresh
            ]
        )
        with patch("proximl.datasets.upload", new_callable=AsyncMock) as mock_upload:
            await dataset.connect()
        # After refresh, status should be downloading
        assert dataset.status == "downloading"
        mock_upload.assert_called_once()

    @mark.asyncio
    async def test_dataset_connect_downloading_missing_properties(
        self, dataset, mock_proximl
    ):
        """Test connect() raises error when downloading status missing properties."""
        dataset._dataset["status"] = "downloading"
        api_response = dict(
            dataset_uuid="1",
            name="first one",
            status="downloading",
            # Missing auth_token, hostname, or source_uri
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        with raises(SpecificationError) as exc_info:
            await dataset.connect()
        assert "missing required connection properties" in str(exc_info.value.message).lower()

    @mark.asyncio
    async def test_dataset_connect_exporting_missing_properties(
        self, dataset, mock_proximl
    ):
        """Test connect() raises error when exporting status missing properties."""
        dataset._dataset["status"] = "exporting"
        api_response = dict(
            dataset_uuid="1",
            name="first one",
            status="exporting",
            # Missing auth_token, hostname, or output_uri
        )
        mock_proximl._query = AsyncMock(return_value=api_response)
        with raises(SpecificationError) as exc_info:
            await dataset.connect()
        assert "missing required connection properties" in str(exc_info.value.message).lower()

    def test_dataset_billed_size_property(self, dataset, mock_proximl):
        """Test billed_size property access."""
        dataset._billed_size = 50000
        assert dataset.billed_size == 50000

    @mark.asyncio
    async def test_dataset_wait_for_archived_succeeded(
        self, dataset, mock_proximl
    ):
        mock_proximl._query = AsyncMock(
            side_effect=ApiError(404, dict(errorMessage="Dataset Not Found"))
        )
        await dataset.wait_for("archived")
        mock_proximl._query.assert_called()

    @mark.asyncio
    async def test_dataset_wait_for_unexpected_api_error(
        self, dataset, mock_proximl
    ):
        mock_proximl._query = AsyncMock(
            side_effect=ApiError(404, dict(errorMessage="Dataset Not Found"))
        )
        with raises(ApiError):
            await dataset.wait_for("ready")
        mock_proximl._query.assert_called()
