from io import BytesIO
import json
import urllib.parse
from collections import OrderedDict

from PIL import Image
import numpy as np

import pytest


@pytest.fixture(scope='module')
def flask_app():
    from terracotta.server import create_app
    return create_app()


@pytest.fixture(scope='module')
def client(flask_app):
    with flask_app.test_client() as client:
        yield client


def test_get_keys(client, use_testdb):
    rv = client.get('/keys')

    expected_response = [
        {'key': 'key1'},
        {'key': 'akey'},
        {'key': 'key2', 'description': 'key2'}
    ]
    assert rv.status_code == 200
    assert expected_response == json.loads(rv.data)['keys']


def test_get_metadata(client, use_testdb):
    rv = client.get('/metadata/val11/x/val12/')
    assert rv.status_code == 200
    assert ['extra_data'] == json.loads(rv.data)['metadata']


def test_get_metadata_nonexisting(client, use_testdb):
    rv = client.get('/metadata/val11/x/NONEXISTING/')
    assert rv.status_code == 404


def test_get_datasets(client, use_testdb):
    rv = client.get('/datasets')
    assert rv.status_code == 200
    datasets = json.loads(rv.data, object_pairs_hook=OrderedDict)['datasets']
    assert len(datasets) == 4
    assert OrderedDict([('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')]) in datasets


def test_get_datasets_pagination(client, use_testdb):
    # no page (implicit 0)
    rv = client.get('/datasets?limit=2')
    assert rv.status_code == 200
    response = json.loads(rv.data, object_pairs_hook=OrderedDict)
    assert response['limit'] == 2
    assert response['page'] == 0

    first_datasets = response['datasets']
    assert len(first_datasets) == 2
    assert OrderedDict([('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')]) in first_datasets

    # second page
    rv = client.get('/datasets?limit=2&page=1')
    assert rv.status_code == 200
    response = json.loads(rv.data, object_pairs_hook=OrderedDict)
    assert response['limit'] == 2
    assert response['page'] == 1

    last_datasets = response['datasets']
    assert len(last_datasets) == 2
    assert OrderedDict([('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')]) not in last_datasets

    # page out of range
    rv = client.get('/datasets?limit=2&page=1000')
    assert rv.status_code == 200
    assert not json.loads(rv.data)['datasets']

    # invalid page
    rv = client.get('/datasets?page=-1')
    assert rv.status_code == 400

    # invalid limit
    rv = client.get('/datasets?limit=-1')
    assert rv.status_code == 400


def test_get_datasets_selective(client, use_testdb):
    rv = client.get('/datasets?key1=val21')
    assert rv.status_code == 200
    assert len(json.loads(rv.data)['datasets']) == 3

    rv = client.get('/datasets?key1=val21&key2=val23')
    assert rv.status_code == 200
    assert len(json.loads(rv.data)['datasets']) == 1


def test_get_datasets_unknown_key(client, use_testdb):
    rv = client.get('/datasets?UNKNOWN=val21')
    assert rv.status_code == 400


def test_get_singleband_greyscale(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def test_get_singleband_extra_args(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?foo=bar&baz=quz')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def test_get_singleband_cmap(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=jet')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def test_get_singleband_preview(client, use_testdb):
    import terracotta
    settings = terracotta.get_settings()

    rv = client.get(f'/singleband/val11/x/val12/preview.png?colormap=jet')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def urlsafe_json(payload):
    payload_json = json.dumps(payload)
    return urllib.parse.quote_plus(payload_json, safe=r',.[]{}:"')


def test_get_singleband_explicit_cmap(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz
    explicit_cmap = {1: (0, 0, 0), 2.0: (255, 255, 255, 20), 3: '#ffffff', 4: 'abcabc'}

    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=explicit'
                    f'&explicit_color_map={urlsafe_json(explicit_cmap)}')
    assert rv.status_code == 200, rv.data.decode('utf-8')

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def test_get_singleband_explicit_cmap_invalid(client, use_testdb, raster_file_xyz):
    x, y, z = raster_file_xyz
    explicit_cmap = {1: (0, 0, 0), 2: (255, 255, 255), 3: '#ffffff', 4: 'abcabc'}

    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?'
                    f'explicit_color_map={urlsafe_json(explicit_cmap)}')
    assert rv.status_code == 400

    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=jet'
                    f'&explicit_color_map={urlsafe_json(explicit_cmap)}')
    assert rv.status_code == 400

    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=explicit')
    assert rv.status_code == 400

    explicit_cmap[3] = 'omgomg'
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=explicit'
                    f'&explicit_color_map={urlsafe_json(explicit_cmap)}')
    assert rv.status_code == 400

    explicit_cmap = [(255, 255, 255)]
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=explicit'
                    f'&explicit_color_map={urlsafe_json(explicit_cmap)}')
    assert rv.status_code == 400

    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=explicit'
                    f'&explicit_color_map=foo')
    assert rv.status_code == 400


def test_get_singleband_stretch(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz

    for stretch_range in ('[0,1]', '[0,null]', '[null, 1]', '[null,null]', 'null'):
        rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?stretch_range={stretch_range}')
        assert rv.status_code == 200

        img = Image.open(BytesIO(rv.data))
        assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def test_get_singleband_out_of_bounds(client, use_testdb):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = (0, 0, 10)
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE
    assert np.all(np.asarray(img) == 0)


def test_get_singleband_unknown_cmap(client, use_testdb, raster_file_xyz):
    x, y, z = raster_file_xyz
    rv = client.get(f'/singleband/val11/x/val12/{z}/{x}/{y}.png?colormap=UNKNOWN')
    assert rv.status_code == 400


def test_get_rgb(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz
    rv = client.get(f'/rgb/val21/x/{z}/{x}/{y}.png?r=val22&g=val23&b=val24')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == (*settings.DEFAULT_TILE_SIZE, 3)


def test_get_rgb_preview(client, use_testdb):
    import terracotta
    settings = terracotta.get_settings()

    rv = client.get(f'/rgb/val21/x/preview.png?r=val22&g=val23&b=val24')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == (*settings.DEFAULT_TILE_SIZE, 3)


def test_get_rgb_extra_args(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz
    rv = client.get(f'/rgb/val21/x/{z}/{x}/{y}.png?r=val22&g=val23&b=val24&foo=bar&baz=quz')
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == (*settings.DEFAULT_TILE_SIZE, 3)


def test_get_rgb_stretch(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    x, y, z = raster_file_xyz

    for stretch_range in ('[0,10000]', '[0,null]', '[null, 10000]', '[null,null]', 'null'):
        rv = client.get(f'/rgb/val21/x/{z}/{x}/{y}.png?r=val22&g=val23&b=val24&'
                        f'r_range={stretch_range}&b_range={stretch_range}&g_range={stretch_range}')
        assert rv.status_code == 200, rv.data

        img = Image.open(BytesIO(rv.data))
        assert np.asarray(img).shape == (*settings.DEFAULT_TILE_SIZE, 3)


def test_get_compute(client, use_testdb, raster_file_xyz):
    import terracotta
    settings = terracotta.get_settings()

    # default tile size
    x, y, z = raster_file_xyz
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?expression=v1*v2&v1=val22&v2=val23'
        '&stretch_range=[0,10000]'
    )
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE

    # custom tile size
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?expression=v1*v2&v1=val22&v2=val23'
        '&stretch_range=[0,10000]'
        '&tile_size=[128,128]'
    )
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == (128, 128)


def test_get_compute_preview(client, use_testdb):
    import terracotta
    settings = terracotta.get_settings()

    rv = client.get(
        f'/compute/val21/x/preview.png'
        '?expression=v1*v2&v1=val22&v2=val23'
        '&stretch_range=[0,10000]'
    )
    assert rv.status_code == 200

    img = Image.open(BytesIO(rv.data))
    assert np.asarray(img).shape == settings.DEFAULT_TILE_SIZE


def test_get_compute_invalid(client, use_testdb, raster_file_xyz):
    x, y, z = raster_file_xyz

    # too few keys
    rv = client.get(
        f'/compute/val21/{z}/{x}/{y}.png'
        '?expression=v1*v2&v1=val22&v2=val23'
        '&stretch_range=[0,10000]'
    )
    assert rv.status_code == 400

    # invalid expression
    rv = client.get(
        f'/compute/val21/x/preview.png'
        '?expression=__builtins__["dir"](v1)&v1=val22'
        '&stretch_range=[0,10000]'
    )
    assert rv.status_code == 400

    # no stretch range
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?expression=v1*v2&v1=val22&v2=val23'
    )
    assert rv.status_code == 400

    # no expression
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?stretch_range=[0,10000)'
    )
    assert rv.status_code == 400

    # missing operand
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?expression=v1*v2'
        '&stretch_range=[0,10000)'
    )
    assert rv.status_code == 400

    # invalid stretch range (syntax)
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?expression=v1*v2&v1=val22&v2=val23'
        '&stretch_range=[0,10000)'
    )
    assert rv.status_code == 400

    # invalid stretch range (value)
    rv = client.get(
        f'/compute/val21/x/{z}/{x}/{y}.png'
        '?expression=v1*v2&v1=val22&v2=val23'
        '&stretch_range=[10000,0]'
    )
    assert rv.status_code == 400


def test_get_colormap(client):
    rv = client.get('/colormap?stretch_range=[0,1]&num_values=100')
    assert rv.status_code == 200
    assert len(json.loads(rv.data)['colormap']) == 100


def test_get_colormap_invalid(client):
    rv = client.get('/colormap?stretch_range=[0,1')
    assert rv.status_code == 400


def test_get_colormap_extra_args(client):
    rv = client.get('/colormap?stretch_range=[0,1]&num_values=100&foo=bar&baz=quz')
    assert rv.status_code == 200
    assert len(json.loads(rv.data)['colormap']) == 100


def test_get_spec(client):
    from terracotta import __version__

    rv = client.get('/swagger.json')
    assert rv.status_code == 200
    assert json.loads(rv.data)
    assert __version__ in rv.data.decode('utf-8')

    rv = client.get('/apidoc')
    assert rv.status_code == 200
    assert b'Terracotta' in rv.data

[((151, 181), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""module"""'}), "(scope='module')\n", (165, 181), False, 'import pytest\n'), ((271, 301), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""module"""'}), "(scope='module')\n", (285, 301), False, 'import pytest\n'), ((255, 267), 'terracotta.server.create_app', 'create_app', ([], {}), '()\n', (265, 267), False, 'from terracotta.server import create_app\n'), ((1495, 1545), 'json.loads', 'json.loads', (['rv.data'], {'object_pairs_hook': 'OrderedDict'}), '(rv.data, object_pairs_hook=OrderedDict)\n', (1505, 1545), False, 'import json\n'), ((1903, 1953), 'json.loads', 'json.loads', (['rv.data'], {'object_pairs_hook': 'OrderedDict'}), '(rv.data, object_pairs_hook=OrderedDict)\n', (1913, 1953), False, 'import json\n'), ((3118, 3143), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (3141, 3143), False, 'import terracotta\n'), ((3489, 3514), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (3512, 3514), False, 'import terracotta\n'), ((3870, 3895), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (3893, 3895), False, 'import terracotta\n'), ((4234, 4259), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (4257, 4259), False, 'import terracotta\n'), ((4520, 4539), 'json.dumps', 'json.dumps', (['payload'], {}), '(payload)\n', (4530, 4539), False, 'import json\n'), ((4722, 4747), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (4745, 4747), False, 'import terracotta\n'), ((6587, 6612), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (6610, 6612), False, 'import terracotta\n'), ((7075, 7100), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (7098, 7100), False, 'import terracotta\n'), ((7686, 7711), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (7709, 7711), False, 'import terracotta\n'), ((8047, 8072), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (8070, 8072), False, 'import terracotta\n'), ((8394, 8419), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (8417, 8419), False, 'import terracotta\n'), ((8788, 8813), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (8811, 8813), False, 'import terracotta\n'), ((9380, 9405), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (9403, 9405), False, 'import terracotta\n'), ((10165, 10190), 'terracotta.get_settings', 'terracotta.get_settings', ([], {}), '()\n', (10188, 10190), False, 'import terracotta\n'), ((12644, 12663), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (12654, 12663), False, 'import json\n'), ((1140, 1190), 'json.loads', 'json.loads', (['rv.data'], {'object_pairs_hook': 'OrderedDict'}), '(rv.data, object_pairs_hook=OrderedDict)\n', (1150, 1190), False, 'import json\n'), ((1244, 1310), 'collections.OrderedDict', 'OrderedDict', (["[('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')]"], {}), "([('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')])\n", (1255, 1310), False, 'from collections import OrderedDict\n'), ((1703, 1769), 'collections.OrderedDict', 'OrderedDict', (["[('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')]"], {}), "([('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')])\n", (1714, 1769), False, 'from collections import OrderedDict\n'), ((2109, 2175), 'collections.OrderedDict', 'OrderedDict', (["[('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')]"], {}), "([('key1', 'val11'), ('akey', 'x'), ('key2', 'val12')])\n", (2120, 2175), False, 'from collections import OrderedDict\n'), ((3296, 3312), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (3303, 3312), False, 'from io import BytesIO\n'), ((3683, 3699), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (3690, 3699), False, 'from io import BytesIO\n'), ((4061, 4077), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (4068, 4077), False, 'from io import BytesIO\n'), ((4391, 4407), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (4398, 4407), False, 'from io import BytesIO\n'), ((5105, 5121), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (5112, 5121), False, 'from io import BytesIO\n'), ((7248, 7264), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (7255, 7264), False, 'from io import BytesIO\n'), ((7875, 7891), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (7882, 7891), False, 'from io import BytesIO\n'), ((8202, 8218), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (8209, 8218), False, 'from io import BytesIO\n'), ((8599, 8615), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (8606, 8615), False, 'from io import BytesIO\n'), ((9668, 9684), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (9675, 9684), False, 'from io import BytesIO\n'), ((10011, 10027), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (10018, 10027), False, 'from io import BytesIO\n'), ((10395, 10411), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (10402, 10411), False, 'from io import BytesIO\n'), ((655, 674), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (665, 674), False, 'import json\n'), ((838, 857), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (848, 857), False, 'import json\n'), ((2321, 2340), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (2331, 2340), False, 'import json\n'), ((3325, 3340), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (3335, 3340), True, 'import numpy as np\n'), ((3712, 3727), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (3722, 3727), True, 'import numpy as np\n'), ((4090, 4105), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (4100, 4105), True, 'import numpy as np\n'), ((4420, 4435), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (4430, 4435), True, 'import numpy as np\n'), ((5134, 5149), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (5144, 5149), True, 'import numpy as np\n'), ((6892, 6908), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (6899, 6908), False, 'from io import BytesIO\n'), ((7277, 7292), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (7287, 7292), True, 'import numpy as np\n'), ((7347, 7362), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (7357, 7362), True, 'import numpy as np\n'), ((7904, 7919), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (7914, 7919), True, 'import numpy as np\n'), ((8231, 8246), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (8241, 8246), True, 'import numpy as np\n'), ((8628, 8643), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (8638, 8643), True, 'import numpy as np\n'), ((9191, 9207), 'io.BytesIO', 'BytesIO', (['rv.data'], {}), '(rv.data)\n', (9198, 9207), False, 'from io import BytesIO\n'), ((9697, 9712), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (9707, 9712), True, 'import numpy as np\n'), ((10040, 10055), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (10050, 10055), True, 'import numpy as np\n'), ((10424, 10439), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (10434, 10439), True, 'import numpy as np\n'), ((2690, 2709), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (2700, 2709), False, 'import json\n'), ((2832, 2851), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (2842, 2851), False, 'import json\n'), ((6925, 6940), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (6935, 6940), True, 'import numpy as np\n'), ((9224, 9239), 'numpy.asarray', 'np.asarray', (['img'], {}), '(img)\n', (9234, 9239), True, 'import numpy as np\n'), ((12112, 12131), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (12122, 12131), False, 'import json\n'), ((12454, 12473), 'json.loads', 'json.loads', (['rv.data'], {}), '(rv.data)\n', (12464, 12473), False, 'import json\n')]