# coding: utf-8
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# -------------------------------------------------------------------------
import jwt
import pytest
from azure.messaging.webpubsubservice import WebPubSubServiceClient
from azure.core.credentials import AzureKeyCredential

try:
    from urlparse import urlparse
except ImportError:
    from urllib.parse import urlparse


def _decode_token(client, token, path="/client/hubs/hub"):
    return jwt.decode(
        token, client._config.credential.key, algorithms=["HS256"], audience=client._config.endpoint + path
    )


access_key = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789ABCDEFGH"

test_cases = [
    ("Endpoint=https://host;AccessKey={};Version=1.0;".format(access_key), "https://host"),
    ("Endpoint=http://host;AccessKey={};Version=1.0;".format(access_key), "http://host"),
    ("Endpoint=http://host;AccessKey={};Version=1.0;Port=8080;".format(access_key), "http://host:8080"),
    ("AccessKey={};Endpoint=http://host;Version=1.0;".format(access_key), "http://host"),
]


@pytest.mark.parametrize("connection_string,endpoint", test_cases)
def test_parse_connection_string(connection_string, endpoint):
    client = WebPubSubServiceClient.from_connection_string(connection_string, "hub")
    assert client._config.endpoint == endpoint
    assert isinstance(client._config.credential, AzureKeyCredential)
    assert client._config.credential.key == access_key


test_cases = [
    (None, None, None),
    ("ab", [], []),
    ("ab", ["a"], ["a"]),
    ("ab", ["a", "a", "a"], ["a", "a", "a"]),
    ("ab", ["a", "b", "c"], ["a", "b", "c"]),
    ("ab", "", ""),
]


@pytest.mark.parametrize("user_id,roles,groups", test_cases)
def test_generate_uri_contains_expected_payloads_dto(user_id, roles, groups):
    client = WebPubSubServiceClient.from_connection_string(
        "Endpoint=http://localhost;Port=8080;AccessKey={};Version=1.0;".format(access_key), "hub"
    )
    minutes_to_expire = 5
    token = client.get_client_access_token(
        user_id=user_id, roles=roles, minutes_to_expire=minutes_to_expire, groups=groups
    )
    assert token
    assert len(token) == 3
    assert set(token.keys()) == set(["baseUrl", "url", "token"])
    assert "access_token={}".format(token["token"]) == urlparse(token["url"]).query
    token = token["token"]
    decoded_token = _decode_token(client, token)
    assert decoded_token["aud"] == "{}/client/hubs/hub".format(client._config.endpoint)

    # default expire should be around 5 minutes
    assert decoded_token["exp"] - decoded_token["iat"] >= minutes_to_expire * 60 - 5
    assert decoded_token["exp"] - decoded_token["iat"] <= minutes_to_expire * 60 + 5
    if user_id:
        assert decoded_token["sub"] == user_id
    else:
        assert not decoded_token.get("sub")

    if roles:
        assert decoded_token["role"] == roles
    else:
        assert not decoded_token.get("role")

    if groups:
        assert decoded_token["webpubsub.group"] == groups
    else:
        assert not decoded_token.get("webpubsub.group")


test_cases = [
    (
        "Endpoint=http://localhost;Port=8080;AccessKey={};Version=1.0;".format(access_key),
        "hub",
        "ws://localhost:8080/client/hubs/hub",
    ),
    ("Endpoint=https://a;AccessKey={};Version=1.0;".format(access_key), "hub", "wss://a/client/hubs/hub"),
    ("Endpoint=http://a;AccessKey={};Version=1.0;".format(access_key), "hub", "ws://a/client/hubs/hub"),
]


@pytest.mark.parametrize("connection_string,hub,expected_url", test_cases)
def test_generate_url_use_same_kid_with_same_key(connection_string, hub, expected_url):
    client = WebPubSubServiceClient.from_connection_string(connection_string, hub)
    url_1 = client.get_client_access_token()["url"]
    url_2 = client.get_client_access_token()["url"]

    assert url_1.split("?")[0] == url_2.split("?")[0] == expected_url

    token_1 = urlparse(url_1).query[len("access_token=") :]
    token_2 = urlparse(url_2).query[len("access_token=") :]

    decoded_token_1 = _decode_token(client, token_1)
    decoded_token_2 = _decode_token(client, token_2)

    assert len(decoded_token_1) == len(decoded_token_2) == 3
    assert decoded_token_1["aud"] == decoded_token_2["aud"] == expected_url.replace("ws", "http")
    assert abs(decoded_token_1["iat"] - decoded_token_2["iat"]) < 5
    assert abs(decoded_token_1["exp"] - decoded_token_2["exp"]) < 5


test_cases = [
    ("Endpoint=http://localhost;Port=8080;AccessKey={};Version=1.0;".format(access_key)),
    ("Endpoint=https://a;AccessKey={};Version=1.0;".format(access_key)),
    ("Endpoint=http://a;AccessKey={};Version=1.0;".format(access_key)),
]


@pytest.mark.parametrize("connection_string", test_cases)
def test_pass_in_jwt_headers(connection_string):
    client = WebPubSubServiceClient.from_connection_string(connection_string, "hub")
    kid = "1234567890"
    token = client.get_client_access_token(jwt_headers={"kid": kid})["token"]
    assert jwt.get_unverified_header(token)["kid"] == kid


test_cases = [
    (
        "Endpoint=http://localhost;Port=8080;AccessKey={};Version=1.0;".format(access_key),
        "hub",
        "ws://localhost:8080/clients/mqtt/hubs/hub",
    ),
    ("Endpoint=https://a;AccessKey={};Version=1.0;".format(access_key), "hub", "wss://a/clients/mqtt/hubs/hub"),
    ("Endpoint=http://a;AccessKey={};Version=1.0;".format(access_key), "hub", "ws://a/clients/mqtt/hubs/hub"),
]


@pytest.mark.parametrize("connection_string,hub,expected_url", test_cases)
def test_generate_mqtt_token(connection_string, hub, expected_url):
    client = WebPubSubServiceClient.from_connection_string(connection_string, hub)
    url_1 = client.get_client_access_token(client_protocol="MQTT")["url"]

    assert url_1.split("?")[0] == expected_url

    token_1 = urlparse(url_1).query[len("access_token=") :]

    decoded_token_1 = _decode_token(client, token_1, path="/clients/mqtt/hubs/hub")

    assert len(decoded_token_1) == 3
    assert decoded_token_1["aud"] == expected_url.replace("ws", "http")


test_cases = [
    (
        "Endpoint=http://localhost;Port=8080;AccessKey={};Version=1.0;".format(access_key),
        "hub",
        "ws://localhost:8080/clients/socketio/hubs/hub",
    ),
    ("Endpoint=https://a;AccessKey={};Version=1.0;".format(access_key), "hub", "wss://a/clients/socketio/hubs/hub"),
    ("Endpoint=http://a;AccessKey={};Version=1.0;".format(access_key), "hub", "ws://a/clients/socketio/hubs/hub"),
]


@pytest.mark.parametrize("connection_string,hub,expected_url", test_cases)
def test_generate_socketio_token(connection_string, hub, expected_url):
    client = WebPubSubServiceClient.from_connection_string(connection_string, hub)
    url_1 = client.get_client_access_token(client_protocol="SocketIO")["url"]

    assert url_1.split("?")[0] == expected_url

    token_1 = urlparse(url_1).query[len("access_token=") :]

    decoded_token_1 = _decode_token(client, token_1, path="/clients/socketio/hubs/hub")

    assert len(decoded_token_1) == 3
    assert decoded_token_1["aud"] == expected_url.replace("ws", "http")