"""Test basic handler functions."""

import copy
import unittest

import numpy as np
from webtest import AppError
from webtest import TestApp as App

from pydap.exceptions import ConstraintExpressionError
from pydap.handlers.lib import (
    BaseHandler,
    ConstraintExpression,
    IterData,
    apply_projection,
    apply_selection,
)
from pydap.lib import walk
from pydap.model import BaseType, SequenceType, StructureType
from pydap.parsers import parse_projection
from pydap.tests.datasets import (
    NestedSequence,
    SimpleArray,
    SimpleGrid,
    SimpleSequence,
    SimpleStructure,
    VerySimpleSequence,
)


class TestBaseHandler(unittest.TestCase):
    """Test the base handler as a WSGI app."""

    def setUp(self):
        """Create a basic WSGI app."""
        self.app = App(MockHandler(SimpleArray))

    def test_unconstrained_das(self):
        """DAS responses are always unconstrained."""
        res = self.app.get("/.dds")
        self.assertEqual(
            res.text,
            """Dataset {
    Byte byte[byte = 5];
    String string[string = 2];
    Int16 short;
} SimpleArray;
""",
        )

        res = self.app.get("/.dds?byte")
        self.assertEqual(
            res.text,
            """Dataset {
    Byte byte[byte = 5];
} SimpleArray;
""",
        )

        res = self.app.get("/.das")
        das = res.text
        self.assertEqual(
            das,
            """Attributes {
    byte {
    }
    string {
    }
    short {
    }
}
""",
        )

        # check that DAS is unmodifed with constraint expression
        res = self.app.get("/.das?byte")
        self.assertEqual(res.text, das)

    def test_exception(self):
        """
        Test exception handling.

        By default pydap will capture all exceptions and return a formatted
        error response.

        """
        with self.assertRaises(AppError):
            self.app.get("/.foo")

    def test_exception_non_captured(self):
        """Test exception handling when not captured."""
        app = App(
            MockHandler(SimpleArray), extra_environ={"x-wsgiorg.throw_errors": True}
        )
        with self.assertRaises(KeyError):
            app.get("/.foo")

    def test_missing_dataset(self):
        """Test exception when dataset is not set."""
        app = App(MockHandler(), extra_environ={"x-wsgiorg.throw_errors": True})
        with self.assertRaises(NotImplementedError):
            app.get("/.dds")


class TestApplySelection(unittest.TestCase):
    """Test function that applies selections to the dataset."""

    def setUp(self):
        """Build our own sequence.

        pydap uses lightweight copies of objects that share data. This breaks
        unit tests since the same objects are reused for tests.

        """
        # make a dataset with its own data
        self.dataset = copy.copy(SimpleSequence)
        self.dataset.cast.data = SimpleSequence.cast.data.copy()

    def test_no_selection(self):
        """Test no selection in the query string."""
        dataset = apply_selection("", self.dataset)
        np.testing.assert_array_equal(dataset.cast.data, self.dataset.cast.data)

    def test_simple_selection(self):
        """Test a simple selection applied to the dataset."""
        dataset = apply_selection(["cast.lon=100"], self.dataset)
        np.testing.assert_array_equal(
            dataset.cast.data,
            self.dataset.cast.data[self.dataset.cast.data["lon"] == 100],
        )

    def test_multiple_selections(self):
        """Test multiple selections applied to dataset."""
        dataset = apply_selection(["cast.lon=100", "cast.lat>0"], self.dataset)
        np.testing.assert_array_equal(
            dataset.cast.data,
            self.dataset.cast.data[self.dataset.cast.data["lon"] == 100][
                self.dataset.cast.data["lat"] > 0
            ],
        )


class TestApplyProjectionGrid(unittest.TestCase):
    """Test applying projections on a dataset with a grid."""

    def setUp(self):
        """Build dataset with no shared data."""
        self.dataset = copy.copy(SimpleGrid)
        for var in walk(self.dataset, BaseType):
            var.data = var.data.copy()

    def test_no_projection(self):
        """Test no projections."""
        dataset = apply_projection("", self.dataset)
        self.assertEqual(list(dataset.children()), [])

    def test_simple_projection(self):
        """Test simple projections."""
        dataset = apply_projection(parse_projection("x"), self.dataset)
        self.assertEqual(list(dataset.keys()), ["x"])

    def test_simple_projection_with_index(self):
        """Test simple projections."""
        dataset = apply_projection(parse_projection("x[1]"), self.dataset)
        np.testing.assert_array_equal(dataset.x.data, [1])

    def test_array(self):
        """Test that the grid degenerates into a structure."""
        dataset = apply_projection(
            parse_projection("SimpleGrid.SimpleGrid"), self.dataset
        )
        self.assertIsInstance(dataset.SimpleGrid, StructureType)

    def test_array_slice(self):
        """Test slices applied to a grid."""
        dataset = apply_projection(parse_projection("SimpleGrid[1]"), self.dataset)
        np.testing.assert_array_equal(
            dataset.SimpleGrid.x.data, self.dataset.SimpleGrid[1].x.data
        )
        np.testing.assert_array_equal(
            dataset.SimpleGrid.y.data, self.dataset.SimpleGrid[1].y.data
        )
        np.testing.assert_array_equal(
            dataset.SimpleGrid.SimpleGrid.data,
            self.dataset.SimpleGrid[1:2].SimpleGrid.data,
        )


class TestApplyProjectionSequence(unittest.TestCase):
    """Test applying projections on a dataset with a sequence."""

    def setUp(self):
        """Build dataset with no shared data."""
        self.dataset = copy.copy(VerySimpleSequence)
        self.dataset.sequence.data = VerySimpleSequence.sequence.data.copy()

    def test_sequence_projection(self):
        """Test projection slicing on sequences."""
        dataset = apply_projection(parse_projection("sequence[2]"), self.dataset)
        np.testing.assert_array_equal(
            dataset.sequence.data, VerySimpleSequence.sequence.data[2]
        )


class TestInvalidProjection(unittest.TestCase):
    """Test applying a projection to a structure object."""

    def test_structure_projection(self):
        """Test projection slicing on a structure."""
        with self.assertRaises(ConstraintExpressionError):
            apply_projection(parse_projection("types[0]"), SimpleStructure)


class TestConstraintExpression(unittest.TestCase):
    """Test the constraint expression object."""

    def test_str(self):
        """Test string representation."""
        ce = ConstraintExpression("a>1")
        self.assertEqual(str(ce), "a>1")

    def test_unicode(self):
        """Test unicode representation."""
        ce = ConstraintExpression("a>1")
        self.assertEqual(str(ce), "a>1")

    def test_and(self):
        """Test CE addition."""
        ce1 = ConstraintExpression("a>1")
        ce2 = ConstraintExpression("b>0")
        ce3 = ce1 & ce2
        self.assertEqual(str(ce3), "a>1&b>0")

    def test_or(self):
        """Expressions cannot be ORed."""
        ce1 = ConstraintExpression("a>1")
        ce2 = ConstraintExpression("b>0")
        with self.assertRaises(ConstraintExpressionError):
            ce1 | ce2


class TestIterData(unittest.TestCase):
    """
    Test the ``IterData`` class, used to store flat/nested sequence data.

    A flat ``IterData`` should behave like a Numpy structured array, except
    all operations are stored to be lazily evaluated when the object is
    iterated over.

    """

    def setUp(self):
        """Create a flat IterData."""
        template = SequenceType("a")
        template["b"] = BaseType("b")
        template["c"] = BaseType("c")
        template["d"] = BaseType("d")
        self.data = IterData([(1, 2, 3), (4, 5, 6)], template)

        self.array = np.array(
            np.rec.fromrecords(
                [
                    (1, 2, 3),
                    (4, 5, 6),
                ],
                names=["b", "c", "d"],
            )
        )

    def assertIteratorEqual(self, it1, it2):
        self.assertEqual(list(it1), list(it2))

    def test_repr(self):
        """Test the object representation."""
        self.assertEqual(repr(self.data), "<IterData to stream [(1, 2, 3), (4, 5, 6)]>")

    def test_dtype(self):
        """Test the ``dtype`` property."""
        self.assertEqual(self.data["b"].dtype, self.array["b"].dtype)

    def test_iteration(self):
        """Test iteration over data."""
        self.assertIteratorEqual(map(tuple, self.data), map(tuple, self.array))
        self.assertIteratorEqual(self.data["b"], self.array["b"])

    def test_filter(self):
        """Test filtering the object."""
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] == 1]),
            map(tuple, self.array[self.array["b"] == 1]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] != 1]),
            map(tuple, self.array[self.array["b"] != 1]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] >= 1]),
            map(tuple, self.array[self.array["b"] >= 1]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] <= 1]),
            map(tuple, self.array[self.array["b"] <= 1]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] > 1]),
            map(tuple, self.array[self.array["b"] > 1]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] < 1]),
            map(tuple, self.array[self.array["b"] < 1]),
        )

    def test_slice(self):
        """Test slicing the object."""
        self.assertIteratorEqual(map(tuple, self.data[1:]), map(tuple, self.array[1:]))

    def test_integer_slice(self):
        """Test slicing with an integer.

        Note that the behavior here is different from Numpy arrays, since the
        data access to ``IterData`` is through iteration it has no direct index
        access.

        """
        self.assertIteratorEqual(self.data[0:1], self.array[0:1].tolist())
        self.assertIteratorEqual(self.data[0], self.array[0:1].tolist())

    def test_invalid_child(self):
        """Test accessing a non-existing child."""
        with self.assertRaises(KeyError):
            self.data["e"]

    def test_invalid_key(self):
        """Test accessing using an invalid key."""
        with self.assertRaises(KeyError):
            self.data[(1, 2)]

    def test_selecting_children(self):
        """Test that we can select children."""
        self.assertIteratorEqual(
            map(tuple, self.data[["d", "b"]]), map(tuple, self.array[["d", "b"]])
        )

    def test_invalid_selection(self):
        """Test invalid selections.

        In theory this should never happen, since ``ConstraintExpression``
        object are constructly directly from existing children.

        """
        with self.assertRaises(ConstraintExpressionError):
            self.data[ConstraintExpression("a.e<1")]
        with self.assertRaises(ConstraintExpressionError):
            self.data[ConstraintExpression("a.d<foo")]

    def test_intercomparison_selection(self):
        """Test comparing children in the selection."""
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] == self.data["c"]]),
            map(tuple, self.array[self.array["b"] == self.array["c"]]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] != self.data["c"]]),
            map(tuple, self.array[self.array["b"] != self.array["c"]]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] >= self.data["c"]]),
            map(tuple, self.array[self.array["b"] >= self.array["c"]]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] <= self.data["c"]]),
            map(tuple, self.array[self.array["b"] <= self.array["c"]]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] > self.data["c"]]),
            map(tuple, self.array[self.array["b"] > self.array["c"]]),
        )
        self.assertIteratorEqual(
            map(tuple, self.data[self.data["b"] < self.data["c"]]),
            map(tuple, self.array[self.array["b"] < self.array["c"]]),
        )

    def test_selection_not_in_projection(self):
        """Test selection with variables that are not in the projection."""
        self.data[["d", "b"]]
        filtered = self.data[["d", "b"]][self.data["c"] > 3]
        self.assertEqual(filtered, self.array[["d", "b"]][self.array["c"] > 3])


class TestRegexp(unittest.TestCase):
    """Test regular expression match."""

    def test_regexp(self):
        sequence = SequenceType("sequence")
        sequence["name"] = BaseType("name")
        sequence.data = IterData(
            [
                ("John", "Paul", "George", "Ringo"),
            ],
            sequence,
        )

        filtered = sequence[ConstraintExpression('sequence.name=~"J.*"')]
        self.assertEqual(list(filtered.iterdata()), [("John",)])


class TestNestedIterData(unittest.TestCase):
    """Test ``IterData`` with nested data."""

    def setUp(self):
        """Load data from test dataset."""
        self.data = NestedSequence.location.data

    def test_iteration(self):
        """Test basic iteration."""
        self.assertEqual(
            list(self.data),
            [
                (1, 1, 1, [(10, 11, 12), (21, 22, 23)]),
                (2, 4, 4, [(15, 16, 17)]),
                (3, 6, 9, []),
                (4, 8, 16, [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)]),
            ],
        )

    def test_children_data(self):
        """Test getting data from a simple child."""
        self.assertEqual(list(self.data["lat"]), [1, 2, 3, 4])

    def test_sequence_children_data(self):
        """Test getting data from a sequence child."""
        self.assertEqual(
            list(self.data["time_series"]),
            [
                [(10, 11, 12), (21, 22, 23)],
                [(15, 16, 17)],
                [],
                [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)],
            ],
        )

    def test_deep_children_data(self):
        """Test getting data from a sequence child."""
        self.assertEqual(
            list(self.data["time_series"]["time"]),
            [[10, 21], [15], [], [31, 41, 51, 61]],
        )

    def test_selecting_children(self):
        """Test that we can select children."""
        self.assertEqual(
            list(self.data[["time_series", "elev"]]),
            [
                ([(10, 11, 12), (21, 22, 23)], 1),
                ([(15, 16, 17)], 4),
                ([], 9),
                ([(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)], 16),
            ],
        )

    def test_slice(self):
        """Test slicing the object."""
        self.assertEqual(
            list(self.data[1::2]),
            [
                (2, 4, 4, [(15, 16, 17)]),
                (4, 8, 16, [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)]),
            ],
        )

    def test_children_data_from_slice(self):
        """Test getting children data from a sliced sequence."""
        self.assertEqual(list(self.data[1::2]["lat"]), [2, 4])

    def test_sequence_children_data_from_slice(self):
        """Test getting children data from a sliced sequence."""
        self.assertEqual(
            list(self.data[1::2]["time_series"]),
            [[(15, 16, 17)], [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)]],
        )

    def test_deep_slice(self):
        """Test slicing the inner sequence."""
        self.assertEqual(
            list(self.data["time_series"][::2]), [[(10, 11, 12), (21, 22, 23)], []]
        )

    def test_integer_slice(self):
        """Test slicing with an integer."""
        self.assertEqual(list(self.data["time_series"][1]), [[(15, 16, 17)]])

    def test_filter_data(self):
        """Test filtering the data."""
        self.assertEqual(
            list(self.data[self.data["lat"] > 2]),
            [
                (3, 6, 9, []),
                (4, 8, 16, [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)]),
            ],
        )

    def test_deep_filter(self):
        """Test deep filtering the data."""
        self.assertEqual(
            list(self.data[self.data["time_series"]["slp"] > 11]),
            [
                (1, 1, 1, [(21, 22, 23)]),
                (2, 4, 4, [(15, 16, 17)]),
                (3, 6, 9, []),
                (4, 8, 16, [(31, 32, 33), (41, 42, 43), (51, 52, 53), (61, 62, 63)]),
            ],
        )


class MockWorkingSet(object):
    """A fake working set for testing handlers."""

    def iter_entry_points(self, group):
        """Return a mock entry point."""
        yield MockEntryPoint(MockHandler)


class MockHandler(BaseHandler):
    """A fake handler for testing."""

    extensions = r"^.*\.foo$"

    def __init__(self, dataset=None):
        BaseHandler.__init__(self, dataset)
        self.additional_headers = [("X-debug", "True")]


class MockEntryPoint(object):
    """A fake entry point for testing."""

    def __init__(self, handler):
        self.handler = handler
        self.name = "test"
        self.module_name = "pydap.handlers.test"
        self.attrs = ("TestHandler",)

    def load(self):
        """Return the wrapped handler."""
        return self.handler
