# Copyright 2022 IQM client developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for user authentication and token management in IQM client.
"""
# pylint: disable=too-many-arguments
import builtins
import io
import json
from time import sleep

from mockito import expect, unstub, verifyNoUnwantedInteractions, when
import pytest
import requests

from iqm.iqm_client import ClientAuthenticationError, ClientConfigurationError, Credentials, IQMClient
from iqm.iqm_client.iqm_client import _time_left_seconds
from tests.conftest import (
    MockJsonResponse,
    expect_logout,
    get_jobs_args,
    make_token,
    post_jobs_args,
    prepare_tokens,
    submit_circuits_args,
)


def test_get_initial_tokens_with_credentials_from_arguments(credentials):
    """
    Tests that if the client is initialized with credentials, they are used correctly
    """
    tokens = prepare_tokens(300, 3600, **credentials)
    expected_credentials = Credentials(
        access_token=tokens['access_token'], refresh_token=tokens['refresh_token'], **credentials
    )

    client = IQMClient('https://example.com', **credentials)
    assert client._credentials == expected_credentials

    unstub()


def test_get_token_from_arguments(monkeypatch):
    """
    Tests that plaintext token is read from the arguments
    """
    monkeypatch.setenv('IQM_TOKEN', '2345')
    client = IQMClient(url='https://example.com', token='1234')
    assert client._token == '1234'


def test_get_initial_tokens_with_credentials_from_env_variables(credentials, monkeypatch):
    """
    Tests that credentials are read from environment variables if they are not given as arguments
    """
    tokens = prepare_tokens(300, 3600, **credentials)
    expected_credentials = Credentials(
        access_token=tokens['access_token'], refresh_token=tokens['refresh_token'], **credentials
    )
    monkeypatch.setenv('IQM_AUTH_SERVER', credentials['auth_server_url'])
    monkeypatch.setenv('IQM_AUTH_USERNAME', credentials['username'])
    monkeypatch.setenv('IQM_AUTH_PASSWORD', credentials['password'])

    client = IQMClient(url='https://example.com')
    assert client._credentials == expected_credentials
    unstub()


def test_get_token_from_env_variable(monkeypatch):
    """
    Tests that plaintext token is read from the environment variable if they are not given as arguments
    """
    monkeypatch.setenv('IQM_TOKEN', '1234')
    client = IQMClient(url='https://example.com')
    assert client._token == '1234'


def test_get_initial_tokens_with_incomplete_credentials_from_env_variables(credentials, monkeypatch):
    """
    Tests configuration error is reported if IQM_AUTH_SERVER is set, but no credentials provided
    """
    monkeypatch.setenv('IQM_AUTH_SERVER', credentials['auth_server_url'])
    with pytest.raises(ClientConfigurationError) as e:
        IQMClient('https://example.com')
    assert str(e.value) == 'Auth server URL is set but no username or password'
    unstub()


def test_add_authorization_header_when_credentials_are_provided(
    client_with_credentials, existing_job_url, existing_run_id, pending_compilation_job_result
):
    """
    Tests that ``get_run`` requests are sent with Authorization header when credentials are provided
    """
    expect(requests, times=1).get(
        existing_job_url, **post_jobs_args(access_token=client_with_credentials._credentials.access_token)
    ).thenReturn(pending_compilation_job_result)

    result = client_with_credentials.get_run(existing_run_id)
    assert result.status == 'pending compilation'

    verifyNoUnwantedInteractions()
    unstub()


def test_add_authorization_header_on_submit_circuits_when_credentials_are_provided(
    client_with_credentials,
    jobs_url,
    minimal_run_request,
    existing_run_id,
    submit_success,
    quantum_architecture_url,
    quantum_architecture_success,
):
    """
    Tests that ``submit_circuits`` requests are sent with Authorization header when credentials are provided
    """
    expect(requests, times=1).get(quantum_architecture_url, ...).thenReturn(quantum_architecture_success)
    expect(requests, times=1).post(
        jobs_url, **post_jobs_args(minimal_run_request, access_token=client_with_credentials._credentials.access_token)
    ).thenReturn(submit_success)

    assert existing_run_id == client_with_credentials.submit_circuits(**submit_circuits_args(minimal_run_request))

    verifyNoUnwantedInteractions()
    unstub()


def test_submit_circuits_raises_when_auth_failed(
    client_with_credentials,
    jobs_url,
    minimal_run_request,
    submit_failed_auth,
    quantum_architecture_url,
    quantum_architecture_success,
):
    """
    Tests that ``submit_circuits`` raises ClientAuthenticationError when authentication fails
    """
    expect(requests, times=1).get(quantum_architecture_url, ...).thenReturn(quantum_architecture_success)
    expect(requests, times=1).post(
        jobs_url, **post_jobs_args(minimal_run_request, access_token=client_with_credentials._credentials.access_token)
    ).thenReturn(submit_failed_auth)

    with pytest.raises(ClientAuthenticationError) as e:
        client_with_credentials.submit_circuits(**submit_circuits_args(minimal_run_request))
    assert str(e.value).startswith('Authentication failed')

    verifyNoUnwantedInteractions()
    unstub()


def test_add_authorization_header_on_get_jobs_when_external_token_is_provided(
    client_with_external_token, existing_job_url, tokens_dict, existing_run_id, pending_execution_job_result
):
    """
    Tests that get jobs requests are sent with Authorization header when external token is provided
    """
    expect(requests, times=1).get(
        existing_job_url, **get_jobs_args(access_token=tokens_dict['access_token'])
    ).thenReturn(pending_execution_job_result)

    result = client_with_external_token.get_run(existing_run_id)
    assert result.status == 'pending execution'

    verifyNoUnwantedInteractions()
    unstub()


def test_no_authorization_header_on_get_jobs_when_credentials_are_not_provided(
    sample_client, existing_job_url, existing_run_id, pending_execution_job_result
):
    """
    Tests that get jobs requests are sent without Authorization header when no credentials are provided
    """
    expect(requests, times=1).get(existing_job_url, **get_jobs_args(access_token=None)).thenReturn(
        pending_execution_job_result
    )

    result = sample_client.get_run(existing_run_id)
    assert result.status == 'pending execution'

    verifyNoUnwantedInteractions()
    unstub()


def test_raises_client_authentication_error_if_authentication_fails(base_url, credentials):
    """
    Tests that authentication failure raises ClientAuthenticationError
    """
    prepare_tokens(300, 3600, status_code=401, **credentials)
    with pytest.raises(ClientAuthenticationError) as e:
        IQMClient(base_url, **credentials)
    assert str(e.value).startswith('Failed to update tokens')

    verifyNoUnwantedInteractions()
    unstub()


def test_get_quantum_architecture_raises_if_no_auth_provided(sample_client, quantum_architecture_url):
    """Test retrieving the quantum architecture if server responded with redirect"""
    redirection_response = requests.Response()
    redirection_response.status_code = 302

    expect(requests, times=1).get(quantum_architecture_url, **get_jobs_args(access_token=None)).thenReturn(
        MockJsonResponse(401, {'detail': 'unauthorized'}, [redirection_response])
    )

    with pytest.raises(ClientConfigurationError) as e:
        sample_client.get_quantum_architecture()
    assert str(e.value) == 'Authentication is required.'

    verifyNoUnwantedInteractions()
    unstub()


def test_get_quantum_architecture_raises_if_wrong_auth_provided(client_with_credentials, quantum_architecture_url):
    """Test retrieving the quantum architecture if server responded with auth error"""
    expect(requests, times=1).get(
        quantum_architecture_url, **get_jobs_args(access_token=client_with_credentials._credentials.access_token)
    ).thenReturn(MockJsonResponse(401, {'detail': 'unauthorized'}))

    with pytest.raises(ClientAuthenticationError) as e:
        client_with_credentials.get_quantum_architecture()
    assert str(e.value).startswith('Authentication failed')

    verifyNoUnwantedInteractions()
    unstub()


def test_access_token_is_not_refreshed_if_it_has_not_expired(
    client_with_credentials, credentials, existing_job_url, existing_run_id, pending_execution_job_result
):
    """
    Test that access token is not refreshed if it has not expired
    """
    tokens = prepare_tokens(300, 3600, **credentials)
    assert client_with_credentials._credentials.access_token == tokens['access_token']

    expect(requests, times=3).get(
        existing_job_url, **get_jobs_args(access_token=client_with_credentials._credentials.access_token)
    ).thenReturn(pending_execution_job_result)

    client_with_credentials.get_run(existing_run_id)
    client_with_credentials.get_run(existing_run_id)
    client_with_credentials.get_run(existing_run_id)

    verifyNoUnwantedInteractions()
    unstub()


def test_expired_access_token_is_refreshed_automatically(
    client_with_credentials, credentials, existing_job_url, existing_run_id, pending_execution_job_result
):
    """
    Test that access token is refreshed automatically if it has expired
    """
    # Provide client with an expired access token
    client_with_credentials._credentials.access_token = make_token('Bearer', -300)
    client_with_credentials._credentials.refresh_token = make_token('Refresh', 4200)

    # Prepare for token refresh request
    refreshed_tokens = prepare_tokens(300, 4200, client_with_credentials._credentials.refresh_token, **credentials)

    # Expect get jobs request with refreshed token
    expect(requests, times=1).get(
        existing_job_url, **get_jobs_args(access_token=refreshed_tokens['access_token'])
    ).thenReturn(pending_execution_job_result)

    # Client should refresh tokens for get jobs request
    result = client_with_credentials.get_run(existing_run_id)
    assert result.status == 'pending execution'

    # Verify that access token has been refreshed
    assert client_with_credentials._credentials.access_token == refreshed_tokens['access_token']

    verifyNoUnwantedInteractions()
    unstub()


def test_start_new_session_when_refresh_token_has_expired(
    client_with_credentials, credentials, existing_job_url, existing_run_id, pending_execution_job_result
):
    """
    Test that a new session is started automatically if refresh token has expired
    """
    # Provide client with an expired tokens
    client_with_credentials._credentials.access_token = make_token('Bearer', -300)
    client_with_credentials._credentials.refresh_token = make_token('Refresh', -300)

    # Prepare for new session start instead of token refresh request
    refreshed_tokens = prepare_tokens(300, 3600, **credentials)

    # Expect get jobs request with refreshed token
    expect(requests, times=1).get(
        existing_job_url, **get_jobs_args(access_token=refreshed_tokens['access_token'])
    ).thenReturn(pending_execution_job_result)

    # Client should refresh tokens for get jobs request
    result = client_with_credentials.get_run(existing_run_id)
    assert result.status == 'pending execution'

    assert client_with_credentials._credentials.access_token == refreshed_tokens['access_token']
    assert client_with_credentials._credentials.refresh_token == refreshed_tokens['refresh_token']

    verifyNoUnwantedInteractions()
    unstub()


def test_tokens_are_cleared_at_logout(client_with_credentials, credentials):
    """
    Tests that calling ``close`` will terminate the session and clear tokens
    """
    expect_logout(credentials['auth_server_url'], client_with_credentials._credentials.refresh_token)

    client_with_credentials.close_auth_session()
    assert client_with_credentials._credentials.access_token is None
    assert client_with_credentials._credentials.refresh_token is None

    verifyNoUnwantedInteractions()
    unstub()


def test_cannot_close_external_auth_session(client_with_external_token):
    """
    Tests that calling ``close_auth_session`` while initialized with an external auth session
    raises ClientAuthenticationError
    """
    with pytest.raises(ClientAuthenticationError) as e:
        client_with_external_token.close_auth_session()
    assert 'Unable to close externally managed auth session' == str(e.value)


def test_logout_on_client_destruction(base_url, credentials):
    """
    Tests that client is trying to terminate the authentication session on destruction
    """
    prepare_tokens(300, 300, **credentials)
    client = IQMClient(base_url, **credentials)

    expect_logout(credentials['auth_server_url'], client._credentials.refresh_token)
    del client

    verifyNoUnwantedInteractions()
    unstub()


def test_external_token_updated_if_expired(base_url):
    """
    Tests that client gets updated token from tokens file if old external token has expired
    """
    tokens_path = 'dummy_path'

    def setup_tokens_file(access_token_lifetime):
        access_token = make_token('Bearer', access_token_lifetime)
        tokens_file_contents = json.dumps({'auth_server_url': base_url + '/auth', 'access_token': access_token})
        when(builtins).open(tokens_path, 'r', encoding='utf-8').thenReturn(io.StringIO(tokens_file_contents))

    setup_tokens_file(1)
    client = IQMClient(base_url, tokens_file=tokens_path)
    sleep(2)
    setup_tokens_file(300)
    bearer_token = client._get_bearer_token()
    assert bearer_token is not None
    assert _time_left_seconds(bearer_token.removeprefix('Bearer ')) > 0
    unstub()


def test_tokens_file_overrides_plaintext_token(base_url):
    """
    Tests tokens_file is preferred over plaintext token
    """
    tokens_path = 'dummy_path'
    access_token = make_token('Bearer', lifetime=300)
    tokens_file_contents = json.dumps({'auth_server_url': base_url + '/auth', 'access_token': access_token})
    when(builtins).open(tokens_path, 'r', encoding='utf-8').thenReturn(io.StringIO(tokens_file_contents))
    try:
        client = IQMClient(base_url, tokens_file=tokens_path, token='12345')
        bearer_token = client._get_bearer_token()
        assert bearer_token is not None
        assert bearer_token == f'Bearer {access_token}'
    finally:
        unstub()


def test_plaintext_token_is_used_as_bearer_token_if_tokens_file_is_not_set(base_url):
    """
    Tests that plaintext token is used as a bearer token correctly
    """
    client = IQMClient(base_url, token='1234')
    bearer_token = client._get_bearer_token()
    assert bearer_token == 'Bearer 1234'
