# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import re
import base64
from azure_devtools.scenario_tests import RecordingProcessor
from datetime import datetime, timedelta
from functools import wraps
from urllib.parse import urlparse
import sys

def generate_token_with_custom_expiry(valid_for_seconds):
    return generate_token_with_custom_expiry_epoch((datetime.now() + timedelta(seconds=valid_for_seconds)).timestamp())
 
def generate_token_with_custom_expiry_epoch(expires_on_epoch):
    expiry_json = f'{{"exp": {str(expires_on_epoch)} }}'
    base64expiry = base64.b64encode(
        expiry_json.encode('utf-8')).decode('utf-8').rstrip("=")
    token_template = (f'''eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.
        {base64expiry}.adM-ddBZZlQ1WlN3pdPBOF5G4Wh9iZpxNP_fSvpF4cWs''')
    return token_template


class URIIdentityReplacer(RecordingProcessor):
    """Replace the identity in request uri"""
    def process_request(self, request):
        resource = (urlparse(request.uri).netloc).split('.')[0]
        request.uri = re.sub('/identities/([^/?]+)', '/identities/sanitized', request.uri) 
        request.uri = re.sub(resource, 'sanitized', request.uri)
        request.uri = re.sub('/identities/([^/?]+)', '/identities/sanitized', request.uri) 
        request.uri = re.sub(resource, 'sanitized', request.uri)
        return request
    
    def process_response(self, response):
        if 'url' in response:
            response['url'] = re.sub('/identities/([^/?]+)', '/identities/sanitized', response['url'])
        return response
    
class URIMsalUsernameReplacer(RecordingProcessor):
    """Replace the MSAL username in request uri"""
    def process_request(self, request):
        resource = (urlparse(request.uri).netloc).split('.')[0]
        request.uri = re.sub('common/userrealm/([^/.]+)', 'common/userrealm/sanitized@test', request.uri) 
        request.uri = re.sub(resource, 'sanitized', request.uri)
        return request
    
    def process_response(self, response):
        if 'url' in response:
            response['url'] = re.sub('common/userrealm/([^/.]+)', 'common/userrealm/sanitized@test', response['url'])
        return response
    
class URIReplacerProcessor(RecordingProcessor):
    def __init__(self, keys=None, replacement="sanitized"):
        self._keys = keys if keys else []
        self._replacement = replacement

    def process_request(self, request):
        request.uri = re.sub(
            "https://([^/?])*.communication",
            "https://sanitized.communication",
            request.uri,
        )
        return request

    def process_response(self, response):
        if 'url' in response :
            response['url'] = re.sub(
                "https://([^/?])*.communication",
                "https://sanitized.communication",
                response['url'],
            )
        return response