/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */

package org.apache.atlas;

import com.sun.jersey.api.client.Client;
import com.sun.jersey.api.client.ClientHandlerException;
import com.sun.jersey.api.client.ClientResponse;
import com.sun.jersey.api.client.WebResource;
import org.apache.atlas.model.legacy.EntityResult;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.v1.model.instance.Referenceable;
import org.apache.commons.configuration.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.mockito.Matchers;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriBuilder;

import java.net.ConnectException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Arrays;
import java.util.List;

import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

public class AtlasClientTest {
    @Mock
    private WebResource service;

    @Mock
    private WebResource.Builder resourceBuilderMock;

    @Mock
    private Configuration configuration;

    @Mock
    private Client client;

    @BeforeMethod
    public void setup() {
        MockitoAnnotations.initMocks(this);
    }

    @Test
    public void shouldVerifyServerIsReady() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.VERSION, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());
        when(response.getEntity(String.class)).thenReturn("{\"Version\":\"version-rrelease\",\"Name\":\"apache-atlas\",\"Description\":\"Metadata Management and Data Governance Platform over Hadoop\"}");
        when(builder.method(AtlasClient.API_V1.VERSION.getMethod(), ClientResponse.class, null)).thenReturn(response);

        assertTrue(atlasClient.isServerReady());
    }

    @Test
    public void testCreateEntity() throws Exception {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.CREATE_ENTITY, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.CREATED.getStatusCode());

        String jsonResponse = AtlasType.toV1Json(new EntityResult(Arrays.asList("id"), null, null));

        when(response.getEntity(String.class)).thenReturn(jsonResponse);
        when(response.getLength()).thenReturn(jsonResponse.length());

        String entityJson = AtlasType.toV1Json(new Referenceable("type"));

        when(builder.method(anyString(), Matchers.<Class>any(), anyString())).thenReturn(response);

        List<String> ids = atlasClient.createEntity(entityJson);

        assertEquals(ids.size(), 1);
        assertEquals(ids.get(0), "id");
    }

    @Test
    public void shouldReturnFalseIfServerIsNotReady() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.VERSION, service);

        when(builder.method(AtlasClient.API_V1.VERSION.getMethod(), ClientResponse.class, null)).thenThrow(new ClientHandlerException());
        assertFalse(atlasClient.isServerReady());
    }

    @Test
    public void shouldReturnFalseIfServiceIsUnavailable() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.VERSION, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.SERVICE_UNAVAILABLE.getStatusCode());
        when(response.getClientResponseStatus()).thenReturn(ClientResponse.Status.SERVICE_UNAVAILABLE);

        when(builder.method(AtlasClient.API_V1.VERSION.getMethod(), ClientResponse.class, null)).thenReturn(response);

        assertFalse(atlasClient.isServerReady());
    }

    @Test(expectedExceptions = AtlasServiceException.class)
    public void shouldThrowErrorIfAnyResponseOtherThanServiceUnavailable() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.VERSION, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode());
        when(response.getClientResponseStatus()).thenReturn(ClientResponse.Status.INTERNAL_SERVER_ERROR);

        when(builder.method(AtlasClient.API_V1.VERSION.getMethod(), ClientResponse.class, null)).thenReturn(response);

        atlasClient.isServerReady();
        fail("Should throw exception");
    }

    @Test
    public void shouldGetAdminStatus() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"Active\"}";

        when(response.getEntity(String.class)).thenReturn(activeStatus);
        when(response.getLength()).thenReturn(activeStatus.length());
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null)).thenReturn(response);

//         Fix after AtlasBaseClient
//        atlasClient.setService();

        String status = atlasClient.getAdminStatus();

        assertEquals(status, "Active");
    }

    @Test(expectedExceptions = AtlasServiceException.class)
    public void shouldReturnStatusAsUnknownOnException() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode());
        when(response.getClientResponseStatus()).thenReturn(ClientResponse.Status.INTERNAL_SERVER_ERROR);
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null)).thenReturn(response);

        atlasClient.getAdminStatus();

        fail("Should fail with AtlasServiceException");
    }

    @Test
    public void shouldReturnStatusAsUnknownIfJSONIsInvalid() throws AtlasServiceException {
        setupRetryParams();

        AtlasClient         atlasClient = new AtlasClient(service, configuration);
        WebResource.Builder builder     = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      response    = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());
        when(response.getEntity(String.class)).thenReturn("{\"status\":\"Active\"}");
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null)).thenReturn(response);

        String status = atlasClient.getAdminStatus();

        assertEquals(status, AtlasClient.UNKNOWN_STATUS);
    }

    @Test
    public void shouldReturnBaseURLAsPassedInURL() {
        AtlasClient atlasClient = new AtlasClient(service, configuration);
        String      serviceURL  = atlasClient.determineActiveServiceURL(new String[] {"http://localhost:21000"}, client);

        assertEquals(serviceURL, "http://localhost:21000");
    }

    @Test
    public void shouldSelectActiveAmongMultipleServersIfHAIsEnabled() {
        setupRetryParams();

        when(client.resource(UriBuilder.fromUri("http://localhost:31000").build())).thenReturn(service);
        when(client.resource(UriBuilder.fromUri("http://localhost:41000").build())).thenReturn(service);

        WebResource.Builder builder       = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      firstResponse = mock(ClientResponse.class);

        when(firstResponse.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String passiveStatus = "{\"Status\":\"PASSIVE\"}";

        when(firstResponse.getEntity(String.class)).thenReturn(passiveStatus);
        when(firstResponse.getLength()).thenReturn(passiveStatus.length());

        ClientResponse secondResponse = mock(ClientResponse.class);

        when(secondResponse.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"ACTIVE\"}";

        when(secondResponse.getEntity(String.class)).thenReturn(activeStatus);
        when(secondResponse.getLength()).thenReturn(activeStatus.length());
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null))
                .thenReturn(firstResponse).thenReturn(firstResponse).thenReturn(firstResponse)
                .thenReturn(secondResponse);

        AtlasClient atlasClient = new AtlasClient(service, configuration);
        String      serviceURL  = atlasClient.determineActiveServiceURL(new String[] {"http://localhost:31000", "http://localhost:41000"}, client);

        assertEquals(serviceURL, "http://localhost:41000");
    }

    @Test
    public void shouldRetryUntilServiceBecomesActive() {
        setupRetryParams();

        when(client.resource(UriBuilder.fromUri("http://localhost:31000").build())).thenReturn(service);

        WebResource.Builder builder  = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      response = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());
        when(response.getEntity(String.class)).thenReturn("{\"Status\":\"BECOMING_ACTIVE\"}");

        ClientResponse nextResponse = mock(ClientResponse.class);

        when(nextResponse.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"ACTIVE\"}";

        when(response.getEntity(String.class)).thenReturn(activeStatus);
        when(response.getLength()).thenReturn(activeStatus.length());
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null))
                .thenReturn(response).thenReturn(response).thenReturn(nextResponse);

        AtlasClient atlasClient = new AtlasClient(service, configuration);
        String      serviceURL  = atlasClient.determineActiveServiceURL(new String[] {"http://localhost:31000", "http://localhost:41000"}, client);

        assertEquals(serviceURL, "http://localhost:31000");
    }

    @Test
    public void shouldRetryIfCannotConnectToServiceInitially() {
        setupRetryParams();

        when(client.resource(UriBuilder.fromUri("http://localhost:31000").build())).thenReturn(service);

        WebResource.Builder builder  = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      response = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());
        when(response.getEntity(String.class)).thenReturn("{\"Status\":\"BECOMING_ACTIVE\"}");

        ClientResponse nextResponse = mock(ClientResponse.class);

        when(nextResponse.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"ACTIVE\"}";

        when(response.getEntity(String.class)).thenReturn(activeStatus);
        when(response.getLength()).thenReturn(activeStatus.length());
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null))
                .thenThrow(new ClientHandlerException("Simulating connection exception"))
                .thenReturn(response)
                .thenReturn(nextResponse);

        AtlasClient atlasClient = new AtlasClient(service, configuration);

        atlasClient.setService(service);
        atlasClient.setConfiguration(configuration);

        String serviceURL = atlasClient.determineActiveServiceURL(new String[] {"http://localhost:31000", "http://localhost:41000"}, client);

        assertEquals(serviceURL, "http://localhost:31000");
    }

    @Test(expectedExceptions = IllegalArgumentException.class)
    public void shouldThrowExceptionIfActiveServerIsNotFound() {
        setupRetryParams();

        when(client.resource(UriBuilder.fromUri("http://localhost:31000").build())).thenReturn(service);

        WebResource.Builder builder  = setupBuilder(AtlasClient.API_V1.STATUS, service);
        ClientResponse      response = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());
        when(response.getEntity(String.class)).thenReturn("{\"Status\":\"BECOMING_ACTIVE\"}");
        when(builder.method(AtlasClient.API_V1.STATUS.getMethod(), ClientResponse.class, null))
                .thenThrow(new ClientHandlerException("Simulating connection exception"))
                .thenReturn(response)
                .thenReturn(response);

        AtlasClient atlasClient = new AtlasClient(service, configuration);
        String      serviceURL  = atlasClient.determineActiveServiceURL(new String[] {"http://localhost:31000", "http://localhost:41000"}, client);

        assertNull(serviceURL);
    }

    @Test
    public void shouldRetryAPICallsOnClientHandlerException() throws AtlasServiceException, URISyntaxException {
        setupRetryParams();

        ResourceCreator resourceCreator = mock(ResourceCreator.class);
        WebResource     resourceObject  = mock(WebResource.class);

        when(resourceObject.getURI())
                .thenReturn(new URI("http://localhost:31000/api/atlas/types"))
                .thenReturn(new URI("http://localhost:41000/api/atlas/types"))
                .thenReturn(new URI("http://localhost:41000/api/atlas/types"));

        WebResource.Builder builder  = getBuilder(resourceObject);
        ClientResponse      response = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"ACTIVE\"}";

        when(response.getEntity(String.class)).thenReturn(activeStatus);
        when(response.getLength()).thenReturn(activeStatus.length());

        when(builder.method(AtlasClient.API_V1.LIST_TYPES.getMethod(), ClientResponse.class, null))
                .thenThrow(new ClientHandlerException("simulating exception in calling API", new ConnectException()))
                .thenReturn(response);

        when(resourceCreator.createResource()).thenReturn(resourceObject);

        AtlasClient atlasClient = getClientForTest("http://localhost:31000", "http://localhost:41000");

        atlasClient.setService(service);
        atlasClient.setConfiguration(configuration);

        atlasClient.callAPIWithRetries(AtlasClient.API_V1.LIST_TYPES, null, resourceCreator);

        verify(client).destroy();
        verify(client).resource(UriBuilder.fromUri("http://localhost:31000").build());
        verify(client).resource(UriBuilder.fromUri("http://localhost:41000").build());
    }

    @Test
    public void shouldRetryWithSameClientIfSingleAddressIsUsed() throws URISyntaxException, AtlasServiceException {
        setupRetryParams();

        ResourceCreator resourceCreator = mock(ResourceCreator.class);
        WebResource     resourceObject  = mock(WebResource.class);

        when(resourceObject.getURI()).thenReturn(new URI("http://localhost:31000/api/atlas/types"));

        WebResource.Builder builder = getBuilder(resourceObject);

        ClientResponse response = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"ACTIVE\"}";

        when(response.getEntity(String.class)).thenReturn(activeStatus);
        when(response.getLength()).thenReturn(activeStatus.length());

        when(builder.method(AtlasClient.API_V1.LIST_TYPES.getMethod(), ClientResponse.class, null))
                .thenThrow(new ClientHandlerException("simulating exception in calling API", new ConnectException()))
                .thenReturn(response);

        when(resourceCreator.createResource()).thenReturn(resourceObject);
        when(configuration.getString("atlas.http.authentication.type", "simple")).thenReturn("simple");

        AtlasClient atlasClient = getClientForTest("http://localhost:31000");

        atlasClient.setService(resourceObject);
        atlasClient.setConfiguration(configuration);

        atlasClient.callAPIWithRetries(AtlasClient.API_V1.LIST_TYPES, null, resourceCreator);

        verify(client).destroy();
        verify(client, times(2)).resource(UriBuilder.fromUri("http://localhost:31000").build());
    }

    @Test
    public void shouldRetryAPICallsOnServiceUnavailable() throws AtlasServiceException, URISyntaxException {
        setupRetryParams();

        ResourceCreator resourceCreator = mock(ResourceCreator.class);
        WebResource     resourceObject  = mock(WebResource.class);

        when(resourceObject.getURI())
                .thenReturn(new URI("http://localhost:31000/api/atlas/types"))
                .thenReturn(new URI("http://localhost:41000/api/atlas/types"))
                .thenReturn(new URI("http://localhost:41000/api/atlas/types"));

        WebResource.Builder builder = getBuilder(resourceObject);

        ClientResponse firstResponse = mock(ClientResponse.class);

        when(firstResponse.getStatus()).thenReturn(Response.Status.SERVICE_UNAVAILABLE.getStatusCode());
        when(firstResponse.getClientResponseStatus()).thenReturn(ClientResponse.Status.SERVICE_UNAVAILABLE);

        ClientResponse response = mock(ClientResponse.class);

        when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());

        String activeStatus = "{\"Status\":\"ACTIVE\"}";

        when(response.getEntity(String.class)).thenReturn(activeStatus);
        when(response.getLength()).thenReturn(activeStatus.length());

        when(builder.method(AtlasClient.API_V1.LIST_TYPES.getMethod(), ClientResponse.class, null))
                .thenThrow(new ClientHandlerException("simulating exception in calling API", new ConnectException()))
                .thenReturn(firstResponse)
                .thenReturn(response);

        when(resourceCreator.createResource()).thenReturn(resourceObject);

        AtlasClient atlasClient = getClientForTest("http://localhost:31000", "http://localhost:41000");

        atlasClient.setService(resourceObject);
        atlasClient.setConfiguration(configuration);

        atlasClient.callAPIWithRetries(AtlasClient.API_V1.LIST_TYPES, null, resourceCreator);

        verify(client).destroy();
        verify(client).resource(UriBuilder.fromUri("http://localhost:31000").build());
        verify(client).resource(UriBuilder.fromUri("http://localhost:41000").build());
    }

    private WebResource.Builder setupBuilder(AtlasClient.API_V1 api, WebResource webResource) {
        when(webResource.path(api.getPath())).thenReturn(service);
        when(webResource.path(api.getNormalizedPath())).thenReturn(service);

        return getBuilder(service);
    }

    private WebResource.Builder getBuilder(WebResource resourceObject) {
        when(resourceObject.getRequestBuilder()).thenReturn(resourceBuilderMock);
        when(resourceObject.path(anyString())).thenReturn(resourceObject);
        when(resourceBuilderMock.accept(AtlasBaseClient.JSON_MEDIA_TYPE)).thenReturn(resourceBuilderMock);
        when(resourceBuilderMock.accept(MediaType.APPLICATION_JSON)).thenReturn(resourceBuilderMock);
        when(resourceBuilderMock.type(AtlasBaseClient.JSON_MEDIA_TYPE)).thenReturn(resourceBuilderMock);
        when(resourceBuilderMock.type(MediaType.MULTIPART_FORM_DATA)).thenReturn(resourceBuilderMock);

        return resourceBuilderMock;
    }

    private void setupRetryParams() {
        when(configuration.getInt(AtlasClient.ATLAS_CLIENT_HA_RETRIES_KEY, AtlasClient.DEFAULT_NUM_RETRIES)).thenReturn(3);
        when(configuration.getInt(AtlasClient.ATLAS_CLIENT_HA_SLEEP_INTERVAL_MS_KEY, AtlasClient.DEFAULT_SLEEP_BETWEEN_RETRIES_MS)).thenReturn(1);
    }

    private AtlasClient getClientForTest(final String... baseUrls) {
        return new AtlasClient(null, (String) null, baseUrls) {
            boolean firstCall = true;

            @Override
            protected Client getClient(Configuration configuration, UserGroupInformation ugi, String doAsUser) {
                return client;
            }

            @Override
            protected String determineActiveServiceURL(String[] baseUrls, Client client) {
                String returnUrl = baseUrls[0];

                if (baseUrls.length > 1 && !firstCall) {
                    returnUrl = baseUrls[1];
                }

                firstCall = false;

                return returnUrl;
            }

            @Override
            protected Configuration getClientProperties() {
                return configuration;
            }
        };
    }
}
