# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import base64
import json
import time
from urllib.parse import urlparse
from unittest import mock

from azure.core.credentials import AccessToken, AccessTokenInfo


FAKE_CLIENT_ID = "fake-client-id"
INVALID_CHARACTERS = "|\\`;{&' "
INVALID_SUBSCRIPTION_CHARACTERS = "|\\`;{&'"
ACCESS_TOKEN_CLASSES = (AccessToken, AccessTokenInfo)
GET_TOKEN_METHODS = ("get_token", "get_token_info")


def build_id_token(
    iss="issuer",
    sub="subject",
    aud="client-id",
    username="username",
    tenant_id="tenant-id",
    object_id="object-id",
    **claims
):
    token_claims = id_token_claims(
        iss=iss, sub=sub, aud=aud, tid=tenant_id, oid=object_id, preferred_username=username, **claims
    )
    jwt_payload = base64.b64encode(json.dumps(token_claims).encode()).decode("utf-8")
    return "header.{}.signature".format(jwt_payload)


def build_adfs_id_token(iss="issuer", sub="subject", aud="client-id", username="username", **claims):
    token_claims = id_token_claims(iss=iss, sub=sub, aud=aud, upn=username, **claims)
    jwt_payload = base64.b64encode(json.dumps(token_claims).encode()).decode("utf-8")
    return "header.{}.signature".format(jwt_payload)


def id_token_claims(iss, sub, aud, exp=None, iat=None, **claims):
    return dict(
        {"iss": iss, "sub": sub, "aud": aud, "exp": exp or int(time.time()) + 3600, "iat": iat or int(time.time())},
        **claims
    )


def build_aad_response(  # simulate a response from Microsoft Entra ID
    uid=None,
    utid=None,  # If present, they will form client_info
    access_token=None,
    expires_in=3600,
    token_type="Bearer",
    refresh_token=None,
    foci=None,
    id_token=None,  # or something generated by build_id_token()
    error=None,
    **kwargs
):
    response = {}
    if uid and utid:  # Mimic the Microsoft Entra ID behavior for "client_info=1" request
        response["client_info"] = base64.b64encode(json.dumps({"uid": uid, "utid": utid}).encode()).decode("utf-8")
    if error:
        response["error"] = error
    if access_token:
        response.update({"access_token": access_token, "expires_in": expires_in, "token_type": token_type})
    if refresh_token:
        response["refresh_token"] = refresh_token
    if id_token:
        response["id_token"] = id_token
    if foci:
        response["foci"] = foci
    response.update(kwargs)
    return response


class Request:
    def __init__(
        self,
        base_url=None,
        url=None,
        authority=None,
        url_substring=None,
        method=None,
        required_headers=None,
        required_data=None,
        required_params=None,
    ):
        self.authority = authority
        self.base_url = base_url
        self.method = method
        self.url = url
        self.url_substring = url_substring
        self.required_headers = required_headers or {}
        self.required_data = required_data or {}
        self.required_params = required_params or {}

    def assert_matches(self, request):
        discrepancies = []

        def add_discrepancy(name, expected, actual):
            discrepancies.append("{}:\n\t expected: {}\n\t   actual: {}".format(name, expected, actual))

        if self.base_url and not request.url.startswith(self.base_url):
            add_discrepancy("base url", self.base_url, request.url)

        if self.url and self.url != request.url:
            add_discrepancy("url", self.url, request.url)

        if self.url_substring and self.url_substring not in request.url:
            add_discrepancy("url substring", self.url_substring, request.url)

        parsed = urlparse(request.url)
        if self.authority and parsed.netloc != self.authority:
            add_discrepancy("authority", self.authority, parsed.netloc)

        if self.method and request.method != self.method:
            add_discrepancy("method", self.method, request.method)

        for param, expected_value in self.required_params.items():
            actual_value = request.query.get(param)
            if actual_value != expected_value:
                add_discrepancy(param, expected_value, actual_value)

        for header, expected_value in self.required_headers.items():
            actual_value = request.headers.get(header)

            # UserAgentPolicy appends the value of $AZURE_HTTP_USER_AGENT, which is set in
            # pipelines, so we accept a user agent which merely contains the expected value
            if header.lower() == "user-agent":
                if expected_value not in actual_value:
                    add_discrepancy("user-agent", "contains " + expected_value, actual_value)
            elif actual_value != expected_value:
                add_discrepancy(header, expected_value, actual_value)

        for field, expected_value in self.required_data.items():
            actual_value = request.body.get(field)
            if actual_value != expected_value:
                add_discrepancy("form field", expected_value, actual_value)

        assert not discrepancies, "Unexpected request\n\t" + "\n\t".join(discrepancies)


def mock_response(status_code=200, headers=None, json_payload=None):
    response = mock.Mock(status_code=status_code, headers=headers or {})
    if json_payload is not None:
        response.text = lambda encoding=None: json.dumps(json_payload)
        response.headers["content-type"] = "application/json"
        response.content_type = "application/json"
    else:
        response.text = lambda encoding=None: ""
        response.headers["content-type"] = "text/plain"
        response.content_type = "text/plain"
    return response


def get_discovery_response(endpoint="https://a/b"):
    """Get a mock response containing the values MSAL requires from tenant and instance discovery.

    The response is incomplete and its values aren't necessarily valid, particularly for instance discovery, but it's
    sufficient. MSAL will send token requests to "{endpoint}/oauth2/v2.0/token_endpoint" after receiving a tenant
    discovery response created by this method.
    """
    aad_metadata_endpoint_names = (
        "authorization_endpoint",
        "token_endpoint",
        "tenant_discovery_endpoint",
        "device_authorization_endpoint",
    )
    payload = {name: endpoint + "/oauth2/v2.0/" + name for name in aad_metadata_endpoint_names}
    payload["metadata"] = ""
    return mock_response(json_payload=payload)


def validating_transport(requests, responses):
    if len(requests) != len(responses):
        raise ValueError("each request must have one response")

    sessions = zip(requests, responses)
    sessions = (s for s in sessions)  # 2.7's zip returns a list, and nesting a generator doesn't break it for 3.x

    def validate_request(request, **kwargs):
        assert "tenant_id" not in kwargs
        assert "claims" not in kwargs
        try:
            expected_request, response = next(sessions)
        except StopIteration:
            assert False, "unexpected request: {} {}".format(request.method, request.url)
        expected_request.assert_matches(request)
        return response

    return mock.Mock(send=mock.Mock(wraps=validate_request))


def msal_validating_transport(requests, responses, **kwargs):
    """a validating transport with default responses to MSAL's discovery requests"""
    return validating_transport([Request()] * 2 + requests, [get_discovery_response(**kwargs)] * 2 + responses)


def new_msal_validating_transport(requests, responses, **kwargs):
    """a transport with default responses to MSAL's discovery requests without validation"""
    """msal made some optimizations to make less calls to discovery endpoint"""
    return validating_transport([Request()] + requests, [get_discovery_response(**kwargs)] + responses)


def urlsafeb64_decode(s):
    if isinstance(s, str):
        s = s.encode("ascii")

    padding_needed = 4 - len(s) % 4
    return base64.urlsafe_b64decode(s + b"=" * padding_needed)


def get_token_payload_contents(token: str):
    _, payload, _ = token.split(".")
    decoded_payload = urlsafeb64_decode(payload).decode()
    return json.loads(decoded_payload)
