from typing import Dict, List
from unittest import TestCase, mock

import numpy as np

from runeq import Config, errors, stream


def mock_get_json_response(
    bodies: List[dict],
    calls: List,
    status_code=200,
    headers: List[Dict[str, str]] = None,
):
    """Return a function that can be used to mock .get_json_response()

    Args:
        bodies: list of JSON bodies to return from response.json()
        calls: list. Every time the function is called, args and kwargs will
            be appended to this list.
        status_code: status code to return for each response
        headers: the response header to apply
    """
    headers = headers or [{}] * len(bodies)

    num = 0

    def _func(*args, **kwargs):
        nonlocal num
        # add inputs to the list of calls that was provided
        calls.append((args, kwargs))

        resp = mock.MagicMock()
        resp.headers = headers[num]
        resp.status_code = status_code
        resp.ok = status_code < 400
        resp.json.return_value = bodies[num]
        num += 1
        return resp

    return _func


def mock_get_csv_response(
    bodies: List[str],
    calls: List,
    status_code=200,
    headers: List[Dict[str, str]] = None,
):
    """Return a function that can be used to mock .get_csv_response()

    Args:
        bodies: list of response bodies to return from response.text
        calls: list. Every time the function is called, args and kwargs will
            be appended to this list.
        status_code: status code to return for each response
        headers: the response header to apply
    """
    headers = headers or [{}] * len(bodies)

    num = 0

    def _func(*args, **kwargs):
        nonlocal num
        # keep track of the kwargs that were used to call this
        calls.append((args, kwargs))

        resp = mock.MagicMock()
        resp.headers = headers[num]
        resp.status_code = status_code
        resp.ok = status_code < 400
        resp.text = bodies[num]
        num += 1
        return resp

    return _func


class TestStreamV1Client(TestCase):
    """
    Test stream.V1Client and the associated accessors.

    """

    def setUp(self) -> None:
        """
        Initialize a client, set up basic mocking.

        """
        self.cfg = Config(
            client_key_id="abc",
            client_access_key="abc123",
        )
        self.client = stream.V1Client(self.cfg)
        self.use_np_orig = stream.v1.USE_NUMPY
        stream.v1.USE_NUMPY = False

    def tearDown(self) -> None:
        """
        Tear down monkey-patching.

        """
        stream.v1.USE_NUMPY = self.use_np_orig

    @mock.patch("runeq.stream.v1.requests")
    def test_get_json_response(self, requests):
        """
        Test the signature of JSON requests.

        """
        for test_num, case in enumerate(
            (
                (self.client.Accel, "/v1/accel.json"),
                (self.client.BandPower, "/v1/band_power.json"),
                (self.client.Event, "/v1/event.json"),
                (self.client.HeartRate, "/v1/heartrate.json"),
                (self.client.LFP, "/v1/lfp.json"),
                (self.client.ProbabilitySymptom, "/v1/probability_symptom.json"),
                (self.client.Rotation, "/v1/rotation.json"),
                (self.client.Span, "/v1/span.json"),
                (self.client.State, "/v1/state.json"),
            )
        ):
            resource_creator, endpoint = case
            resource = resource_creator(leslie="knope")
            resource.get_json_response(ron="swanson", test_num=test_num)
            requests.get.assert_has_calls(
                [
                    mock.call(
                        self.cfg.stream_url + endpoint,
                        headers=self.cfg.auth_headers,
                        params={
                            "leslie": "knope",
                            "ron": "swanson",
                            "test_num": test_num,
                        },
                    ),
                ]
            )

    @mock.patch("runeq.stream.v1.requests")
    def test_get_csv_response(self, requests):
        """
        Test the signature of CSV requests.

        """
        for test_num, case in enumerate(
            (
                (self.client.Accel, "/v1/accel.csv"),
                (self.client.BandPower, "/v1/band_power.csv"),
                (self.client.HeartRate, "/v1/heartrate.csv"),
                (self.client.LFP, "/v1/lfp.csv"),
                (self.client.ProbabilitySymptom, "/v1/probability_symptom.csv"),
                (self.client.Rotation, "/v1/rotation.csv"),
                (self.client.State, "/v1/state.csv"),
            )
        ):
            resource_creator, endpoint = case
            resource = resource_creator(leslie="knope")
            resource.get_csv_response(ron="swanson", test_num=test_num)
            requests.get.assert_has_calls(
                [
                    mock.call(
                        self.cfg.stream_url + endpoint,
                        stream=True,
                        headers=self.cfg.auth_headers,
                        params={
                            "leslie": "knope",
                            "ron": "swanson",
                            "test_num": test_num,
                        },
                    ),
                ]
            )

    def test_iter_json_data_with_token(self):
        """
        Test the iterator over JSON responses, paginating with the next page
        token header.

        """
        for test_num, resource_creator in enumerate(
            (
                self.client.Accel,
                self.client.BandPower,
                self.client.Event,
                self.client.HeartRate,
                self.client.LFP,
                self.client.ProbabilitySymptom,
                self.client.Rotation,
                self.client.Span,
                self.client.State,
            )
        ):
            resource = resource_creator()

            mock_responses = [
                {"success": True, "result": [], "next_page": 1},
                {"success": True, "result": []},
            ]

            calls = []
            resource.get_json_response = mock_get_json_response(
                mock_responses,
                calls,
                200,
                [
                    {"X-Rune-Next-Page-Token": "MTIzNDU2MDAwMA=="},
                    {},
                ],
            )

            iterator = resource.iter_json_data(test_num=test_num)
            self.assertEqual(len(list(iterator)), 2)

            # Check that all parameters were kept the same across calls,
            # except for "next_page_token"
            self.assertEqual(
                calls,
                [
                    ((), {"test_num": test_num}),
                    ((), {"test_num": test_num, "next_page_token": "MTIzNDU2MDAwMA=="}),
                ],
            )

    def test_iter_json_data(self):
        """
        Test the iterator over JSON responses, following pagination with
        the page number.

        """
        results = [{"a": 1}, {"b": 2}]
        mock_responses = [
            {"success": True, "result": results[0], "next_page": 1},
            {"success": True, "result": results[1]},
        ]

        for test_num, resource_creator in enumerate(
            (
                self.client.Accel,
                self.client.BandPower,
                self.client.Event,
                self.client.HeartRate,
                self.client.LFP,
                self.client.ProbabilitySymptom,
                self.client.Rotation,
                self.client.State,
            )
        ):
            resource = resource_creator()

            #
            # Successful Requests
            #
            calls = []
            resource.get_json_response = mock_get_json_response(mock_responses, calls)

            # Check the results
            num_results = 0
            iterator = resource.iter_json_data(test_num=test_num)
            for i, actual in enumerate(iterator):
                self.assertEqual(results[i], actual)
                num_results += 1

            self.assertEqual(num_results, 2)

            # Check that all parameters were kept the same across calls,
            # except for "page" (which must be incremented)
            self.assertEqual(
                calls,
                [((), {"test_num": test_num}), ((), {"test_num": test_num, "page": 1})],
            )

            #
            # Request Error
            # Iterator should check the response status for each request
            #
            err_details = {
                "message": "i am an intentional error!",
                "type": "TestError",
            }
            resource.get_json_response = mock_get_json_response(
                [{"success": False, "error": err_details}],
                [],
                status_code=404,
            )

            with self.assertRaises(errors.APIError) as e:
                next(resource.iter_json_data())

            err = e.exception
            self.assertEqual(err.status_code, 404)
            self.assertEqual(err.details, err_details)

    def test_iter_csv_data_with_token(self):
        """
        Test the iterator over CSV responses, which follows new pagination

        """
        mock_responses = [
            "good,better\nskiing,hiking\n",
            "good,better\ncupcakes,brownies\n",
            "",
        ]

        for test_num, resource_creator in enumerate(
            (
                self.client.Accel,
                self.client.BandPower,
                self.client.HeartRate,
                self.client.LFP,
                self.client.ProbabilitySymptom,
                self.client.Rotation,
                self.client.State,
            )
        ):
            resource = resource_creator()

            #
            # Successful Requests
            #
            calls = []
            resource.get_csv_response = mock_get_csv_response(
                mock_responses,
                calls,
                200,
                [
                    {"X-Rune-Next-Page-Token": "MTIzNDU2MDAwMA=="},
                    {"X-Rune-Next-Page-Token": "MTIzNDU2MDAwMA=="},
                    {},
                ],
            )

            # Check the results
            iterator = resource.iter_csv_text(test_num=test_num)
            self.assertEqual(len(list(iterator)), 2)

            # Check that all parameters were kept the same across calls,
            # except for "next_page_token" (which will normally be different
            # for each response)
            self.assertEqual(
                calls,
                [
                    ((), {"test_num": test_num}),
                    ((), {"test_num": test_num, "next_page_token": "MTIzNDU2MDAwMA=="}),
                    ((), {"test_num": test_num, "next_page_token": "MTIzNDU2MDAwMA=="}),
                ],
            )

    def test_iter_csv_data(self):
        """
        Test the iterator over CSV responses, which follows pagination
        """
        mock_responses = [
            "good,better\nskiing,hiking\n",
            "good,better\ncupcakes,brownies\n",
            "",
        ]

        for test_num, resource_creator in enumerate(
            (
                self.client.Accel,
                self.client.BandPower,
                self.client.HeartRate,
                self.client.LFP,
                self.client.ProbabilitySymptom,
                self.client.Rotation,
                self.client.State,
            )
        ):
            resource = resource_creator()

            #
            # Successful Requests
            #
            calls = []
            resource.get_csv_response = mock_get_csv_response(
                mock_responses,
                calls,
            )

            # Check the results
            num_results = 0
            iterator = resource.iter_csv_text(test_num=test_num)
            for i, actual in enumerate(iterator):
                self.assertEqual(mock_responses[i], actual)
                num_results += 1

            # although there are 3 responses, the last (empty) body should not
            # be returned by the iterator
            self.assertEqual(num_results, 2)

            # Check that all parameters were kept the same across calls,
            # except for "page" (which must be incremented)
            self.assertEqual(
                calls,
                [
                    ((), {"test_num": test_num}),
                    ((), {"test_num": test_num, "page": 1}),
                    ((), {"test_num": test_num, "page": 2}),
                ],
            )

            #
            # Request Error
            # Iterator should check the response status for each request
            #
            err_details = {
                "message": "i am an intentional error!",
                "type": "TestError",
            }
            # note: CSV endpoints return JSON on API errors
            resource.get_csv_response = mock_get_json_response(
                [{"success": False, "error": err_details}],
                [],
                status_code=404,
            )

            with self.assertRaises(errors.APIError) as e:
                next(resource.iter_csv_text())

            err = e.exception
            self.assertEqual(err.status_code, 404)
            self.assertEqual(err.details, err_details)

    def test_iter_points(self):
        """
        Test iterating over data as points. Uses the CSV endpoint.

        """
        mock_responses = [
            "lower,higher,label\n1,2,ints\n3.5,6.7,floats\n",
            "lower,higher,label\n,8.9,missing data\n",
            "",
        ]
        expected = [
            {"lower": 1, "higher": 2, "label": "ints"},
            {"lower": 3.5, "higher": 6.7, "label": "floats"},
            {"lower": None, "higher": 8.9, "label": "missing data"},
        ]

        for test_num, resource_creator in enumerate(
            (
                self.client.Accel,
                self.client.BandPower,
                self.client.HeartRate,
                self.client.LFP,
                self.client.ProbabilitySymptom,
                self.client.Rotation,
                self.client.State,
            )
        ):
            resource = resource_creator()
            # replace get_csv_response on the resource
            resource.get_csv_response = mock_get_csv_response(
                mock_responses,
                [],
            )

            for i, point in enumerate(resource.points()):
                self.assertDictEqual(expected[i], point)

            # replace get_csv_response again, to restart the mock responses
            resource.get_csv_response = mock_get_csv_response(
                mock_responses,
                [],
            )

            for i, point in enumerate(resource):
                self.assertDictEqual(expected[i], point)
                # check dtype for "higher", which always has a numeric
                # value in the test data
                self.assertNotIsInstance(point["higher"], np.float64)

    def test_iter_points_numpy(self):
        """
        Test iterating over data as points, using Numpy to convert.

        """
        stream.v1.USE_NUMPY = True

        mock_responses = [
            "lower,higher,label\n1,2,ints\n3.5,6.7,floats\n",
            "lower,higher,label\n,8.9,missing data\n",
            "",
        ]
        expected = [
            {"lower": 1, "higher": 2, "label": "ints"},
            {"lower": 3.5, "higher": 6.7, "label": "floats"},
            {"lower": np.NaN, "higher": 8.9, "label": "missing data"},
        ]

        for test_num, resource_creator in enumerate(
            (
                self.client.Accel,
                self.client.BandPower,
                self.client.HeartRate,
                self.client.LFP,
                self.client.ProbabilitySymptom,
                self.client.Rotation,
                self.client.State,
            )
        ):
            resource = resource_creator()
            # replace get_csv_response on the resource
            resource.get_csv_response = mock_get_csv_response(
                mock_responses,
                [],
            )

            for i, point in enumerate(resource.points()):
                self.assertDictEqual(expected[i], point)
                # check dtype for "higher", which always has a numeric
                # value in the test data
                self.assertIsInstance(point["higher"], np.float64)
